diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index a110cc5b15d..eb0de9d9b60 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -3,7 +3,9 @@ Calling + translation logic for anthropic's `/v1/messages` endpoint """ import copy +import io import json +import re from collections.abc import Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast @@ -20,6 +22,11 @@ from litellm.anthropic_beta_headers_manager import ( from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator +from litellm.llms.base_llm.base_utils import ( + RESPONSE_FORMAT_ENVELOPE_NAME_KEYS, + ResponseFormatSchema, + response_format_envelope_contents, +) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -57,7 +64,7 @@ from litellm.types.utils import ( from ...base import BaseLLM from ..common_utils import AnthropicError, process_anthropic_headers -from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig +from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -69,6 +76,67 @@ def _loads_stream_chunk(payload: str) -> dict[str, object]: return json.loads(payload) +_EnvelopeVerdict = Literal["undecided", "pass_through", "hold"] +_FIRST_OBJECT_KEY: Final = re.compile(r'^\s*\{\s*"((?:[^"\\]|\\.)*)"\s*:') +_OBJECT_KEY_PREFIX: Final = re.compile(r'^\s*(?:\{\s*(?:"(?:[^"\\]|\\.)*\\?(?:"\s*)?)?)?$') +_ENVELOPE_NAME_KEY_PATTERN: Final = "|".join(sorted(RESPONSE_FORMAT_ENVELOPE_NAME_KEYS)) +_FIRST_NAME_VALUE: Final = re.compile(r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*"((?:[^"\\]|\\.)*)"') +_NAME_VALUE_PREFIX: Final = re.compile( + r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*(?:"(?:[^"\\]|\\.)*\\?)?$' +) +_STRING_VALUE_THEN_KEY: Final = re.compile(r'^"(?:[^"\\]|\\.)*"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:') +_UNDECIDED_PREFIX_LIMIT: Final = 4096 + + +def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict: + next_key: Final = _STRING_VALUE_THEN_KEY.match(string_value_start) + if next_key is None: + return "undecided" + return "pass_through" if next_key.group(1) in schema.property_names else "hold" + + +def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict: + value_start: Final = after_first_key.lstrip() + if not value_start: + return "undecided" + if value_start.startswith("{"): + return "hold" + if value_start.startswith('"') and schema.required_names: + return _key_after_a_string_verdict(value_start, schema) + return "pass_through" + + +def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict: + first_key: Final = _FIRST_OBJECT_KEY.match(buffered_json) + if first_key is None: + return "undecided" if _OBJECT_KEY_PREFIX.match(buffered_json) else "pass_through" + if first_key.group(1) in RESPONSE_FORMAT_ENVELOPE_NAME_KEYS: + first_value: Final = _FIRST_NAME_VALUE.match(buffered_json) + if first_value is not None and first_value.group(1) == tool_name: + return "hold" + if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json): + return "undecided" + if first_key.group(1) in schema.property_names or not schema.property_names or schema.accepts_a_lone_undeclared_key: + return "pass_through" + return _first_value_verdict(buffered_json[first_key.end() :], schema) + + +def _held_answer(held: str, schema: ResponseFormatSchema) -> str: + try: + payload: Final[object] = json.loads(held) + except (json.JSONDecodeError, RecursionError): + return held + contents: Final = response_format_envelope_contents(payload, schema, RESPONSE_FORMAT_TOOL_NAME) + return held if contents is None else json.dumps(contents) + + +def _envelope_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict: + verdict: Final = _prefix_verdict(buffered_json, schema, tool_name) + if verdict == "undecided" and len(buffered_json) > _UNDECIDED_PREFIX_LIMIT: + return "hold" + return verdict + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -81,6 +149,7 @@ async def make_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, + response_format_schema: ResponseFormatSchema | None = None, ) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_aclient @@ -116,6 +185,7 @@ async def make_call( json_mode=json_mode, speed=speed, tool_name_reverse_map=tool_name_reverse_map, + response_format_schema=response_format_schema, ) # LOGGING @@ -141,6 +211,7 @@ def make_sync_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, + response_format_schema: ResponseFormatSchema | None = None, ) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_client # re-use a module level client @@ -184,6 +255,7 @@ def make_sync_call( json_mode=json_mode, speed=speed, tool_name_reverse_map=tool_name_reverse_map, + response_format_schema=response_format_schema, ) # LOGGING @@ -241,6 +313,9 @@ class AnthropicChatCompletion(BaseLLM): tool_name_reverse_map=( litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY) if isinstance(litellm_params, dict) else None ), + response_format_schema=( + JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode and optional_params else None + ), ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -485,6 +560,9 @@ class AnthropicChatCompletion(BaseLLM): if isinstance(litellm_params, dict) else None ), + response_format_schema=( + JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None + ), ) return CustomStreamWrapper( completion_stream=completion_stream, @@ -550,6 +628,7 @@ class ModelResponseIterator: json_mode: bool | None = False, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, + response_format_schema: ResponseFormatSchema | None = None, ): self.streaming_response = streaming_response self.response_iterator = self.streaming_response @@ -573,6 +652,9 @@ class ModelResponseIterator: self.is_response_format_tool: bool = False # Track if we've converted any response_format tools (affects finish_reason) self.converted_response_format_tool: bool = False + self.response_format_schema = response_format_schema + self._response_format_buffer: io.StringIO = io.StringIO() + self._response_format_verdict: _EnvelopeVerdict = "undecided" # For handling partial JSON chunks from fragmentation # See: https://github.com/BerriAI/litellm/issues/17473 @@ -955,6 +1037,8 @@ class ModelResponseIterator: except (json.JSONDecodeError, TypeError): pass self._current_server_tool_id = None + if self.is_response_format_tool: + text = self._flush_response_format_buffer() # Reset response_format tool tracking when block stops self.is_response_format_tool = False # Reset current content block type @@ -1072,6 +1156,10 @@ class ModelResponseIterator: tool_name: Final = tool_use.get("function", {}).get("name", "") self.is_response_format_tool = tool_name == RESPONSE_FORMAT_TOOL_NAME + if self.is_response_format_tool and self.response_format_schema is not None: + self.converted_response_format_tool = True + return self._stream_response_format_fragment(tool_use, self.response_format_schema), None + # Convert tool to content if we're tracking a response_format tool if self.is_response_format_tool: message: Final = AnthropicConfig.convert_tool_response_to_message(tool_calls=[tool_use]) @@ -1083,6 +1171,31 @@ class ModelResponseIterator: return text, tool_use + def _stream_response_format_fragment( + self, tool_use: ChatCompletionToolCallChunk, schema: ResponseFormatSchema + ) -> str: + if tool_use.get("id") is not None: + self._response_format_buffer = io.StringIO() + self._response_format_verdict = "undecided" + return "" + fragment: Final = tool_use.get("function", {}).get("arguments") or "" + if self._response_format_verdict == "pass_through": + return fragment + self._response_format_buffer.write(fragment) + if self._response_format_verdict == "hold": + return "" + buffered: Final = self._response_format_buffer.getvalue() + self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME) + return buffered if self._response_format_verdict == "pass_through" else "" + + def _flush_response_format_buffer(self) -> str: + if self.response_format_schema is None or self._response_format_verdict == "pass_through": + return "" + held: Final = self._response_format_buffer.getvalue() + self._response_format_buffer = io.StringIO() + self._response_format_verdict = "undecided" + return _held_answer(held, self.response_format_schema) if held else "" + def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]: """ Handle message_delta event for finish_reason, usage, and container. diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 54a8780ca5b..c43e5f9b05b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast import httpx -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -35,7 +35,13 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import place_mid_conversation_system, split_leading_system_run, ) -from litellm.llms.base_llm.base_utils import type_to_response_format_param +from litellm.llms.base_llm.base_utils import ( + ResponseFormatSchema, + as_json_object, + response_format_schema, + type_to_response_format_param, + unwrap_response_format_envelope, +) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_ADVISOR_TOOL_TYPE, @@ -1746,7 +1752,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: _input_schema.update(cast(AnthropicInputSchema, json_schema)) - _tool: Final = AnthropicMessagesTool(name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema) + _tool: Final = AnthropicMessagesTool( + name=RESPONSE_FORMAT_TOOL_NAME, + input_schema=_input_schema, + ) return _tool def should_strip_billing_metadata(self) -> bool: @@ -2199,6 +2208,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, json_mode: bool | None, tool_calls: list[ChatCompletionToolCallChunk], + response_format_schema: ResponseFormatSchema | None = None, ) -> tuple[ LitellmMessage | None, list[ChatCompletionToolCallChunk], @@ -2218,11 +2228,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): json_tool: Final = tool_calls[json_indices[0]] if json_tool.get("function", {}).get("arguments") is None: return None, tool_calls, None - _message: Final = AnthropicConfig.convert_tool_response_to_message(tool_calls=[json_tool]) + _message: Final = AnthropicConfig.convert_tool_response_to_message( + tool_calls=[json_tool], response_format_schema=response_format_schema + ) return _message, [], None first_json: Final = tool_calls[json_indices[0]] - json_msg: Final = AnthropicConfig.convert_tool_response_to_message([first_json]) + json_msg: Final = AnthropicConfig.convert_tool_response_to_message( + [first_json], response_format_schema=response_format_schema + ) extra_content: Final[str | None] = json_msg.content if json_msg is not None else None filtered_tools: Final = [t for i, t in enumerate(tool_calls) if i not in json_indices] return None, filtered_tools, extra_content @@ -2658,6 +2672,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): prefix_prompt: str | None = None, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, + response_format_schema: ResponseFormatSchema | None = None, ): _hidden_params: Final[dict] = {} _hidden_params["additional_headers"] = process_anthropic_headers(dict(raw_response.headers)) @@ -2715,6 +2730,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): json_mode_message, tool_calls_for_message, json_extra_content = self._resolve_json_mode_non_streaming( json_mode=json_mode, tool_calls=tool_calls, + response_format_schema=response_format_schema, ) merged_text = text_content or "" if json_extra_content: @@ -2823,12 +2839,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): prefix_prompt=prefix_prompt, speed=speed, tool_name_reverse_map=tool_name_reverse_map, + response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None, ) return model_response @staticmethod def _convert_tool_response_to_message( tool_calls: list[ChatCompletionToolCallChunk], + response_format_schema: ResponseFormatSchema | None = None, ) -> LitellmMessage | None: """ In JSON mode, Anthropic API returns JSON schema as a tool call, we need to convert it to a message to follow the OpenAI format @@ -2836,28 +2854,32 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ ## HANDLE JSON MODE - anthropic returns single function call json_mode_content_str: Final[str | None] = tool_calls[0]["function"].get("arguments") + if json_mode_content_str is None: + return None try: - if json_mode_content_str is not None: - args: Final = json.loads(json_mode_content_str) - if isinstance(args, dict) and (values := args.get("values")) is not None: - _message = litellm.Message(content=json.dumps(values)) - return _message - else: - # a lot of the times the `values` key is not present in the tool response - # relevant issue: https://github.com/BerriAI/litellm/issues/6741 - _message = litellm.Message(content=json.dumps(args)) - return _message - except json.JSONDecodeError: + payload: Final[object] = json.loads(json_mode_content_str) + except (json.JSONDecodeError, RecursionError): # json decode error does occur, return the original tool response str return litellm.Message(content=json_mode_content_str) - return None + if response_format_schema is not None: + return litellm.Message( + content=json.dumps( + unwrap_response_format_envelope(payload, response_format_schema, RESPONSE_FORMAT_TOOL_NAME) + ) + ) + # a lot of the times the `values` key is not present in the tool response + # relevant issue: https://github.com/BerriAI/litellm/issues/6741 + payload_object: Final = as_json_object(payload) + values: Final = None if payload_object is None else payload_object.get("values") + return litellm.Message(content=json.dumps(payload if values is None else values)) @classmethod def convert_tool_response_to_message( cls, tool_calls: list[ChatCompletionToolCallChunk], # mutable-ok: mirrors override contract + response_format_schema: ResponseFormatSchema | None = None, ) -> LitellmMessage | None: - return cls._convert_tool_response_to_message(tool_calls) + return cls._convert_tool_response_to_message(tool_calls, response_format_schema) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return AnthropicError( @@ -2867,6 +2889,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) +_REQUEST_TOOLS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +class JsonModeRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + tools: object = None + + def tool_schema(self) -> ResponseFormatSchema | None: + try: + tools: Final = _REQUEST_TOOLS.validate_python(self.tools) + except ValidationError: + return None + json_tool: Final = next((tool for tool in tools if tool.get("name") == RESPONSE_FORMAT_TOOL_NAME), None) + return None if json_tool is None else response_format_schema(json_tool.get("input_schema")) + + def _valid_user_id(user_id: str) -> bool: """ Validate that user_id is not an email or phone number. diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index 6065c56d083..ca4dad5ee3e 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -11,6 +11,7 @@ import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion +from litellm.llms.anthropic.chat.transformation import JsonModeRequest from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -162,6 +163,9 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): logging_obj=logging_obj, timeout=timeout, json_mode=json_mode, + response_format_schema=( + JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None + ), ) from litellm.llms.anthropic.common_utils import ( process_anthropic_headers, diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index ba4dad40d13..4df5635b257 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -3,13 +3,15 @@ Utility functions for base LLM classes. """ import copy +import itertools import json from abc import ABC, abstractmethod -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from dataclasses import dataclass from typing import Any, Final from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk @@ -157,6 +159,136 @@ def convert_tool_response_to_message( _convert_tool_response_to_message = convert_tool_response_to_message +@dataclass(frozen=True, slots=True) +class ResponseFormatSchema: + property_names: frozenset[str] + required_names: frozenset[str] = frozenset() + allows_undeclared_keys: bool = False + + @property + def accepts_a_lone_undeclared_key(self) -> bool: + return self.allows_undeclared_keys and not self.required_names + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, JsonValue]) +_NAMES: Final = TypeAdapter(frozenset[str]) + + +def as_json_object(value: object) -> Mapping[str, JsonValue] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _names(value: object) -> frozenset[str]: + try: + return _NAMES.validate_python(value) + except ValidationError: + return frozenset() + + +def _allows_undeclared_keys(document: Mapping[str, JsonValue]) -> bool: + additional_properties: Final = document.get("additionalProperties") + return ( + additional_properties is True + or as_json_object(additional_properties) is not None + or bool(as_json_object(document.get("patternProperties"))) + ) + + +def response_format_schema(value: object) -> ResponseFormatSchema | None: + document: Final = as_json_object(value) + if document is None: + return None + properties: Final = as_json_object(document.get("properties", {})) + if properties is None: + return None + return ResponseFormatSchema( + property_names=frozenset(properties), + required_names=_names(document.get("required", ())), + allows_undeclared_keys=_allows_undeclared_keys(document), + ) + + +RESPONSE_FORMAT_ENVELOPE_NAME_KEYS: Final = frozenset({"name", "tool_name"}) + + +@dataclass(frozen=True, slots=True) +class _EnvelopeLayer: + inner: Mapping[str, JsonValue] + names_the_tool: bool + + +def _json_object_or_encoded_object(value: JsonValue) -> Mapping[str, JsonValue] | None: + if not isinstance(value, str): + return as_json_object(value) + try: + return as_json_object(json.loads(value)) + except (json.JSONDecodeError, RecursionError): + return None + + +def _label_keys(payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> frozenset[str]: + if not schema.required_names: + return frozenset() + declared: Final = schema.property_names | schema.required_names + return frozenset(key for key, value in payload_object.items() if isinstance(value, str) and key not in declared) + + +def _envelope_layer( + payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str +) -> _EnvelopeLayer | None: + name_keys: Final = frozenset( + key for key in RESPONSE_FORMAT_ENVELOPE_NAME_KEYS if payload_object.get(key) == tool_name + ) + label_keys: Final = name_keys or _label_keys(payload_object, schema) + wrapper_keys: Final = tuple(key for key in payload_object if key not in label_keys) + if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.property_names: + return None + wrapped: Final = payload_object[wrapper_keys[0]] + inner: Final = _json_object_or_encoded_object(wrapped) if name_keys else as_json_object(wrapped) + return None if inner is None else _EnvelopeLayer(inner=inner, names_the_tool=bool(name_keys)) + + +def _envelope_layers( + payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str +) -> Iterator[_EnvelopeLayer]: + layer = _envelope_layer(payload_object, schema, tool_name) # rebind-ok: the recursion gate needs a loop + while layer is not None: + yield layer + layer = _envelope_layer(layer.inner, schema, tool_name) + + +def _schema_object_inside_envelope( + payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str +) -> Mapping[str, JsonValue] | None: + layers: Final = _envelope_layers(payload_object, schema, tool_name) + candidates: Final = layers if schema.property_names else itertools.islice(layers, 1) + return next((layer.inner for layer in candidates if _holds_the_schema_object(layer, schema)), None) + + +def _holds_the_schema_object(layer: _EnvelopeLayer, schema: ResponseFormatSchema) -> bool: + if layer.names_the_tool: + return True + return not schema.accepts_a_lone_undeclared_key and not schema.property_names.isdisjoint(layer.inner) + + +def response_format_envelope_contents(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object | None: + payload_object: Final = as_json_object(payload) + if payload_object is None: + return None + schema_object: Final = _schema_object_inside_envelope(payload_object, schema, tool_name) + if schema_object is not None: + return schema_object + return None if "values" in schema.property_names else payload_object.get("values") + + +def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object: + contents: Final = response_format_envelope_contents(payload, schema, tool_name) + return payload if contents is None else contents + + def _dict_to_response_format_helper(response_format: dict, ref_template: str | None = None) -> dict: if ref_template is not None and response_format.get("type") == "json_schema": # Deep copy to avoid modifying original diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 4682a47de6b..d5565fed54d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1,9 +1,10 @@ +import json import types from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Final, cast import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger @@ -16,6 +17,8 @@ from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) +from litellm.llms.anthropic.chat.transformation import JsonModeRequest +from litellm.llms.base_llm.base_utils import ResponseFormatSchema from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -174,15 +177,34 @@ class AmazonCohereChatConfig: return optional_params +_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) + + +def response_format_schema_from_request(data: str, json_mode: bool | None) -> ResponseFormatSchema | None: + if not json_mode: + return None + try: + body: Final = _REQUEST_BODY.validate_python(json.loads(data)) + except (json.JSONDecodeError, ValidationError): + return None + return JsonModeRequest.model_validate(body).tool_schema() + + def _stream_decoder( bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None, *, model: str, json_mode: bool | None, sync_stream: bool, + response_format_schema: ResponseFormatSchema | None = None, ) -> "AWSEventStreamDecoder": if bedrock_invoke_provider == "anthropic": - return AmazonAnthropicClaudeStreamDecoder(model=model, sync_stream=sync_stream, json_mode=json_mode) + return AmazonAnthropicClaudeStreamDecoder( + model=model, + sync_stream=sync_stream, + json_mode=json_mode, + response_format_schema=response_format_schema, + ) if bedrock_invoke_provider == "deepseek_r1": return AmazonDeepSeekR1StreamDecoder(model=model, sync_stream=sync_stream) if bedrock_invoke_provider == "moonshot": @@ -250,7 +272,11 @@ async def make_call( ) else: decoder: Final = _stream_decoder( - bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=False + bedrock_invoke_provider, + model=model, + json_mode=json_mode, + sync_stream=False, + response_format_schema=response_format_schema_from_request(data, json_mode), ) completion_stream = decoder.aiter_bytes( response.aiter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers @@ -341,7 +367,11 @@ def make_sync_call( ) else: decoder: Final = _stream_decoder( - bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=True + bedrock_invoke_provider, + model=model, + json_mode=json_mode, + sync_stream=True, + response_format_schema=response_format_schema_from_request(data, json_mode), ) completion_stream = decoder.iter_bytes( response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers @@ -878,6 +908,7 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): model: str, sync_stream: bool, json_mode: bool | None = None, + response_format_schema: ResponseFormatSchema | None = None, ) -> None: """ Child class of AWSEventStreamDecoder that handles the streaming response from the Anthropic family of models @@ -889,6 +920,7 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): streaming_response=None, sync_stream=sync_stream, json_mode=json_mode, + response_format_schema=response_format_schema, ) def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py index 64ba1ad21cb..d5950d49d2b 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py @@ -10,8 +10,11 @@ import; extracting the helper into a leaf module resolves the warning and keeps the parent module's import surface narrow. """ +from collections.abc import Mapping from typing import Final +from pydantic import TypeAdapter, ValidationError + # Keys inside ``output_config`` that Vertex AI Claude rejects regardless of # the target model. Add an entry only when a 400 "Extra inputs are not # permitted" is reproducible against the live Vertex endpoint for every model. @@ -31,6 +34,24 @@ def _model_accepts_output_config_effort(model: str) -> bool: return AnthropicConfig.model_supports_effort_param(model, "vertex_ai") +_OUTPUT_CONFIG: Final = TypeAdapter(Mapping[str, object]) + + +def _as_output_config(value: object) -> Mapping[str, object] | None: + try: + return _OUTPUT_CONFIG.validate_python(value) + except ValidationError: + return None + + +def _fold_output_format_into_output_config( + output_config: Mapping[str, object] | None, output_format: object +) -> Mapping[str, object] | None: + if output_format is None: + return output_config + return {**(output_config or {}), "format": output_format} + + def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: """ Strip Vertex-unsupported keys from ``output_config`` / @@ -44,15 +65,18 @@ def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: * Keys in ``VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS`` are always filtered. * ``output_config`` left empty after filtering is removed so the request body has no empty dict. - * ``output_format`` is forwarded as-is (Vertex AI Claude accepts it). + * ``output_format`` is moved into ``output_config.format``, because + Vertex AI Claude rejects the top-level field as deprecated. When both + are set, ``output_format`` wins, as it does in the Anthropic adapter. * Non-dict values for ``output_config`` are dropped to avoid sending malformed payloads downstream. """ - output_config: Final = data.get("output_config") + request: Final[dict[str, object]] = data # mutable-ok: the caller's request body is sanitized in place + output_config: Final = _fold_output_format_into_output_config( + _as_output_config(request.get("output_config")), request.pop("output_format", None) + ) if output_config is None: - return - if not isinstance(output_config, dict): - data.pop("output_config", None) + request.pop("output_config", None) return drop_keys: Final = set(VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS) @@ -68,6 +92,6 @@ def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: sanitized: Final = {k: v for k, v in output_config.items() if k not in drop_keys} if sanitized: - data["output_config"] = sanitized + request["output_config"] = sanitized else: - data.pop("output_config", None) + request.pop("output_config", None) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 2fab6f438f6..72c08c6b8d7 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -159,7 +159,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): """ Override parent method so VertexAI uses tool-based structured outputs unless the vertex map entry advertises native structured output - (``output_format``, which Vertex AI Claude forwards for those models). + (sent to Vertex AI Claude as ``output_config.format``). """ from litellm.llms.anthropic.common_utils import AnthropicModelInfo 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..6080840a5d2 --- /dev/null +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -0,0 +1,1920 @@ +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 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" +_UNMAPPED_MODEL: Final = "claude-scripted-9" +_BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" +_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"], +} +_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} +_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", +) +_STREAM_FAMILY: Final = ("anthropic_unmapped", "azure_ai_unmapped", "bedrock_invoke", "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 _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 _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 scenario.model( + model=f"bedrock/invoke/{_BEDROCK_INVOKE_MODEL}", + api_base=api_base, + api_key=_BEDROCK_TOKEN, + aws_region_name="us-east-1", + ) + 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 + expected_frames: Final = 1 if shape in _HELD_STREAM_SHAPES else _fragment_count(shape, marker_hex) + assert len(frames) == expected_frames, (len(frames), expected_frames, 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: + expected_frames: Final = 1 if unwrapped else _fragment_count(shape, marker_hex) + assert len(_content_frames(response)) == expected_frames, 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: + expected_frames: Final = 1 if unwrapped else _fragment_count(shape, marker_hex) + assert len(_content_frames(response)) == expected_frames, 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", _STREAM_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("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 + + +@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") + with gateway.scenario() as scenario: + with wire_server(_failing(status)) as failing: + port: Final = int(failing.url.rsplit(":", 1)[1]) + 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)) + + with wire_server(_peer) as wire: + port: Final = int(wire.url.rsplit(":", 1)[1]) + 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/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 808afb29878..ad5c8301f14 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -13,7 +13,8 @@ import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call +from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call +from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, @@ -315,9 +316,7 @@ def test_streaming_thinking_blocks_are_replayable_after_signature_delta(): assert reasoning_content == "Step 1. Step 2." assert thinking_blocks == (*expected_delta_blocks, expected_signature_block) assert "".join(block.get("thinking") or "" for block in thinking_blocks) == reasoning_content - assert parsed_chunks[1].choices[0].delta.provider_specific_fields == { - "thinking_blocks": [expected_delta_blocks[0]] - } + assert parsed_chunks[1].choices[0].delta.provider_specific_fields == {"thinking_blocks": [expected_delta_blocks[0]]} assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == { "thinking_blocks": [expected_signature_block] } @@ -356,7 +355,11 @@ def test_streamed_signed_thinking_round_trips_to_the_next_turn_once(): "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'}, }, {"type": "content_block_stop", "index": 1}, - {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 30}}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 30}, + }, {"type": "message_stop"}, ) user_message: Final = {"role": "user", "content": "What's the weather in Paris?"} @@ -843,13 +846,9 @@ async def test_anthropic_stream_options_return_usage_after_the_finish_chunk( }, {"type": "message_stop"}, ) - body: Final = "".join( - f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events - ) + body: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) route: Final = respx.post("https://api.anthropic.com/v1/messages").mock( - return_value=httpx.Response( - 200, text=body, headers={"content-type": "text/event-stream"} - ) + return_value=httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) ) chunks: Final = await _collect_anthropic_stream( sync_mode, @@ -860,9 +859,7 @@ async def test_anthropic_stream_options_return_usage_after_the_finish_chunk( stream_options={"include_usage": True}, ) finish_indices: Final = tuple( - index - for index, chunk in enumerate(chunks) - if chunk.choices and chunk.choices[0].finish_reason is not None + index for index, chunk in enumerate(chunks) if chunk.choices and chunk.choices[0].finish_reason is not None ) finish_reasons: Final = tuple( chunk.choices[0].finish_reason @@ -870,15 +867,11 @@ async def test_anthropic_stream_options_return_usage_after_the_finish_chunk( if chunk.choices and chunk.choices[0].finish_reason is not None ) usage_indices: Final = tuple( - index - for index, chunk in enumerate(chunks) - if getattr(chunk, "usage", None) is not None + index for index, chunk in enumerate(chunks) if getattr(chunk, "usage", None) is not None ) usage_chunk: Final = chunks[usage_indices[0]] - assert "".join( - chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices - ) == "Blue sky." + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "Blue sky." assert finish_reasons == ("stop",) assert finish_indices == (len(chunks) - 2,) assert usage_indices == (len(chunks) - 1,) @@ -1460,17 +1453,11 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing(): unlike Vertex which already deferred parsing until the buffer could close. A fragment that can't close a JSON value must not trigger a decode attempt. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) iterator.chunk_type = "accumulated_json" - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: - result = iterator._handle_accumulated_json_chunk( - '{"type":"content_block_delta","index":0,"delta":' - ) + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: + result = iterator._handle_accumulated_json_chunk('{"type":"content_block_delta","index":0,"delta":') assert result is None assert spy.call_count == 0, "incomplete buffer should not be parsed" @@ -1482,21 +1469,15 @@ def test_accumulated_json_does_not_reparse_every_fragment(): fragment. """ text = "x" * 200_000 - blob = json.dumps( - {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}} - ) + blob = json.dumps({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}) fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) iterator.chunk_type = "accumulated_json" parsed = None - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: for fragment in fragments: out = iterator._handle_accumulated_json_chunk(fragment) if out is not None: @@ -1520,9 +1501,7 @@ def test_accumulated_json_concatenated_envelopes_do_not_wedge(): and keeps the remainder, so both values surface across two calls. """ obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) iterator.chunk_type = "accumulated_json" first = iterator._handle_accumulated_json_chunk(obj + obj) @@ -1543,9 +1522,7 @@ def test_accumulated_json_heuristic_passes_but_value_still_incomplete(): heuristic must let the parse attempt through, and pop_next_value finding nothing must propagate as None rather than raising. """ - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) iterator.chunk_type = "accumulated_json" result = iterator._handle_accumulated_json_chunk('{"type": {"nested": 1}') @@ -1559,9 +1536,7 @@ def test_accumulated_json_setter_and_sync_end_of_stream_drain(): underlying stream ends, instead of being silently dropped. """ obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False) iterator.chunk_type = "accumulated_json" iterator.accumulated_json = obj # exercises the setter @@ -1576,9 +1551,7 @@ def test_accumulated_json_async_end_of_stream_drain(): import asyncio obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' - iterator = ModelResponseIterator( - streaming_response=MagicMock(), sync_stream=False, json_mode=False - ) + iterator = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=False, json_mode=False) iterator.chunk_type = "accumulated_json" iterator.accumulated_json = obj mock_async_iterator = MagicMock() @@ -1772,7 +1745,9 @@ def test_web_search_calls_are_cumulative_through_incomplete_search(): ) assert list(first_start.choices[0].delta.provider_specific_fields["web_search_calls"]) == ["srvtoolu_A"] - assert first_result.choices[0].delta.provider_specific_fields["web_search_calls"]["srvtoolu_A"].status == "completed" + assert ( + first_result.choices[0].delta.provider_specific_fields["web_search_calls"]["srvtoolu_A"].status == "completed" + ) calls = second_start.choices[0].delta.provider_specific_fields["web_search_calls"] assert list(calls) == ["srvtoolu_A", "srvtoolu_B"] assert calls["srvtoolu_A"].status == "completed" @@ -2651,9 +2626,7 @@ class TestAnthropicChatCompletionPreCallLogging: calls = {"pre_call": []} logging_obj = MagicMock() logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs) - with patch.object( - AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []} - ): + with patch.object(AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}): try: AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj)) except Exception: @@ -2739,6 +2712,272 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): assert assembled._hidden_params["provider_response_model"] == served_model +_STREAMED_RISK_SCHEMA: Final = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], +} +_STREAMED_RISK_OBJECT: Final = {"tier": 3, "summary": "Two-hour read-only window on the loyalty APIs"} +_STREAMED_CONTACT_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}, "email": {"type": "string"}}, + "additionalProperties": False, +} +_STREAMED_BOOK_SCHEMA: Final = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "parameters": { + "type": "object", + "properties": {"title": {"type": "string"}, "author": {"type": "string"}}, + "required": ["title", "author"], + }, + }, + "required": ["name", "parameters"], + "additionalProperties": False, +} +_STREAMED_CONTACT: Final = {"name": "Ada", "email": "ada@example.com"} +_STREAMED_PERSON_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}, + "required": ["name", "age"], +} +_STREAMED_PERSON: Final = {"name": "Ada", "age": 36} + + +def _json_tool_events(fragments: tuple[str, ...]) -> tuple[dict[str, object], ...]: + return ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-emulated-structured-output", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {}}, + }, + *( + {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": fragment}} + for fragment in fragments + ), + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 20}, + }, + ) + + +def _streamed_json_mode_contents( + fragments: tuple[str, ...], schema: dict[str, object] = _STREAMED_RISK_SCHEMA +) -> tuple[str, ...]: + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + json_mode=True, + response_format_schema=response_format_schema(schema), + ) + chunks = tuple(iterator.chunk_parser(event) for event in _json_tool_events(fragments)) + assert all(chunk.choices[0].delta.tool_calls is None for chunk in chunks) + assert chunks[-1].choices[0].finish_reason == "stop" + return tuple(chunk.choices[0].delta.content or "" for chunk in chunks) + + +def test_streamed_json_mode_envelope_is_unwrapped_when_the_block_ends(): + envelope = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _STREAMED_RISK_OBJECT}) + fragments = (envelope[:9], envelope[9:40], envelope[40:]) + + contents = _streamed_json_mode_contents(fragments) + + assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT + assert contents[:-2] == ("",) * (len(contents) - 2) + + +def test_streamed_json_mode_schema_object_streams_as_it_arrives(): + answer = json.dumps(_STREAMED_RISK_OBJECT) + fragments = (answer[:4], answer[4:12], answer[12:]) + + contents = _streamed_json_mode_contents(fragments) + + assert "".join(contents) == answer + assert contents[3] == answer[:12] + assert contents[4] == answer[12:] + assert contents[5] == "" + + +def test_streamed_json_mode_truncated_output_is_passed_through_untouched(): + contents = _streamed_json_mode_contents(('{"na', 'me": "json_tool_call", "parameters": {"tier": 3')) + + assert "".join(contents) == '{"name": "json_tool_call", "parameters": {"tier": 3' + + +def test_streamed_json_mode_tool_name_envelope_is_unwrapped_when_name_is_a_schema_property(): + envelope = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _STREAMED_CONTACT}) + + contents = _streamed_json_mode_contents((envelope[:6], envelope[6:24], envelope[24:]), _STREAMED_CONTACT_SCHEMA) + + assert json.loads("".join(contents)) == _STREAMED_CONTACT + assert contents[:-2] == ("",) * (len(contents) - 2) + + +def test_streamed_json_mode_name_property_answer_streams_once_its_name_value_is_known(): + answer = json.dumps(_STREAMED_CONTACT) + + contents = _streamed_json_mode_contents((answer[:6], answer[6:24], answer[24:]), _STREAMED_CONTACT_SCHEMA) + + assert "".join(contents) == answer + assert contents[2] == "" + assert contents[3] == answer[:24] + assert contents[4] == answer[24:] + + +def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close(): + answer = json.dumps({"name": "Ada"}) + + contents = _streamed_json_mode_contents((answer[:11], answer[11:]), _STREAMED_CONTACT_SCHEMA) + + assert "".join(contents) == answer + assert contents[2] == "" + assert contents[3] == answer + + +@pytest.mark.parametrize( + "schema, answer", + [ + pytest.param( + {"type": "object", "additionalProperties": {"type": "integer"}}, + {"alpha": 1, "beta": 2, "gamma": 3}, + id="schema_without_properties", + ), + pytest.param( + _STREAMED_RISK_SCHEMA, + {"rev": 2, **_STREAMED_RISK_OBJECT}, + id="undeclared_first_key_holding_a_number", + ), + pytest.param( + {"type": "object", "properties": {"score": {"type": "integer"}}}, + {"note": "read-only window", "meta": {"score": 7}}, + id="string_first_key_on_a_schema_without_required_keys", + ), + pytest.param( + {"type": "object", "properties": {"summary": {"type": "string"}}, "additionalProperties": {"type": "object"}}, + {"user_1": {"summary": "likes tea", "visits": 3}}, + id="undeclared_key_the_schema_allows", + ), + ], +) +def test_streamed_json_mode_answer_that_cannot_be_an_envelope_streams_as_it_arrives( + schema: dict[str, object], answer: dict[str, object] +): + encoded = json.dumps(answer) + + contents = _streamed_json_mode_contents((encoded[:12], encoded[12:24], encoded[24:]), schema) + + assert "".join(contents) == encoded + assert contents[2:5] == (encoded[:12], encoded[12:24], encoded[24:]) + + +def test_streamed_json_mode_string_label_answer_streams_once_the_next_key_is_declared(): + answer = json.dumps({"note": "read-only window", **_STREAMED_RISK_OBJECT}) + + contents = _streamed_json_mode_contents((answer[:12], answer[12:40], answer[40:])) + + assert "".join(contents) == answer + assert contents[2] == "" + assert contents[3] == answer[:40] + assert contents[4] == answer[40:] + + +@pytest.mark.parametrize( + "envelope", + [ + pytest.param( + {"parameter_name": "RiskAssessment", "parameter_value": _STREAMED_RISK_OBJECT}, id="schema_name_label" + ), + pytest.param( + {"$FUNCTION_NAME": RESPONSE_FORMAT_TOOL_NAME, "PARAMETER_VALUE": _STREAMED_RISK_OBJECT}, + id="foreign_tool_name_key", + ), + ], +) +def test_streamed_json_mode_string_label_envelope_is_unwrapped_when_the_block_ends(envelope: dict[str, object]): + encoded = json.dumps(envelope) + + contents = _streamed_json_mode_contents((encoded[:12], encoded[12:40], encoded[40:])) + + assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT + assert contents[:-2] == ("",) * (len(contents) - 2) + + +def test_streamed_json_mode_two_level_envelope_is_unwrapped_when_the_block_ends(): + envelope = json.dumps({"parameter": {"value": _STREAMED_RISK_OBJECT}}) + + contents = _streamed_json_mode_contents((envelope[:9], envelope[9:30], envelope[30:])) + + assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT + assert contents[:-2] == ("",) * (len(contents) - 2) + + +def test_streamed_json_mode_accepted_answer_naming_the_tool_is_kept_whole(): + answer = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": {"title": "Dune", "author": "Frank Herbert"}}) + + contents = _streamed_json_mode_contents((answer[:9], answer[9:40], answer[40:]), _STREAMED_BOOK_SCHEMA) + + assert json.loads("".join(contents)) == json.loads(answer) + assert contents[:-2] == ("",) * (len(contents) - 2) + + +@pytest.mark.parametrize( + "answer, schema", + [ + pytest.param( + '{"name":"json_tool_call","parameters":{"title":"Dune","author":"Frank Herbert"}}', + _STREAMED_BOOK_SCHEMA, + id="answer_naming_the_tool", + ), + pytest.param('{"city":"café 東京","country":"France"}', _STREAMED_RISK_SCHEMA, id="string_labels_only"), + ], +) +def test_streamed_json_mode_held_answer_without_an_envelope_keeps_the_model_text( + answer: str, schema: dict[str, object] +): + contents = _streamed_json_mode_contents((answer[:9], answer[9:30], answer[30:]), schema) + + assert "".join(contents) == answer + assert contents[:-2] == ("",) * (len(contents) - 2) + + +def test_envelope_verdict_holds_a_prefix_that_stays_undecided_past_the_limit(): + schema = response_format_schema(_STREAMED_RISK_SCHEMA) + assert schema is not None + never_closing_key = '{"' + "k" * 5000 + + assert _envelope_verdict(never_closing_key[:100], schema, RESPONSE_FORMAT_TOOL_NAME) == "undecided" + assert _envelope_verdict(never_closing_key, schema, RESPONSE_FORMAT_TOOL_NAME) == "hold" + + +def test_streamed_json_mode_answer_with_a_foreign_name_matches_the_non_streamed_answer(): + answer = {"name": "Person", "parameters": _STREAMED_PERSON} + encoded = json.dumps(answer) + schema = response_format_schema(_STREAMED_PERSON_SCHEMA) + assert schema is not None + + contents = _streamed_json_mode_contents((encoded[:16], encoded[16:34], encoded[34:]), _STREAMED_PERSON_SCHEMA) + + assert "".join(contents) == encoded + assert unwrap_response_format_envelope(answer, schema, RESPONSE_FORMAT_TOOL_NAME) == answer + + @respx.mock def test_anthropic_streaming_function_call_arguments_are_assembled( monkeypatch: pytest.MonkeyPatch, @@ -2784,9 +3023,7 @@ def test_anthropic_streaming_function_call_arguments_are_assembled( }, {"type": "message_stop"}, ) - body: Final = "".join( - f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events - ) + body: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) respx.post("https://api.anthropic.com/v1/messages").mock( return_value=httpx.Response( 200, @@ -2820,18 +3057,12 @@ def test_anthropic_streaming_function_call_arguments_are_assembled( ) chunks: Final = tuple(stream) tool_call_deltas: Final = tuple( - chain.from_iterable( - chunk.choices[0].delta.tool_calls or () for chunk in chunks if chunk.choices - ) + chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks if chunk.choices) ) - assert tuple( - tool_call.function.name - for tool_call in tool_call_deltas - if tool_call.function.name is not None - ) == ("get_current_weather",) - arguments: Final = "".join( - tool_call.function.arguments or "" for tool_call in tool_call_deltas + assert tuple(tool_call.function.name for tool_call in tool_call_deltas if tool_call.function.name is not None) == ( + "get_current_weather", ) + arguments: Final = "".join(tool_call.function.arguments or "" for tool_call in tool_call_deltas) assert json.loads(arguments) == {"location": "Boston", "unit": "fahrenheit"} assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=messages) assert assembled.choices[0].message.tool_calls[0].function.name == "get_current_weather" 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 e8f3399d4cc..b682d8661f3 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messag from litellm.llms.anthropic.chat import ModelResponseIterator from litellm.llms.anthropic.common_utils import process_anthropic_headers from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.base_llm.base_utils import ResponseFormatSchema, response_format_schema from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -36,7 +37,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im ) from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk -from litellm.types.utils import ChatCompletionToolCallChunk, ServerToolUse, Usage +from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage def test_response_format_transformation_unit_test(): @@ -7295,6 +7296,412 @@ def test_convert_tool_response_to_message_no_arguments(): assert message is None +_RISK_SCHEMA: Final = { + "$defs": {"Owner": {"title": "Owner", "type": "object", "properties": {"team": {"type": "string"}}}}, + "title": "RiskAssessment", + "type": "object", + "properties": { + "tier": {"title": "Tier", "type": "integer"}, + "summary": {"title": "Summary", "type": "string"}, + "owner": {"$ref": "#/$defs/Owner"}, + }, + "required": ["tier", "summary"], +} +_RISK_OBJECT: Final = {"tier": 3, "summary": "Two-hour read-only window on the loyalty APIs"} +_PERSON_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}, + "required": ["name", "age"], +} +_OPTIONAL_FIELDS_SCHEMA: Final = { + "type": "object", + "properties": {"note": {"type": "string"}, "score": {"type": "integer"}}, +} +_TREE_SCHEMA: Final = { + "type": "object", + "properties": {"value": {"type": "integer"}, "child": {"type": "object"}}, + "required": ["value"], +} +_OPEN_SCORE_SCHEMA: Final = { + "type": "object", + "properties": {"score": {"type": "integer"}}, + "additionalProperties": True, +} +_VALUES_SCHEMA: Final = {"type": "object", "properties": {"values": {"type": "array"}}, "required": ["values"]} +_REMOTE_REF_SCHEMA: Final = {**_PERSON_SCHEMA, "$ref": "https://schemas.example.com/person.json"} +_LOOSE_SCHEMA: Final = {"type": "object", "properties": {"name": {"type": "string"}, "email": {"type": "string"}}} +_CLOSED_CONTACT_SCHEMA: Final = {**_LOOSE_SCHEMA, "additionalProperties": False} +_CONTACT: Final = {"name": "Ada", "email": "ada@example.com"} +_NAMED_BOOK_SCHEMA: Final = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "parameters": { + "type": "object", + "properties": {"title": {"type": "string"}, "author": {"type": "string"}}, + "required": ["title", "author"], + "additionalProperties": False, + }, + }, + "required": ["name", "parameters"], + "additionalProperties": False, +} +_NAMED_BOOK: Final = {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": {"title": "Dune", "author": "Frank Herbert"}} +_RECURSIVE_NODE_SCHEMA: Final = { + "type": "object", + "properties": { + "id": {"type": "string"}, + "title": {"type": "string"}, + "parent": {"anyOf": [{"$ref": "#"}, {"type": "null"}]}, + }, + "required": ["id", "title"], +} +_CHILD_WITHOUT_TITLE: Final = {"id": "child-7", "parent": {"id": "root-1", "title": "Root", "parent": None}} +_JSON_OBJECT_MODE_SCHEMA: Final = {"type": "object", "properties": {}, "additionalProperties": True} +_SUMMARY_PROPERTIES: Final = {"summary": {"type": "string"}} +_PER_USER_SCHEMA: Final = { + "type": "object", + "properties": _SUMMARY_PROPERTIES, + "additionalProperties": {"type": "object"}, +} +_PER_USER_PATTERN_SCHEMA: Final = { + "type": "object", + "properties": _SUMMARY_PROPERTIES, + "patternProperties": {"^user_": {"type": "object"}}, +} +_PER_USER_ANSWER: Final = {"user_1": {"summary": "likes tea", "visits": 3}} + + +def _schema(document: dict[str, object]) -> ResponseFormatSchema: + schema = response_format_schema(document) + assert schema is not None + return schema + + +def _json_tool_call(arguments: object) -> list[ChatCompletionToolCallChunk]: + return [ + ChatCompletionToolCallChunk( + id="toolu_1", + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=RESPONSE_FORMAT_TOOL_NAME, arguments=json.dumps(arguments) + ), + index=0, + ) + ] + + +@pytest.mark.parametrize( + "schema, tool_input, expected_content", + [ + pytest.param( + _RISK_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _RISK_OBJECT}, + _RISK_OBJECT, + id="name_and_parameters_envelope", + ), + pytest.param(_RISK_SCHEMA, {"input": _RISK_OBJECT}, _RISK_OBJECT, id="input_wrapper"), + pytest.param(_RISK_SCHEMA, {"arguments": _RISK_OBJECT}, _RISK_OBJECT, id="arguments_wrapper"), + pytest.param( + _PERSON_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": {"name": "Ada", "age": 36}}, + {"name": "Ada", "age": 36}, + id="envelope_key_shared_with_a_schema_property", + ), + pytest.param( + _OPTIONAL_FIELDS_SCHEMA, + {"input": {"score": 7}}, + {"score": 7}, + id="wrapper_around_an_all_optional_schema", + ), + pytest.param( + _RISK_SCHEMA, + {"parameter": {"value": _RISK_OBJECT}}, + _RISK_OBJECT, + id="two_level_envelope", + ), + pytest.param( + _RISK_SCHEMA, + {"parameter_name": "RiskAssessment", "parameter_value": _RISK_OBJECT}, + _RISK_OBJECT, + id="string_label_beside_the_wrapper", + ), + pytest.param( + _RISK_SCHEMA, + {"$FUNCTION_NAME": RESPONSE_FORMAT_TOOL_NAME, "PARAMETER_VALUE": _RISK_OBJECT}, + _RISK_OBJECT, + id="foreign_tool_name_key_beside_the_wrapper", + ), + pytest.param( + _RISK_SCHEMA, + {"note": "read-only window", **_RISK_OBJECT}, + {"note": "read-only window", **_RISK_OBJECT}, + id="string_label_beside_schema_keys", + ), + pytest.param( + {**_RISK_SCHEMA, "required": "tier"}, + {"parameter_name": "RiskAssessment", "parameter_value": _RISK_OBJECT}, + {"parameter_name": "RiskAssessment", "parameter_value": _RISK_OBJECT}, + id="string_label_when_required_is_malformed", + ), + pytest.param( + _RISK_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "arguments": json.dumps(_RISK_OBJECT)}, + _RISK_OBJECT, + id="tool_name_envelope_holding_encoded_json", + ), + pytest.param( + _RISK_SCHEMA, + {"tier": 1, "parameters": _RISK_OBJECT}, + {"tier": 1, "parameters": _RISK_OBJECT}, + id="schema_key_beside_a_wrapper_key", + ), + pytest.param( + _RECURSIVE_NODE_SCHEMA, + _CHILD_WITHOUT_TITLE, + _CHILD_WITHOUT_TITLE, + id="recursive_schema_keeps_the_answer_over_its_complete_parent", + ), + pytest.param( + {**_RISK_SCHEMA, "$schema": 5}, + {"input": _RISK_OBJECT}, + _RISK_OBJECT, + id="non_string_dollar_schema", + ), + pytest.param( + _OPEN_SCORE_SCHEMA, + {"meta": {"score": 7}, "label": "keep"}, + {"meta": {"score": 7}, "label": "keep"}, + id="valid_answer_with_extra_keys_and_a_schema_shaped_nested_object", + ), + pytest.param(_VALUES_SCHEMA, {"values": [1, 2]}, {"values": [1, 2]}, id="values_is_a_schema_property"), + pytest.param( + _REMOTE_REF_SCHEMA, + {"input": {"name": "Ada", "age": 36}}, + {"name": "Ada", "age": 36}, + id="remote_ref_schema_falls_back_to_key_matching", + ), + pytest.param( + _CLOSED_CONTACT_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _CONTACT}, + _CONTACT, + id="tool_name_envelope_when_the_schema_lists_no_required_keys", + ), + pytest.param( + _LOOSE_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _CONTACT}, + _CONTACT, + id="tool_name_envelope_on_an_open_schema", + ), + pytest.param(_NAMED_BOOK_SCHEMA, _NAMED_BOOK, _NAMED_BOOK, id="accepted_answer_whose_name_is_the_tool_name"), + pytest.param( + _LOOSE_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "email": "ada@example.com"}, + {"name": RESPONSE_FORMAT_TOOL_NAME, "email": "ada@example.com"}, + id="tool_name_value_without_a_nested_object_stays", + ), + pytest.param(_RISK_SCHEMA, _RISK_OBJECT, _RISK_OBJECT, id="plain_schema_object"), + pytest.param( + _TREE_SCHEMA, + {"value": 1, "child": {"value": 2}}, + {"value": 1, "child": {"value": 2}}, + id="schema_object_whose_property_holds_a_schema_shaped_object", + ), + pytest.param(_RISK_SCHEMA, {"input": {"tier": 3}}, {"tier": 3}, id="inner_object_missing_required_key"), + pytest.param(_RISK_SCHEMA, {"values": {"tier": 2}}, {"tier": 2}, id="values_wrapper_missing_required_key"), + pytest.param( + _RISK_SCHEMA, + {"first": _RISK_OBJECT, "second": _RISK_OBJECT}, + {"first": _RISK_OBJECT, "second": _RISK_OBJECT}, + id="two_candidate_objects", + ), + pytest.param( + _OPTIONAL_FIELDS_SCHEMA, + {"meta": {"source": "change-request"}}, + {"meta": {"source": "change-request"}}, + id="inner_object_without_schema_keys", + ), + pytest.param( + _JSON_OBJECT_MODE_SCHEMA, + {"input": _RISK_OBJECT}, + {"input": _RISK_OBJECT}, + id="json_object_mode_without_schema_properties", + ), + pytest.param( + _JSON_OBJECT_MODE_SCHEMA, + {"values": _RISK_OBJECT}, + _RISK_OBJECT, + id="json_object_mode_values_wrapper", + ), + pytest.param( + _JSON_OBJECT_MODE_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _RISK_OBJECT}, + _RISK_OBJECT, + id="json_object_mode_tool_name_envelope", + ), + pytest.param(_PER_USER_SCHEMA, _PER_USER_ANSWER, _PER_USER_ANSWER, id="key_allowed_by_additional_properties"), + pytest.param( + _PER_USER_PATTERN_SCHEMA, _PER_USER_ANSWER, _PER_USER_ANSWER, id="key_allowed_by_pattern_properties" + ), + pytest.param(_OPEN_SCORE_SCHEMA, {"input": {"score": 7}}, {"input": {"score": 7}}, id="key_allowed_by_true"), + pytest.param( + {**_PER_USER_SCHEMA, "required": ["summary"]}, + {"input": {"summary": "likes tea"}}, + {"summary": "likes tea"}, + id="allowed_key_on_a_schema_with_required_keys", + ), + pytest.param( + _PER_USER_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": {"summary": "likes tea"}}, + {"summary": "likes tea"}, + id="tool_name_envelope_on_a_schema_allowing_extra_keys", + ), + ], +) +def test_convert_tool_response_to_message_unwraps_only_an_unambiguous_schema_envelope( + schema: dict, tool_input: dict, expected_content: dict +): + message = AnthropicConfig.convert_tool_response_to_message( + tool_calls=_json_tool_call(tool_input), + response_format_schema=_schema(schema), + ) + + assert message is not None + assert json.loads(message.content) == expected_content + + +_TOO_DEEP_TO_DECODE: Final = "[" * 100_000 + "]" * 100_000 + + +@pytest.mark.parametrize( + "arguments", + [json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "arguments": _TOO_DEEP_TO_DECODE}), _TOO_DEEP_TO_DECODE], + ids=["encoded_inner_object", "whole_arguments"], +) +def test_convert_tool_response_to_message_keeps_arguments_too_deep_to_decode(arguments: str): + message = AnthropicConfig.convert_tool_response_to_message( + tool_calls=[ + ChatCompletionToolCallChunk( + id="toolu_1", + type="function", + function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, arguments=arguments), + index=0, + ) + ], + response_format_schema=_schema(_RISK_SCHEMA), + ) + + assert message is not None + assert message.content == arguments + + +@pytest.mark.parametrize( + "document", + [ + {"type": "object", "properties": ["tier", "summary"]}, + {"type": "object", "properties": "tier"}, + ["tier", "summary"], + ], + ids=["properties_list", "properties_string", "schema_list"], +) +def test_response_format_schema_rejects_a_malformed_schema(document: object): + assert response_format_schema(document) is None + + +def test_json_mode_with_user_tool_call_unwraps_the_envelope_into_content(): + user_tool_call: Final = ChatCompletionToolCallChunk( + id="toolu_user", + type="function", + function=ChatCompletionToolCallFunctionChunk(name="lookup_change_window", arguments='{"system": "ledger"}'), + index=1, + ) + + replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming( + json_mode=True, + tool_calls=[*_json_tool_call({"input": _RISK_OBJECT}), user_tool_call], + response_format_schema=_schema(_RISK_SCHEMA), + ) + + assert replacement is None + assert remaining_tool_calls == [user_tool_call] + assert extra_content is not None + assert json.loads(extra_content) == _RISK_OBJECT + + +def test_vertex_claude_tool_emulated_response_format_returns_the_object_the_model_wrapped(): + config = VertexAIAnthropicConfig() + optional_params = config.map_openai_params( + non_default_params={ + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_change_window", + "parameters": { + "type": "object", + "properties": {"system": {"type": "string"}}, + "required": ["system"], + }, + }, + } + ], + "response_format": { + "type": "json_schema", + "json_schema": {"name": "RiskAssessment", "schema": _RISK_SCHEMA}, + }, + }, + optional_params={}, + model="claude-emulated-structured-output", + drop_params=False, + ) + json_mode = optional_params.pop("json_mode") + request = config.transform_request( + model="claude-emulated-structured-output", + messages=[{"role": "user", "content": "Assess this change request"}], + optional_params={**optional_params, "is_vertex_request": True}, + litellm_params={}, + headers={}, + ) + json_tool = next(tool for tool in request["tools"] if tool["name"] == RESPONSE_FORMAT_TOOL_NAME) + assert json_tool["input_schema"]["required"] == ["tier", "summary"] + assert "output_format" not in request + assert "output_config" not in request + completion_response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-emulated-structured-output", + "content": [ + { + "type": "tool_use", + "id": "toolu_1", + "name": RESPONSE_FORMAT_TOOL_NAME, + "input": {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _RISK_OBJECT}, + } + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 30}, + } + + result = config.transform_response( + model="claude-emulated-structured-output", + raw_response=httpx.Response(200, json=completion_response), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data=request, + messages=[{"role": "user", "content": "Assess this change request"}], + optional_params=optional_params, + litellm_params={}, + encoding=None, + json_mode=json_mode, + ) + + assert json.loads(result.choices[0].message.content) == _RISK_OBJECT + assert result.choices[0].message.tool_calls is None + assert result.choices[0].finish_reason == "stop" + + def test_anthropic_tool_with_image(): import json @@ -7606,9 +8013,7 @@ async def test_anthropic_api_max_completion_tokens(model: str): print("request_body: ", request_body) assert request_body == { - "messages": [ - {"role": "user", "content": [{"type": "text", "text": "Hello!"}]} - ], + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}], "max_tokens": 10, "model": model.split("/")[-1], } @@ -7826,9 +8231,7 @@ async def test_anthropic_api_prompt_caching_with_content_str(): "cache_control": {"type": "ephemeral"}, }, ] - translated_system_message = litellm.AnthropicConfig().translate_system_message( - messages=system_message - ) + translated_system_message = litellm.AnthropicConfig().translate_system_message(messages=system_message) assert translated_system_message == [ # System Message @@ -7898,9 +8301,9 @@ async def test_anthropic_api_prompt_caching_with_content_str(): assert len(translated_messages) == len(expected_messages) for idx, i in enumerate(translated_messages): - assert ( - i == expected_messages[idx] - ), "Error on idx={}. Got={}, Expected={}".format(idx, i, expected_messages[idx]) + assert i == expected_messages[idx], "Error on idx={}. Got={}, Expected={}".format( + idx, i, expected_messages[idx] + ) @pytest.fixture diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 45a222096b7..6fb5c746354 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -20,8 +20,10 @@ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, AWSEventStreamDecoder, + _stream_decoder, make_call, make_sync_call, + response_format_schema_from_request, ) from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_stream_event_statuses from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -1618,6 +1620,91 @@ def test_invoke_decoder_reports_only_missing_aws_dependency(missing): assert error.value is failure +_INVOKE_RISK_SCHEMA: Final = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], +} +_INVOKE_RISK_OBJECT: Final = {"tier": 2, "summary": "Weekend database migration with a tested rollback"} + + +def _invoke_json_mode_request(tools: Sequence[Mapping[str, object]]) -> str: + return json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "messages": [{"role": "user", "content": "Assess this change request"}], + "tools": list(tools), + "tool_choice": {"type": "tool", "name": "json_tool_call"}, + } + ) + + +def _invoke_json_tool_events(fragments: Sequence[str]) -> tuple[dict[str, object], ...]: + deltas = tuple( + {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": fragment}} + for fragment in fragments + ) + return ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 1, "output_tokens": 0}}}, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "t1", "name": "json_tool_call", "input": {}}, + }, + *deltas, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 9}}, + ) + + +def _invoke_streamed_json_mode_contents(data: str, fragments: Sequence[str], json_mode: bool = True) -> tuple[str, ...]: + decoder = _stream_decoder( + "anthropic", + model="anthropic.claude-sonnet-4", + json_mode=True, + sync_stream=True, + response_format_schema=response_format_schema_from_request(data, json_mode), + ) + chunks = tuple(decoder._chunk_parser(event) for event in _invoke_json_tool_events(fragments)) + assert all(not chunk.choices[0].delta.tool_calls for chunk in chunks) + return tuple(chunk.choices[0].delta.content or "" for chunk in chunks) + + +def test_invoke_streamed_json_mode_envelope_is_unwrapped_with_the_schema_from_the_request_body(): + envelope: Final = json.dumps({"tool_name": "json_tool_call", "parameters": _INVOKE_RISK_OBJECT}) + data: Final = _invoke_json_mode_request([{"name": "json_tool_call", "input_schema": _INVOKE_RISK_SCHEMA}]) + + contents = _invoke_streamed_json_mode_contents(data, (envelope[:17], envelope[17:48], envelope[48:])) + + assert json.loads("".join(contents)) == _INVOKE_RISK_OBJECT + + +@pytest.mark.parametrize( + ("data", "json_mode"), + [ + pytest.param("not json", True, id="unparseable_body"), + pytest.param(json.dumps([1, 2]), True, id="body_is_not_an_object"), + pytest.param( + _invoke_json_mode_request([{"name": "other_tool", "input_schema": _INVOKE_RISK_SCHEMA}]), + True, + id="no_json_tool", + ), + pytest.param( + _invoke_json_mode_request([{"name": "json_tool_call", "input_schema": _INVOKE_RISK_SCHEMA}]), + False, + id="json_mode_off", + ), + ], +) +def test_invoke_stream_keeps_the_wrapper_when_the_request_body_carries_no_json_tool_schema(data: str, json_mode: bool): + envelope: Final = json.dumps({"tool_name": "json_tool_call", "parameters": _INVOKE_RISK_OBJECT}) + + assert response_format_schema_from_request(data, json_mode) is None + contents = _invoke_streamed_json_mode_contents(data, (envelope[:17], envelope[17:]), json_mode) + + assert json.loads("".join(contents)) == {"tool_name": "json_tool_call", "parameters": _INVOKE_RISK_OBJECT} + + @pytest.mark.asyncio async def test_converse_stream_thinking_merges_into_content_with_think_tags() -> None: frames: Final = ( diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 1452f77604f..113ddcfe2a9 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -514,6 +514,29 @@ def test_messages_request_strips_effort_for_haiku_45(): assert opus_result["output_config"] == {"effort": "high"} +def test_messages_request_moves_output_format_into_output_config(): + config = VertexAIPartnerModelsAnthropicMessagesConfig() + output_format: Final = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "integer"}}, "required": ["answer"]}, + } + + result = config.transform_anthropic_messages_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "What is 6 times 7?"}], + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_format": output_format, + "output_config": {"effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_format" not in result + assert result["output_config"] == {"effort": "high", "format": output_format} + + def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): """ Regression test: repeated provider config lookups for the same Vertex Claude model diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index c37b0d3f5bd..c82e6e541c5 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -252,9 +252,8 @@ def test_vertex_ai_claude_sonnet_4_5_structured_output_fix(): This test verifies that: 1. Claude Sonnet 4.5 uses tool-based structured outputs when ``response_format`` is given to the OpenAI-compat path (the path that triggered #18625). - 2. ``output_format`` is forwarded to Vertex AI when present — Vertex now - accepts the field; the prior blanket-strip behavior was the silent drop - of Anthropic Structured Outputs that this PR fixes. + 2. ``output_format`` reaches Vertex AI as ``output_config.format``, since + Vertex rejects the top-level field as deprecated. """ config = VertexAIAnthropicConfig() @@ -340,12 +339,8 @@ def test_vertex_ai_claude_sonnet_4_5_structured_output_fix(): headers={}, ) - # output_format is now forwarded to Vertex (Vertex parity has shifted — - # it accepts the field and uses it to enforce the JSON schema). The - # prior behavior silently stripped it, hiding Structured Outputs from - # callers who explicitly requested them. - assert "output_format" in final_data - assert final_data["output_format"]["type"] == "json_schema" + assert "output_format" not in final_data + assert final_data["output_config"]["format"]["type"] == "json_schema" assert ( "model" not in final_data ), "model is still stripped (Vertex routes by URL)" @@ -735,7 +730,7 @@ def test_vertex_ai_anthropic_output_config_non_dict_dropped(): def test_vertex_ai_anthropic_output_format_and_output_config_effort_preserved(): - """Both ``output_format`` and ``output_config.effort`` are forwarded on Vertex 4.6/4.7.""" + """``output_format`` is folded into ``output_config`` next to ``effort`` on Vertex 4.6/4.7.""" config = VertexAIAnthropicConfig() messages = [{"role": "user", "content": "Extract structured data"}] @@ -782,10 +777,8 @@ def test_vertex_ai_anthropic_output_format_and_output_config_effort_preserved(): headers={}, ) - # output_format flows through unchanged — Vertex AI Claude accepts it. - assert result["output_format"] == output_format - # output_config.effort now flows through (Vertex accepts it on 4.6/4.7). - assert result["output_config"] == {"effort": "high"} + assert "output_format" not in result + assert result["output_config"] == {"effort": "high", "format": output_format} assert result["max_tokens"] == 2048 assert "model" not in result, "model is still stripped (Vertex routes by URL)" finally: @@ -859,6 +852,97 @@ def test_sanitize_strips_effort_for_haiku_45(): assert data["output_config"] == {"effort": "high"} +_STRUCTURED_FORMAT: Final = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}, "required": ["result"]}, +} +_CALLER_FORMAT: Final = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "integer"}}}, +} + + +@pytest.mark.parametrize( + "model, data, expected_output_config", + [ + pytest.param( + "claude-opus-4-6", + {"output_format": _STRUCTURED_FORMAT}, + {"format": _STRUCTURED_FORMAT}, + id="output_format_alone", + ), + pytest.param( + "claude-opus-4-6", + {"output_format": _STRUCTURED_FORMAT, "output_config": {"effort": "high"}}, + {"effort": "high", "format": _STRUCTURED_FORMAT}, + id="merged_with_effort", + ), + pytest.param( + "claude-haiku-4-5@20251001", + {"output_format": _STRUCTURED_FORMAT, "output_config": {"effort": "high"}}, + {"format": _STRUCTURED_FORMAT}, + id="effort_still_dropped_for_haiku_45", + ), + pytest.param( + "claude-opus-4-6", + {"output_format": _STRUCTURED_FORMAT, "output_config": {"format": _CALLER_FORMAT}}, + {"format": _STRUCTURED_FORMAT}, + id="output_format_wins_over_output_config_format", + ), + pytest.param( + "claude-opus-4-6", + {"output_format": _STRUCTURED_FORMAT, "output_config": {"format": None}}, + {"format": _STRUCTURED_FORMAT}, + id="null_output_config_format_counts_as_absent", + ), + pytest.param( + "claude-opus-4-6", + {"output_format": _STRUCTURED_FORMAT, "output_config": "not-a-dict"}, + {"format": _STRUCTURED_FORMAT}, + id="malformed_output_config_replaced", + ), + ], +) +def test_sanitize_folds_output_format_into_output_config(model: str, data: dict, expected_output_config: dict): + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.output_params_utils import ( + sanitize_vertex_anthropic_output_params, + ) + + request: Final = copy.deepcopy(data) + sanitize_vertex_anthropic_output_params(request, model) + + assert "output_format" not in request + assert request["output_config"] == expected_output_config + + +def test_vertex_ai_native_response_format_reaches_vertex_as_output_config_format(local_model_cost_map): + config = VertexAIAnthropicConfig() + schema: Final = {"type": "object", "properties": {"result": {"type": "string"}}, "required": ["result"]} + optional_params: Final = config.map_openai_params( + non_default_params={ + "response_format": {"type": "json_schema", "json_schema": {"name": "result", "schema": schema}}, + "reasoning_effort": "high", + }, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + + request: Final = config.transform_request( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Return the result"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert "output_format" not in request + assert request["output_config"]["format"] == optional_params["output_format"] + assert request["output_config"]["format"]["schema"]["properties"] == schema["properties"] + assert request["output_config"]["effort"] == optional_params["output_config"]["effort"] + assert "tool_choice" not in request + + def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_model_cost_map): """Regression: Fable 5.1 rejects forced tool use, so the vertex map entry advertises native structured output and ``response_format`` must map to