From 5c102bcc8a0e046d238e45ac77e1bcb47ba98ae6 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:20:49 -0700 Subject: [PATCH] fix(anthropic): preserve schema-shaped fallback responses from main --- litellm/llms/anthropic/chat/handler.py | 115 ++++- litellm/llms/anthropic/chat/transformation.py | 75 ++- litellm/llms/azure_ai/anthropic/handler.py | 4 + litellm/llms/base_llm/base_utils.py | 136 +++++- litellm/llms/bedrock/chat/invoke_handler.py | 40 +- .../chat/test_anthropic_chat_handler.py | 292 ++++++++++- .../test_anthropic_chat_transformation.py | 462 +++++++++++++++++- .../llms/bedrock/chat/test_invoke_handler.py | 89 ++++ 8 files changed, 1186 insertions(+), 27 deletions(-) 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 616ac727faa..0650e981d3d 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -7,7 +7,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 @@ -47,7 +47,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 REQUEST_PARAMS, BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_ADVISOR_TOOL_TYPE, @@ -1805,7 +1811,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: @@ -2258,6 +2267,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], @@ -2277,11 +2287,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 @@ -2717,6 +2731,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)) @@ -2774,6 +2789,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: @@ -2885,12 +2901,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 @@ -2898,28 +2916,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=response_format_schema) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return AnthropicError( @@ -2929,6 +2951,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 64a10f777a0..9256ae20a61 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, JsonValue +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk @@ -159,6 +161,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/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 6d7d62e47aa..89fa884794a 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -11,7 +11,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, @@ -2367,3 +2368,292 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): assert chunk._hidden_params["provider_response_model"] == served_model assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}]) 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 + + +def test_json_envelope_split_before_its_value_keeps_the_schema_object() -> None: + envelope: Final = json.dumps({"parameters": _STREAMED_RISK_OBJECT}) + split: Final = envelope.index(":") + 1 + contents: Final = _streamed_json_mode_contents((envelope[:split], envelope[split:])) + assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT 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 0e2f0e58a26..da5a5539bba 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -35,9 +35,10 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) +from litellm.llms.base_llm.base_utils import ResponseFormatSchema, response_format_schema 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(): @@ -8088,3 +8089,462 @@ def _transform_json_mode_body(body: JsonModeMessageBody) -> litellm.ModelRespons json_mode=True, prefix_prompt=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( + {"type": "object", "properties": {"values": {"type": "object", "additionalProperties": {"type": "integer"}}}, "required": ["values"]}, + {"values": {}}, + {"values": {}}, + id="empty_open_values_object_is_preserved", + ), + 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(_RISK_SCHEMA, [1, 2], [1, 2], id="non_object_answer_is_preserved"), + 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_json_tool_without_arguments_has_no_message() -> None: + calls: Final = [ChatCompletionToolCallChunk( + id="toolu_empty", type="function", index=0, + function=ChatCompletionToolCallFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME), + )] + assert AnthropicConfig.convert_tool_response_to_message(calls, _schema(_RISK_SCHEMA)) is None diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index d5fe5d8cf65..5b48ff14822 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -19,9 +19,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, + _stream_decoder, AWSEventStreamDecoder, 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 @@ -1442,3 +1444,90 @@ def test_invoke_decoder_reports_only_missing_aws_dependency(missing): assert error.value.__cause__ is failure else: 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}