mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(anthropic): preserve schema-shaped fallback responses from main
This commit is contained in:
Joshua Valluru
parent
91870f27cb
commit
5c102bcc8a
8 changed files with 1186 additions and 27 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue