mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(anthropic): stream json_schema answers as content on Bedrock Mantle and Claude Platform, and join split answers (#46331)
* fix(vertex_ai): unwrap tool-emulated json_schema envelopes and send output_config.format for Claude
Vertex AI Claude deployments that emulate response_format through a forced
tool sometimes nest the schema object under an extra key such as
`parameters` or `input`, and that envelope reached message.content with a
200. The converter now unwraps exactly one nested object that matches the
schema, leaving a payload that is itself schema-shaped alone.
Native structured-output rows sent the deprecated top-level `output_format`,
which Vertex rejects with a 400. The Vertex request now carries it as
`output_config.format`, with an explicit caller value winning.
* fix(anthropic): unwrap tool-emulated json_schema envelopes on streams and tool-name wrappers
* test(anthropic): expect the json tool description through a shared constant
* fix(anthropic): drop the json tool description that makes sonnet-5 omit required keys
* fix(anthropic): keep schema-accepted answers whole and decide held streams at the second key
An answer the json_schema accepts is returned as is before any envelope rule runs, the tool-name
wrapper is unwrapped only when the schema rejects it, a malformed schema turns the unwrap off, a
streamed answer whose first key is a schema property named name or tool_name waits for its second
key before streaming, the held fragments go through a StringIO, and a prefix that stays undecided
past 4 KB is held to the end
* test(anthropic): keep the stream family cells to providers whose stream iterator gets the schema
* fix(vertex_ai): unwrap claude json_schema answers nested under up to three wrapper keys
A tool-emulated json_schema answer was unwrapped only when the real object sat one key down. In live QA, 1 of 58 answers (/v1/responses, claude-sonnet-5) came back as {"parameter": {"value": {...}}} and reached the client still wrapped
The unwrap now follows single-key wrappers up to three levels down and stops at the first object that has a key the schema declares. That object is returned only when it has every required key. A payload wrapped deeper than three levels, a level with more than one key, and a level whose key is one of the schema's own property names all come back exactly as they arrived
The streamed path needs no change because the buffered answer goes through the same unwrap when the tool block closes, and a test now pins that both paths agree
* fix(bedrock): unwrap streamed claude json_schema answers on claude platform and mantle
A streamed json_schema answer on bedrock/claude_platform and Bedrock Mantle Claude reached the client as tool_calls deltas holding the json_tool_call wrapper, with finish_reason tool_calls and no content. The shared HTTP handler built the async stream iterator without json_mode, and neither stream path passed the response_format schema
The handler now asks the provider config for the iterator through get_model_response_iterator_for_request, which carries json_mode and the request's optional params on both the sync and async paths. Claude Platform and Mantle override it to hand the json tool schema to the Anthropic stream iterator, so the streamed answer is the bare schema object with finish_reason stop, the same as the non-streamed one
Databricks is the only other provider on this handler whose iterator reads json_mode, so its async Claude json mode stream now returns the answer as content like its sync stream already did
The stream family integration test covers claude_platform_unmapped and bedrock_mantle again
* fix(anthropic): unwrap json tool envelopes by structure instead of schema validation
The unwrap now reads only the schema's top-level property names. A lone
undeclared key holding an object, alone or beside a name equal to the json
tool's own name, is the envelope, and the unwrap repeats while the inner
object shares no key with the schema. jsonschema and referencing are no
longer imported on the request path. The stream holds only an answer that
opens with the tool's name or with an undeclared key holding an object, and a
schema without top-level properties streams as it arrives.
* fix(llms): keep other providers' async streams as they were
The request-aware iterator hook passed json_mode on async streams to every provider on the shared HTTP handler, which changed Databricks and any other provider whose iterator reads it. The default hook now makes the same two calls the handler made before, so only Claude Platform and Bedrock Mantle change, and the Databricks test that pinned the wider change is removed
* fix(anthropic): walk json tool envelopes iteratively
* test(llms): pin which stream iterators get json_mode from the shared handler
A provider config that does not override the request-aware iterator hook still gets json_mode on a sync stream and still does not get it on an async stream, the same as before the hook existed
* fix(anthropic): keep json tool arguments too deep to decode instead of raising
* fix(anthropic): unwrap json tool envelopes beside undeclared string labels
* fix(anthropic): keep json answers whole on schemas that allow extra keys and stream held text as sent
A schema that allows undeclared keys and requires none now keeps an undeclared wrapper key on both the stream and non-stream paths. When no envelope is found, the held stream text is now sent exactly as the model wrote it. On Vertex /v1/messages, output_format now wins over output_config.format when both are set, which matches the Anthropic adapter.
* fix(anthropic): join a json_schema answer the model splits into one tool block per field
With tool emulation the model sometimes answers with one json_tool_call block per schema field, each naming the field and holding its value. The non-stream path returned only the first block and the stream sent every block's text back to back, so the client got one field or invalid JSON. Both paths now build one schema object from those blocks, and any other set of blocks keeps today's output
* refactor(anthropic): set the streamed finish reason in one expression
* test(anthropic): cover label envelopes, dynamic-key schemas, as-written held streams and output_format precedence
* fix(anthropic): unwrap a streamed values wrapper like the non-stream answer
* fix(anthropic): join streamed split answers on schemas without required keys, keep client tool calls and length on json_schema streams
* fix(anthropic): join streamed split answers that carry text labels on schemas without required keys
* test(bedrock): script the wrapped json_schema stream with the wrapper Bedrock Mantle sends
* test(bedrock): script the wrapped json_schema stream with the wrapper Bedrock Mantle sent raw
* fix(anthropic): join a split json_schema answer whose fields all hold the same string
* fix(anthropic): join a split json_schema answer sent as one partial object per field
* fix(anthropic): join field blocks after an empty streamed answer without a leading comma
This commit is contained in:
Mateo Wang
GitHub
parent
4f34856a65
commit
457b5d26c7
13 changed files with 1072 additions and 53 deletions
|
|
@ -25,7 +25,10 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc
|
|||
from litellm.llms.base_llm.base_utils import (
|
||||
RESPONSE_FORMAT_ENVELOPE_NAME_KEYS,
|
||||
ResponseFormatSchema,
|
||||
names_a_response_format_property,
|
||||
partial_response_format_answer,
|
||||
response_format_envelope_contents,
|
||||
split_response_format_answer,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -85,6 +88,9 @@ _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*:')
|
||||
_LABEL_THEN_KEY: Final = re.compile(r'"((?:[^"\\]|\\.)*)"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:\s*')
|
||||
_LEADING_LABELS: Final = re.compile(r'(?:"(?:[^"\\]|\\.)*"\s*,\s*"(?:[^"\\]|\\.)*"\s*:\s*)*')
|
||||
_CLOSED_STRING: Final = re.compile(r'"((?:[^"\\]|\\.)*)"')
|
||||
_UNDECIDED_PREFIX_LIMIT: Final = 4096
|
||||
|
||||
|
||||
|
|
@ -95,14 +101,45 @@ def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatS
|
|||
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:
|
||||
def _is_a_label(text: str, schema: ResponseFormatSchema, tool_name: str) -> bool:
|
||||
return text in schema.property_names or text == tool_name
|
||||
|
||||
|
||||
def _label_verdict(label: str, next_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
if _is_a_label(label, schema, tool_name):
|
||||
return "hold"
|
||||
return "pass_through" if next_key in schema.property_names else "undecided"
|
||||
|
||||
|
||||
def _leading_labels_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
last_value_start: Final = _LEADING_LABELS.sub("", string_value_start, count=1)
|
||||
labelled: Final = string_value_start[: len(string_value_start) - len(last_value_start)]
|
||||
verdicts: Final[tuple[_EnvelopeVerdict, ...]] = tuple(
|
||||
_label_verdict(pair.group(1), pair.group(2), schema, tool_name) for pair in _LABEL_THEN_KEY.finditer(labelled)
|
||||
)
|
||||
decided: Final[_EnvelopeVerdict | None] = next((verdict for verdict in verdicts if verdict != "undecided"), None)
|
||||
if decided is not None:
|
||||
return decided
|
||||
last_label: Final = _CLOSED_STRING.match(last_value_start)
|
||||
if last_label is not None and _is_a_label(last_label.group(1), schema, tool_name):
|
||||
return "hold"
|
||||
return "undecided" if not last_value_start or last_value_start.startswith('"') else "pass_through"
|
||||
|
||||
|
||||
def _string_value_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
if schema.required_names:
|
||||
return _key_after_a_string_verdict(string_value_start, schema)
|
||||
return _leading_labels_verdict(string_value_start, schema, tool_name)
|
||||
|
||||
|
||||
def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
|
||||
value_start: Final = after_first_key.lstrip()
|
||||
if not value_start:
|
||||
return "undecided"
|
||||
if value_start.startswith("{"):
|
||||
if value_start.startswith('"'):
|
||||
return _string_value_verdict(value_start, schema, tool_name)
|
||||
if value_start.startswith("{") and not schema.accepts_a_lone_undeclared_key:
|
||||
return "hold"
|
||||
if value_start.startswith('"') and schema.required_names:
|
||||
return _key_after_a_string_verdict(value_start, schema)
|
||||
return "pass_through"
|
||||
|
||||
|
||||
|
|
@ -116,9 +153,11 @@ def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, 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:
|
||||
if first_key.group(1) == "values" and "values" not in schema.property_names:
|
||||
return "hold"
|
||||
if first_key.group(1) in schema.property_names or not schema.property_names:
|
||||
return "pass_through"
|
||||
return _first_value_verdict(buffered_json[first_key.end() :], schema)
|
||||
return _first_value_verdict(buffered_json[first_key.end() :], schema, tool_name)
|
||||
|
||||
|
||||
def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
|
||||
|
|
@ -130,6 +169,13 @@ def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
|
|||
return held if contents is None else json.dumps(contents)
|
||||
|
||||
|
||||
def _split_answer_text(blocks: tuple[str, ...], schema: ResponseFormatSchema) -> str:
|
||||
answer: Final = split_response_format_answer(blocks, schema)
|
||||
if answer is not None:
|
||||
return json.dumps(answer)
|
||||
return "".join(_held_answer(block, schema) for block in blocks)
|
||||
|
||||
|
||||
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:
|
||||
|
|
@ -655,6 +701,9 @@ class ModelResponseIterator:
|
|||
self.response_format_schema = response_format_schema
|
||||
self._response_format_buffer: io.StringIO = io.StringIO()
|
||||
self._response_format_verdict: _EnvelopeVerdict = "undecided"
|
||||
self._split_answer_blocks: tuple[str, ...] = ()
|
||||
self._withheld_answer_close: str = ""
|
||||
self._streamed_answer_keys: frozenset[str] | None = None
|
||||
|
||||
# For handling partial JSON chunks from fragmentation
|
||||
# See: https://github.com/BerriAI/litellm/issues/17473
|
||||
|
|
@ -917,7 +966,7 @@ class ModelResponseIterator:
|
|||
# Track current content block type for filtering deltas
|
||||
self.current_content_block_type = content_block_start["content_block"]["type"]
|
||||
if content_block_start["content_block"]["type"] == "text":
|
||||
text = content_block_start["content_block"]["text"]
|
||||
text = self._release_answer_close() + content_block_start["content_block"]["text"]
|
||||
elif (
|
||||
content_block_start["content_block"]["type"] == "tool_use"
|
||||
or content_block_start["content_block"]["type"] == "server_tool_use"
|
||||
|
|
@ -1049,6 +1098,7 @@ class ModelResponseIterator:
|
|||
pass
|
||||
elif type_chunk == "message_delta":
|
||||
finish_reason, usage, container = self._handle_message_delta(chunk)
|
||||
text = self._release_answer_close() + self._drain_split_answer_blocks()
|
||||
if container:
|
||||
provider_specific_fields["container"] = container
|
||||
# Store container_id and re-emit code_interpreter_results
|
||||
|
|
@ -1155,6 +1205,8 @@ class ModelResponseIterator:
|
|||
# (per Anthropic's fine-grained streaming pattern)
|
||||
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:
|
||||
self.tool_index -= 1
|
||||
|
||||
if self.is_response_format_tool and self.response_format_schema is not None:
|
||||
self.converted_response_format_tool = True
|
||||
|
|
@ -1176,25 +1228,69 @@ class ModelResponseIterator:
|
|||
) -> str:
|
||||
if tool_use.get("id") is not None:
|
||||
self._response_format_buffer = io.StringIO()
|
||||
self._response_format_verdict = "undecided"
|
||||
held_back: Final = self._split_answer_blocks or self._streamed_answer_keys is not None
|
||||
self._response_format_verdict = "hold" if held_back else "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 == "pass_through":
|
||||
return self._withhold_answer_close(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 ""
|
||||
return self._withhold_answer_close(buffered) if self._response_format_verdict == "pass_through" else ""
|
||||
|
||||
def _withhold_answer_close(self, text: str) -> str:
|
||||
pending: Final = self._withheld_answer_close + text
|
||||
body: Final = pending.rstrip()
|
||||
cut: Final = len(body) - 1 if body.endswith("}") else len(pending)
|
||||
self._withheld_answer_close = pending[cut:]
|
||||
return pending[:cut]
|
||||
|
||||
def _release_answer_close(self) -> str:
|
||||
close: Final = self._withheld_answer_close
|
||||
self._withheld_answer_close = ""
|
||||
self._streamed_answer_keys = None
|
||||
return close
|
||||
|
||||
def _join_streamed_answer(self, held: str, schema: ResponseFormatSchema) -> str | None:
|
||||
streamed: Final = self._streamed_answer_keys
|
||||
fields: Final = partial_response_format_answer(held, schema)
|
||||
if streamed is None or fields is None or not streamed.isdisjoint(fields):
|
||||
return None
|
||||
self._streamed_answer_keys = streamed | frozenset(fields)
|
||||
separator: Final = ", " if streamed else ""
|
||||
return f"{separator}{json.dumps(fields)[1:-1]}" if fields else ""
|
||||
|
||||
def _flush_response_format_buffer(self) -> str:
|
||||
if self.response_format_schema is None or self._response_format_verdict == "pass_through":
|
||||
schema: Final = self.response_format_schema
|
||||
if schema is None:
|
||||
return ""
|
||||
held: Final = self._response_format_buffer.getvalue()
|
||||
verdict: Final = self._response_format_verdict
|
||||
self._response_format_buffer = io.StringIO()
|
||||
self._response_format_verdict = "undecided"
|
||||
return _held_answer(held, self.response_format_schema) if held else ""
|
||||
if verdict == "pass_through":
|
||||
answer: Final = partial_response_format_answer(held, schema)
|
||||
close: Final = "" if answer is not None else self._release_answer_close()
|
||||
self._streamed_answer_keys = None if answer is None else frozenset(answer)
|
||||
return close
|
||||
joined: Final = self._join_streamed_answer(held, schema)
|
||||
if joined is not None:
|
||||
return joined
|
||||
released: Final = self._release_answer_close()
|
||||
if self._split_answer_blocks or names_a_response_format_property(held, schema):
|
||||
self._split_answer_blocks = (*self._split_answer_blocks, held)
|
||||
return released
|
||||
return released + (_held_answer(held, schema) if held else "")
|
||||
|
||||
def _drain_split_answer_blocks(self) -> str:
|
||||
blocks: Final = self._split_answer_blocks
|
||||
self._split_answer_blocks = ()
|
||||
if self.response_format_schema is None or not blocks:
|
||||
return ""
|
||||
return _split_answer_text(blocks, self.response_format_schema)
|
||||
|
||||
def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
|
||||
"""
|
||||
|
|
@ -1207,11 +1303,13 @@ class ModelResponseIterator:
|
|||
Tuple of (finish_reason, usage, container)
|
||||
"""
|
||||
message_delta: Final = MessageBlockDelta(**chunk)
|
||||
finish_reason = map_finish_reason(finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop")
|
||||
# Override finish_reason to "stop" if we converted response_format tools
|
||||
# (matches OpenAI behavior and non-streaming Anthropic implementation)
|
||||
if self.converted_response_format_tool:
|
||||
finish_reason = "stop"
|
||||
provider_finish_reason: Final = map_finish_reason(
|
||||
finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop"
|
||||
)
|
||||
called_only_the_response_format_tool: Final = (
|
||||
provider_finish_reason == "tool_calls" and self.converted_response_format_tool and self.tool_index < 0
|
||||
)
|
||||
finish_reason: Final = "stop" if called_only_the_response_format_tool else provider_finish_reason
|
||||
usage: Final = (
|
||||
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.llms.base_llm.base_utils import (
|
|||
ResponseFormatSchema,
|
||||
as_json_object,
|
||||
response_format_schema,
|
||||
split_response_format_answer,
|
||||
type_to_response_format_param,
|
||||
unwrap_response_format_envelope,
|
||||
)
|
||||
|
|
@ -2268,13 +2269,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
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], response_format_schema=response_format_schema
|
||||
tool_calls=tool_calls, 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], response_format_schema=response_format_schema
|
||||
[tool_calls[i] for i in json_indices], 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]
|
||||
|
|
@ -2895,6 +2895,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
json_mode_content_str: Final[str | None] = tool_calls[0]["function"].get("arguments")
|
||||
if json_mode_content_str is None:
|
||||
return None
|
||||
split_answer: Final = (
|
||||
None
|
||||
if response_format_schema is None
|
||||
else split_response_format_answer(
|
||||
tuple(tool_call["function"].get("arguments") or "" for tool_call in tool_calls),
|
||||
response_format_schema,
|
||||
)
|
||||
)
|
||||
if split_answer is not None:
|
||||
return litellm.Message(content=json.dumps(split_answer))
|
||||
try:
|
||||
payload: Final[object] = json.loads(json_mode_content_str)
|
||||
except (json.JSONDecodeError, RecursionError):
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import copy
|
|||
import itertools
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -289,6 +289,73 @@ def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchem
|
|||
return payload if contents is None else contents
|
||||
|
||||
|
||||
def _field_names(
|
||||
blocks: tuple[Mapping[str, JsonValue], ...], key: str, schema: ResponseFormatSchema
|
||||
) -> tuple[str, ...] | None:
|
||||
values: Final = (block[key] for block in blocks)
|
||||
names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.property_names)
|
||||
return names if len(frozenset(names)) == len(blocks) else None
|
||||
|
||||
|
||||
def _is_a_label_every_block_repeats(blocks: tuple[Mapping[str, JsonValue], ...], key: str) -> bool:
|
||||
return isinstance(blocks[0][key], str) and all(block[key] == blocks[0][key] for block in blocks)
|
||||
|
||||
|
||||
def _value_key(blocks: tuple[Mapping[str, JsonValue], ...], keys: frozenset[str]) -> str | None:
|
||||
if len(keys) == 1:
|
||||
return next(iter(keys))
|
||||
varying: Final = tuple(key for key in keys if not _is_a_label_every_block_repeats(blocks, key))
|
||||
return varying[0] if len(varying) == 1 else None
|
||||
|
||||
|
||||
def _split_answer_fields(
|
||||
blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
keys: Final = frozenset(blocks[0])
|
||||
if not keys.isdisjoint(schema.property_names) or any(frozenset(block) != keys for block in blocks):
|
||||
return None
|
||||
named: Final = tuple((key, names) for key in keys if (names := _field_names(blocks, key, schema)) is not None)
|
||||
if len(named) != 1:
|
||||
return None
|
||||
((name_key, field_names),) = named
|
||||
value_key: Final = _value_key(blocks, keys - {name_key})
|
||||
if value_key is None:
|
||||
return None
|
||||
return dict(zip(field_names, (block[value_key] for block in blocks), strict=True))
|
||||
|
||||
|
||||
def _merged_partial_answers(
|
||||
blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
keys: Final = frozenset(itertools.chain.from_iterable(blocks))
|
||||
if not keys <= schema.property_names or len(keys) != sum(len(block) for block in blocks):
|
||||
return None
|
||||
return dict(itertools.chain.from_iterable(block.items() for block in blocks))
|
||||
|
||||
|
||||
def split_response_format_answer(
|
||||
encoded_blocks: Sequence[str], schema: ResponseFormatSchema
|
||||
) -> Mapping[str, JsonValue] | None:
|
||||
decoded: Final = tuple(_json_object_or_encoded_object(encoded) for encoded in encoded_blocks)
|
||||
blocks: Final = tuple(block for block in decoded if block is not None)
|
||||
if len(blocks) < 2 or len(blocks) != len(decoded):
|
||||
return None
|
||||
named_fields: Final = _split_answer_fields(blocks, schema)
|
||||
return named_fields if named_fields is not None else _merged_partial_answers(blocks, schema)
|
||||
|
||||
|
||||
def partial_response_format_answer(encoded_block: str, schema: ResponseFormatSchema) -> Mapping[str, JsonValue] | None:
|
||||
block: Final = _json_object_or_encoded_object(encoded_block)
|
||||
return block if block is not None and frozenset(block) <= schema.property_names else None
|
||||
|
||||
|
||||
def names_a_response_format_property(encoded_block: str, schema: ResponseFormatSchema) -> bool:
|
||||
block: Final = _json_object_or_encoded_object(encoded_block)
|
||||
if block is None:
|
||||
return False
|
||||
return any(isinstance(value, str) and value in schema.property_names for value in block.values())
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -417,6 +417,20 @@ class BaseConfig(ABC):
|
|||
) -> Any:
|
||||
pass
|
||||
|
||||
def get_model_response_iterator_for_request(
|
||||
self,
|
||||
streaming_response: "Iterator[str] | AsyncIterator[str] | ModelResponse",
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None,
|
||||
optional_params: Mapping[str, object],
|
||||
) -> object:
|
||||
if sync_stream:
|
||||
return self.get_model_response_iterator(
|
||||
streaming_response=streaming_response, sync_stream=True, json_mode=json_mode
|
||||
)
|
||||
# FIXME: async streams never got json_mode, passing it changes every provider whose iterator reads it
|
||||
return self.get_model_response_iterator(streaming_response=streaming_response, sync_stream=False)
|
||||
|
||||
async def get_async_custom_stream_wrapper(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -7,10 +7,11 @@ The bedrock-mantle endpoint uses the Anthropic Messages API format but is served
|
|||
at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
|
||||
from litellm.llms.anthropic.chat.transformation import JsonModeRequest
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeConfig,
|
||||
)
|
||||
|
|
@ -134,10 +135,22 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
|
|||
sync_stream: bool,
|
||||
json_mode: bool | None = False,
|
||||
) -> Any:
|
||||
return self.get_model_response_iterator_for_request(
|
||||
streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={}
|
||||
)
|
||||
|
||||
def get_model_response_iterator_for_request(
|
||||
self,
|
||||
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None,
|
||||
optional_params: Mapping[str, object],
|
||||
) -> object:
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
|
||||
|
||||
return ModelResponseIterator(
|
||||
streaming_response=streaming_response,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig, JsonModeRequest
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from .common_utils import (
|
||||
BedrockClaudePlatformMixin,
|
||||
|
|
@ -91,14 +93,26 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig):
|
|||
|
||||
def get_model_response_iterator(
|
||||
self,
|
||||
streaming_response: Any,
|
||||
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None = False,
|
||||
) -> Any:
|
||||
return self.get_model_response_iterator_for_request(
|
||||
streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={}
|
||||
)
|
||||
|
||||
def get_model_response_iterator_for_request(
|
||||
self,
|
||||
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None,
|
||||
optional_params: Mapping[str, object],
|
||||
) -> object:
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
|
||||
|
||||
return ModelResponseIterator(
|
||||
streaming_response=streaming_response,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=bool(json_mode),
|
||||
response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import httpx
|
|||
from httpx import USE_CLIENT_DEFAULT
|
||||
from httpx._types import FileContent
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
|
|
@ -464,6 +465,7 @@ async def _aiter_bytes_then_close(response: httpx.Response, *, chunk_size: int)
|
|||
|
||||
|
||||
_DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encoding", "content-length"})
|
||||
_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _decoded_body_headers(response: httpx.Response) -> httpx.Headers:
|
||||
|
|
@ -1000,10 +1002,11 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
|
||||
else:
|
||||
completion_stream = provider_config.get_model_response_iterator(
|
||||
completion_stream = provider_config.get_model_response_iterator_for_request(
|
||||
streaming_response=response.iter_lines(),
|
||||
sync_stream=True,
|
||||
json_mode=json_mode,
|
||||
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -1139,8 +1142,11 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
|
||||
else:
|
||||
completion_stream = provider_config.get_model_response_iterator(
|
||||
streaming_response=response.aiter_lines(), sync_stream=False
|
||||
completion_stream = provider_config.get_model_response_iterator_for_request(
|
||||
streaming_response=response.aiter_lines(),
|
||||
sync_stream=False,
|
||||
json_mode=json_mode,
|
||||
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
|
||||
)
|
||||
if isinstance(completion_stream, BaseModelResponseIterator):
|
||||
completion_stream.http_response = response
|
||||
|
|
|
|||
|
|
@ -171,7 +171,6 @@ _FAMILY: Final = (
|
|||
"claude_platform_unmapped",
|
||||
"vertex_haiku",
|
||||
)
|
||||
_STREAM_FAMILY: Final = ("anthropic_unmapped", "azure_ai_unmapped", "bedrock_invoke", "vertex_haiku")
|
||||
_FRAGMENT_SIZE: Final = 20
|
||||
_TIMING_BYTES: Final = 512 * 1024
|
||||
_TIMING_FRAGMENT: Final = 4096
|
||||
|
|
@ -1169,8 +1168,7 @@ def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_nam
|
|||
content: Final = _stream_content(response) if stream else _chat_content(response)
|
||||
assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text
|
||||
if stream:
|
||||
expected_frames: Final = 1 if unwrapped else _fragment_count(shape, marker_hex)
|
||||
assert len(_content_frames(response)) == expected_frames, response.text
|
||||
assert (len(_content_frames(response)) == 1) is unwrapped, response.text
|
||||
_one_delivery(wire, marker_hex)
|
||||
|
||||
|
||||
|
|
@ -1322,7 +1320,7 @@ def test_tool_path_family_unwraps_the_envelope(gateway: Gateway, member: str) ->
|
|||
assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("member", _STREAM_FAMILY)
|
||||
@pytest.mark.parametrize("member", _FAMILY)
|
||||
def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: str) -> None:
|
||||
marker, marker_hex = _marker("wrapped")
|
||||
with wire_server(_peer) as wire, gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -2810,9 +2810,7 @@ def test_streamed_json_mode_schema_object_streams_as_it_arrives():
|
|||
contents = _streamed_json_mode_contents(fragments)
|
||||
|
||||
assert "".join(contents) == answer
|
||||
assert contents[3] == answer[:12]
|
||||
assert contents[4] == answer[12:]
|
||||
assert contents[5] == ""
|
||||
assert contents[3:] == (answer[:12], answer[12:-1], "", "}")
|
||||
|
||||
|
||||
def test_streamed_json_mode_truncated_output_is_passed_through_untouched():
|
||||
|
|
@ -2836,9 +2834,7 @@ def test_streamed_json_mode_name_property_answer_streams_once_its_name_value_is_
|
|||
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:]
|
||||
assert contents[2:] == ("", answer[:24], answer[24:-1], "", "}")
|
||||
|
||||
|
||||
def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
|
||||
|
|
@ -2847,8 +2843,7 @@ def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
|
|||
contents = _streamed_json_mode_contents((answer[:11], answer[11:]), _STREAMED_CONTACT_SCHEMA)
|
||||
|
||||
assert "".join(contents) == answer
|
||||
assert contents[2] == ""
|
||||
assert contents[3] == answer
|
||||
assert contents[2:] == ("", answer[:-1], "", "}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -2865,12 +2860,11 @@ def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
|
|||
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"}},
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"summary": {"type": "string"}},
|
||||
"additionalProperties": {"type": "object"},
|
||||
},
|
||||
{"user_1": {"summary": "likes tea", "visits": 3}},
|
||||
id="undeclared_key_the_schema_allows",
|
||||
),
|
||||
|
|
@ -2884,7 +2878,61 @@ def test_streamed_json_mode_answer_that_cannot_be_an_envelope_streams_as_it_arri
|
|||
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:])
|
||||
assert contents[2:] == (encoded[:12], encoded[12:24], encoded[24:-1], "}", "")
|
||||
|
||||
|
||||
def test_streamed_json_mode_text_label_then_an_object_streams_once_the_object_opens():
|
||||
encoded = json.dumps({"note": "read-only window", "meta": {"score": 7}})
|
||||
object_opens = encoded.index('{"score"') + 1
|
||||
|
||||
contents = _streamed_json_mode_contents(
|
||||
(encoded[:12], encoded[12:object_opens], encoded[object_opens:]),
|
||||
{"type": "object", "properties": {"score": {"type": "integer"}}},
|
||||
)
|
||||
|
||||
assert contents[2:] == ("", encoded[:object_opens], encoded[object_opens:-1], "}", "")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema, values",
|
||||
[
|
||||
pytest.param(
|
||||
{"type": "object", "additionalProperties": True}, {"alpha": 1, "beta": 2}, id="schema_without_properties"
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"summary": {"type": "string"}},
|
||||
"additionalProperties": {"type": "object"},
|
||||
},
|
||||
{"user_1": {"summary": "likes tea"}},
|
||||
id="undeclared_key_the_schema_allows",
|
||||
),
|
||||
pytest.param(_STREAMED_RISK_SCHEMA, [3, 5, 8], id="values_holding_a_list"),
|
||||
],
|
||||
)
|
||||
def test_streamed_json_mode_values_wrapper_is_unwrapped_like_the_non_stream_answer(
|
||||
schema: dict[str, object], values: object
|
||||
):
|
||||
wrapped = {"values": values}
|
||||
encoded = json.dumps(wrapped)
|
||||
|
||||
contents = _streamed_json_mode_contents((encoded[:6], encoded[6:14], encoded[14:]), schema)
|
||||
|
||||
assert json.loads("".join(contents)) == values
|
||||
assert json.loads("".join(contents)) == unwrap_response_format_envelope(
|
||||
wrapped, response_format_schema(schema), RESPONSE_FORMAT_TOOL_NAME
|
||||
)
|
||||
|
||||
|
||||
def test_streamed_json_mode_declared_values_key_streams_as_it_arrives():
|
||||
schema = {"type": "object", "properties": {"values": {"type": "array"}}, "required": ["values"]}
|
||||
answer = json.dumps({"values": [3, 5, 8]})
|
||||
|
||||
contents = _streamed_json_mode_contents((answer[:6], answer[6:14], answer[14:]), schema)
|
||||
|
||||
assert "".join(contents) == answer
|
||||
assert contents[3:] == (answer[:14], answer[14:-1], "", "}")
|
||||
|
||||
|
||||
def test_streamed_json_mode_string_label_answer_streams_once_the_next_key_is_declared():
|
||||
|
|
@ -2893,9 +2941,7 @@ def test_streamed_json_mode_string_label_answer_streams_once_the_next_key_is_dec
|
|||
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:]
|
||||
assert contents[2:] == ("", answer[:40], answer[40:-1], "}", "")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -2978,6 +3024,226 @@ def test_streamed_json_mode_answer_with_a_foreign_name_matches_the_non_streamed_
|
|||
assert unwrap_response_format_envelope(answer, schema, RESPONSE_FORMAT_TOOL_NAME) == answer
|
||||
|
||||
|
||||
def _json_tool_block_events(index: int, encoded: str) -> tuple[dict[str, object], ...]:
|
||||
return (
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": f"toolu_{index}",
|
||||
"name": RESPONSE_FORMAT_TOOL_NAME,
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
*(
|
||||
{"type": "content_block_delta", "index": index, "delta": {"type": "input_json_delta", "partial_json": part}}
|
||||
for part in (encoded[:7], encoded[7:])
|
||||
),
|
||||
{"type": "content_block_stop", "index": index},
|
||||
)
|
||||
|
||||
|
||||
def _streamed_json_tool_blocks(
|
||||
blocks: tuple[str, ...], schema: dict[str, object] = _STREAMED_RISK_SCHEMA
|
||||
) -> tuple[str, ...]:
|
||||
message_start, *_, message_delta = _json_tool_events(())
|
||||
events = (
|
||||
message_start,
|
||||
*chain.from_iterable(_json_tool_block_events(index, encoded) for index, encoded in enumerate(blocks)),
|
||||
message_delta,
|
||||
)
|
||||
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 events)
|
||||
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)
|
||||
|
||||
|
||||
_STREAMED_TIER_FIELD: Final = json.dumps(
|
||||
{"$FUNCTION_NAME": "RiskAssessment", "$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _STREAMED_RISK_OBJECT["tier"]}
|
||||
)
|
||||
_STREAMED_SUMMARY_FIELD: Final = json.dumps(
|
||||
{
|
||||
"$FUNCTION_NAME": "RiskAssessment",
|
||||
"$PARAMETER_NAME": "summary",
|
||||
"$PARAMETER_VALUE": _STREAMED_RISK_OBJECT["summary"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_streamed_json_mode_answer_split_into_one_block_per_field_arrives_as_one_object():
|
||||
contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD, _STREAMED_SUMMARY_FIELD))
|
||||
|
||||
assert json.loads(contents[-1]) == _STREAMED_RISK_OBJECT
|
||||
assert contents[:-1] == ("",) * (len(contents) - 1)
|
||||
|
||||
|
||||
def test_streamed_json_mode_field_block_alone_keeps_the_model_text():
|
||||
contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD,))
|
||||
|
||||
assert "".join(contents) == _STREAMED_TIER_FIELD
|
||||
|
||||
|
||||
def test_streamed_json_mode_whole_answer_after_a_field_block_keeps_the_order_the_model_sent():
|
||||
whole_answer = json.dumps(_STREAMED_RISK_OBJECT)
|
||||
|
||||
contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD, whole_answer))
|
||||
|
||||
assert "".join(contents) == _STREAMED_TIER_FIELD + whole_answer
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema, answer",
|
||||
[
|
||||
pytest.param(_STREAMED_RISK_SCHEMA, _STREAMED_RISK_OBJECT, id="schema_with_a_required_list"),
|
||||
pytest.param(
|
||||
_STREAMED_CONTACT_SCHEMA, {"name": "unknown", "email": "unknown"}, id="every_field_holds_the_same_string"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_streamed_json_mode_answer_split_into_one_partial_object_per_field_arrives_as_one_object(
|
||||
schema: dict[str, object], answer: dict[str, object]
|
||||
):
|
||||
blocks = tuple(json.dumps({name: value}) for name, value in answer.items())
|
||||
|
||||
contents = _streamed_json_tool_blocks(blocks, schema)
|
||||
|
||||
assert json.loads("".join(contents)) == answer
|
||||
assert contents[3] == blocks[0][:-1]
|
||||
|
||||
|
||||
def test_streamed_json_mode_empty_object_then_field_blocks_arrives_as_one_object():
|
||||
blocks = ("{}", *(json.dumps({name: value}) for name, value in _STREAMED_RISK_OBJECT.items()))
|
||||
|
||||
contents = _streamed_json_tool_blocks(blocks)
|
||||
|
||||
assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT
|
||||
|
||||
|
||||
def test_streamed_json_mode_partial_objects_that_repeat_a_field_keep_the_order_the_model_sent():
|
||||
blocks = (json.dumps({"tier": 3}), json.dumps({"tier": 4}))
|
||||
|
||||
contents = _streamed_json_tool_blocks(blocks)
|
||||
|
||||
assert "".join(contents) == "".join(blocks)
|
||||
|
||||
|
||||
def test_streamed_json_mode_text_after_the_answer_follows_the_whole_answer():
|
||||
answer = json.dumps(_STREAMED_RISK_OBJECT)
|
||||
message_start, *_, message_delta = _json_tool_events(())
|
||||
text_block = (
|
||||
{"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " Done."}},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
)
|
||||
|
||||
chunks = _json_mode_stream_chunks((message_start, *_json_tool_block_events(0, answer), *text_block, message_delta))
|
||||
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer + " Done."
|
||||
|
||||
|
||||
def test_streamed_json_mode_envelope_whose_label_is_a_property_name_is_still_unwrapped():
|
||||
envelope = json.dumps({"$FUNCTION_NAME": "summary", "PARAMETER_VALUE": _STREAMED_RISK_OBJECT})
|
||||
|
||||
contents = _streamed_json_tool_blocks((envelope,))
|
||||
|
||||
assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT
|
||||
|
||||
|
||||
_STREAMED_OPTIONAL_RISK_SCHEMA: Final = {
|
||||
"type": "object",
|
||||
"properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema",
|
||||
[
|
||||
pytest.param(_STREAMED_OPTIONAL_RISK_SCHEMA, id="no_required_keys"),
|
||||
pytest.param(
|
||||
{**_STREAMED_OPTIONAL_RISK_SCHEMA, "additionalProperties": True}, id="no_required_keys_and_open_to_others"
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"leading_label",
|
||||
[
|
||||
pytest.param({}, id="field_name_label_first"),
|
||||
pytest.param({"$FUNCTION_NAME": RESPONSE_FORMAT_TOOL_NAME}, id="tool_name_label_first"),
|
||||
pytest.param({"$FUNCTION_NAME": "RiskAssessment"}, id="schema_name_label_first"),
|
||||
pytest.param({"$FUNCTION_NAME": "RiskAssessment", "$CALL": "first"}, id="two_text_labels_first"),
|
||||
],
|
||||
)
|
||||
def test_streamed_json_mode_split_answer_arrives_as_one_object_when_the_schema_requires_no_key(
|
||||
schema: dict[str, object], leading_label: dict[str, str]
|
||||
):
|
||||
blocks = tuple(
|
||||
json.dumps({**leading_label, "$PARAMETER_NAME": name, "$PARAMETER_VALUE": value})
|
||||
for name, value in _STREAMED_RISK_OBJECT.items()
|
||||
)
|
||||
|
||||
contents = _streamed_json_tool_blocks(blocks, schema)
|
||||
|
||||
assert json.loads(contents[-1]) == _STREAMED_RISK_OBJECT
|
||||
assert contents[:-1] == ("",) * (len(contents) - 1)
|
||||
|
||||
|
||||
def _json_mode_stream_chunks(events: tuple[dict[str, object], ...]) -> tuple[ModelResponseStream, ...]:
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=True,
|
||||
json_mode=True,
|
||||
response_format_schema=response_format_schema(_STREAMED_RISK_SCHEMA),
|
||||
)
|
||||
return tuple(iterator.chunk_parser(event) for event in events)
|
||||
|
||||
|
||||
def test_streamed_json_mode_answer_beside_a_client_tool_call_keeps_that_call_usable():
|
||||
answer = json.dumps(_STREAMED_RISK_OBJECT)
|
||||
message_start, *_, message_delta = _json_tool_events(())
|
||||
weather_call = (
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_weather", "name": "get_weather", "input": {}},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 1,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
)
|
||||
|
||||
chunks = _json_mode_stream_chunks(
|
||||
(message_start, *_json_tool_block_events(0, answer), *weather_call, message_delta)
|
||||
)
|
||||
|
||||
tool_calls = tuple(chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks))
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer
|
||||
assert {call.index for call in tool_calls} == {0}
|
||||
assert tool_calls[0].function.name == "get_weather"
|
||||
assert "".join(call.function.arguments or "" for call in tool_calls) == '{"city": "Paris"}'
|
||||
assert chunks[-1].choices[0].finish_reason == "tool_calls"
|
||||
|
||||
|
||||
def test_streamed_json_mode_answer_cut_at_max_tokens_ends_with_length():
|
||||
cut_answer = '{"tier": 3, "summ'
|
||||
*events, message_delta = _json_tool_events((cut_answer,))
|
||||
cut_message_delta = {**message_delta, "delta": {"stop_reason": "max_tokens", "stop_sequence": None}}
|
||||
|
||||
chunks = _json_mode_stream_chunks((*events, cut_message_delta))
|
||||
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == cut_answer
|
||||
assert chunks[-1].choices[0].finish_reason == "length"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_anthropic_streaming_function_call_arguments_are_assembled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import copy
|
||||
import json
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -7767,6 +7768,278 @@ def test_vertex_claude_tool_emulated_response_format_returns_the_object_the_mode
|
|||
assert result.choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
_SPLIT_RISK_SCHEMA: Final = {
|
||||
"title": "RiskAssessment",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tier": {"type": "integer"},
|
||||
"risk_level": {"type": "string", "enum": ["low", "medium", "high"]},
|
||||
"score": {"type": "integer"},
|
||||
"summary": {"type": "string"},
|
||||
"affected_systems": {"type": "array", "items": {"type": "string"}},
|
||||
"findings": {"type": "array", "items": {"type": "object"}},
|
||||
"rollback_plan": {"type": "string"},
|
||||
"requires_downtime": {"type": "boolean"},
|
||||
},
|
||||
"required": [
|
||||
"tier",
|
||||
"risk_level",
|
||||
"score",
|
||||
"summary",
|
||||
"affected_systems",
|
||||
"findings",
|
||||
"rollback_plan",
|
||||
"requires_downtime",
|
||||
],
|
||||
}
|
||||
_SPLIT_RISK_OBJECT: Final = {
|
||||
"tier": 1,
|
||||
"risk_level": "high",
|
||||
"score": 78,
|
||||
"summary": "Migrating loyalty-points ledger DB to managed cloud with weekend read-only cutover; financial data, high risk.",
|
||||
"affected_systems": ["Loyalty-points ledger DB", "Loyalty rewards API", "Reporting/analytics pipeline"],
|
||||
"findings": [
|
||||
{
|
||||
"title": "Financial data integrity risk",
|
||||
"severity": "high",
|
||||
"detail": "Loyalty points are monetary-equivalent; any sync error causes customer-facing balance discrepancies.",
|
||||
},
|
||||
{
|
||||
"title": "Weekend support coverage",
|
||||
"severity": "low",
|
||||
"detail": "Reduced staff availability over weekend may delay incident response if issues arise.",
|
||||
},
|
||||
],
|
||||
"rollback_plan": "Keep old cluster live/read-only; revert DNS/connection strings to original cluster if migration fails.",
|
||||
"requires_downtime": True,
|
||||
}
|
||||
# Block shape captured from vertex_ai/claude-sonnet-5 tool emulation on 2026-10-13, one tool_use block per field
|
||||
_SPLIT_RISK_BLOCKS: Final = tuple(
|
||||
{"$FUNCTION_NAME": "RiskAssessment", "$PARAMETER_NAME": name, "$PARAMETER_VALUE": value}
|
||||
for name, value in _SPLIT_RISK_OBJECT.items()
|
||||
)
|
||||
_CONTACT_SCHEMA: Final = {
|
||||
"type": "object",
|
||||
"properties": {"name": {"type": "string"}, "email": {"type": "string"}},
|
||||
"required": ["name", "email"],
|
||||
}
|
||||
_OPTIONAL_CONTACT_SCHEMA: Final = {
|
||||
"type": "object",
|
||||
"properties": {"name": {"type": "string"}, "email": {"type": "string"}, "phone": {"type": "string"}},
|
||||
}
|
||||
_FLAGS_SCHEMA: Final = {
|
||||
"type": "object",
|
||||
"properties": {"urgent": {"type": "boolean"}, "reviewed": {"type": "boolean"}},
|
||||
"required": ["urgent", "reviewed"],
|
||||
}
|
||||
|
||||
|
||||
def _json_tool_calls(*blocks: object) -> list[ChatCompletionToolCallChunk]:
|
||||
return list(chain.from_iterable(_json_tool_call(block) for block in blocks))
|
||||
|
||||
|
||||
def test_vertex_claude_json_answer_split_into_one_tool_block_per_field_returns_one_schema_object():
|
||||
config = VertexAIAnthropicConfig()
|
||||
optional_params = config.map_openai_params(
|
||||
non_default_params={
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "RiskAssessment", "schema": _SPLIT_RISK_SCHEMA},
|
||||
},
|
||||
},
|
||||
optional_params={},
|
||||
model="claude-emulated-structured-output",
|
||||
drop_params=False,
|
||||
)
|
||||
json_mode = optional_params.pop("json_mode")
|
||||
completion_response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-emulated-structured-output",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": f"toolu_{index}", "name": RESPONSE_FORMAT_TOOL_NAME, "input": block}
|
||||
for index, block in enumerate(_SPLIT_RISK_BLOCKS)
|
||||
],
|
||||
"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={},
|
||||
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) == _SPLIT_RISK_OBJECT
|
||||
assert result.choices[0].message.tool_calls is None
|
||||
assert result.choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema, blocks, expected_content",
|
||||
[
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
(
|
||||
{"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _RISK_OBJECT["tier"]},
|
||||
{"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": _RISK_OBJECT["summary"]},
|
||||
),
|
||||
_RISK_OBJECT,
|
||||
id="name_and_value_keys_only",
|
||||
),
|
||||
pytest.param(
|
||||
_CONTACT_SCHEMA,
|
||||
(
|
||||
{"function": "Contact", "parameter_name": "name", "parameter_value": "Ada"},
|
||||
{"function": "Contact", "parameter_name": "email", "parameter_value": "ada@example.com"},
|
||||
),
|
||||
{"name": "Ada", "email": "ada@example.com"},
|
||||
id="string_values_beside_a_repeated_label",
|
||||
),
|
||||
pytest.param(
|
||||
_FLAGS_SCHEMA,
|
||||
({"field": "urgent", "value": True}, {"field": "reviewed", "value": True}),
|
||||
{"urgent": True, "reviewed": True},
|
||||
id="every_field_holds_the_same_value",
|
||||
),
|
||||
pytest.param(
|
||||
_CONTACT_SCHEMA,
|
||||
({"field": "name", "value": "unknown"}, {"field": "email", "value": "unknown"}),
|
||||
{"name": "unknown", "email": "unknown"},
|
||||
id="every_field_holds_the_same_string",
|
||||
),
|
||||
pytest.param(
|
||||
_CONTACT_SCHEMA,
|
||||
({"name": "unknown"}, {"email": "unknown"}),
|
||||
{"name": "unknown", "email": "unknown"},
|
||||
id="one_partial_object_per_field",
|
||||
),
|
||||
pytest.param(
|
||||
_OPTIONAL_CONTACT_SCHEMA,
|
||||
({"name": "Ada", "email": "ada@example.com"}, {"phone": "555-0100"}),
|
||||
{"name": "Ada", "email": "ada@example.com", "phone": "555-0100"},
|
||||
id="partial_objects_of_uneven_size_on_a_schema_with_no_required_list",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_json_tool_blocks_that_each_name_one_schema_property_become_one_object(
|
||||
schema: dict[str, object], blocks: tuple[object, ...], expected_content: dict[str, object]
|
||||
):
|
||||
replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
|
||||
json_mode=True,
|
||||
tool_calls=_json_tool_calls(*blocks),
|
||||
response_format_schema=_schema(schema),
|
||||
)
|
||||
|
||||
assert replacement is not None
|
||||
assert json.loads(replacement.content) == expected_content
|
||||
assert remaining_tool_calls == []
|
||||
assert extra_content is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema, blocks",
|
||||
[
|
||||
pytest.param(_RISK_SCHEMA, ({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},), id="one_field_alone"),
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3}, {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 4}),
|
||||
id="same_property_named_twice",
|
||||
),
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
(
|
||||
{"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},
|
||||
{"$PARAMETER_NAME": "color", "$PARAMETER_VALUE": "red"},
|
||||
),
|
||||
id="property_the_schema_does_not_declare",
|
||||
),
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
(
|
||||
{"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},
|
||||
{"parameter_name": "summary", "parameter_value": "Read-only window"},
|
||||
),
|
||||
id="blocks_with_different_keys",
|
||||
),
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
(
|
||||
{"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3, "note": "first"},
|
||||
{"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": "Read-only window", "note": "second"},
|
||||
),
|
||||
id="two_keys_that_could_hold_the_value",
|
||||
),
|
||||
pytest.param(
|
||||
_RISK_SCHEMA,
|
||||
({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3}, ["summary", "Read-only window"]),
|
||||
id="block_that_is_not_an_object",
|
||||
),
|
||||
pytest.param(
|
||||
_PERSON_SCHEMA,
|
||||
({"name": "name", "age": 1}, {"name": "age", "age": 2}),
|
||||
id="blocks_keyed_by_schema_properties",
|
||||
),
|
||||
pytest.param(
|
||||
_CONTACT_SCHEMA,
|
||||
({"name": "Ada"}, {"name": "Grace"}),
|
||||
id="partial_objects_that_repeat_a_field",
|
||||
),
|
||||
pytest.param(
|
||||
_CONTACT_SCHEMA,
|
||||
({"name": "Ada"}, {"nickname": "A"}),
|
||||
id="partial_object_with_a_key_the_schema_does_not_declare",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_json_tool_blocks_that_are_not_one_field_each_keep_the_first_block(
|
||||
schema: dict[str, object], blocks: tuple[object, ...]
|
||||
):
|
||||
replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
|
||||
json_mode=True,
|
||||
tool_calls=_json_tool_calls(*blocks),
|
||||
response_format_schema=_schema(schema),
|
||||
)
|
||||
|
||||
assert replacement is not None
|
||||
assert json.loads(replacement.content) == blocks[0]
|
||||
assert remaining_tool_calls == []
|
||||
assert extra_content is None
|
||||
|
||||
|
||||
def test_json_mode_with_user_tool_call_joins_the_split_answer_into_content():
|
||||
user_tool_call: Final = ChatCompletionToolCallChunk(
|
||||
id="toolu_user",
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(name="lookup_change_window", arguments='{"system": "ledger"}'),
|
||||
index=2,
|
||||
)
|
||||
first_field, second_field = _json_tool_calls(
|
||||
{"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _RISK_OBJECT["tier"]},
|
||||
{"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": _RISK_OBJECT["summary"]},
|
||||
)
|
||||
|
||||
replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
|
||||
json_mode=True,
|
||||
tool_calls=[first_field, user_tool_call, second_field],
|
||||
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_anthropic_tool_with_image():
|
||||
import json
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import httpx
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
|
||||
ONE_PIXEL_PNG = base64.b64decode(
|
||||
|
|
@ -140,3 +140,106 @@ async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sourc
|
|||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
|
||||
} in captured["body"]["messages"][0]["content"]
|
||||
|
||||
|
||||
RISK_ANSWER = {"tier": 2, "summary": "ok"}
|
||||
MANTLE_CLAUDE_ROUTES = ("bedrock_mantle/anthropic.claude-scripted-9", "bedrock/mantle/anthropic.claude-scripted-9")
|
||||
|
||||
|
||||
def _json_schema_stream_kwargs(model):
|
||||
return {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "Assess the risk"}],
|
||||
"stream": True,
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "Risk",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
|
||||
"required": ["tier", "summary"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
},
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "fake-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
|
||||
|
||||
def _wrapped_json_tool_stream(request):
|
||||
wrapper = json.dumps({"$PARAMETER_NAME": RISK_ANSWER})
|
||||
events = (
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "anthropic.claude-scripted-9",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "json_tool_call", "input": {}},
|
||||
},
|
||||
*(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": wrapper[start : start + 20]},
|
||||
}
|
||||
for start in range(0, len(wrapper), 20)
|
||||
),
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 9},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode(),
|
||||
)
|
||||
|
||||
|
||||
def _streamed_answer(chunks):
|
||||
deltas = tuple(chunk.choices[0].delta for chunk in chunks)
|
||||
return {
|
||||
"content": json.loads("".join(delta.content or "" for delta in deltas)),
|
||||
"tool_call_deltas": sum(1 for delta in deltas if delta.tool_calls),
|
||||
"finish_reasons": tuple(chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason),
|
||||
}
|
||||
|
||||
|
||||
BARE_STREAMED_ANSWER = {"content": RISK_ANSWER, "tool_call_deltas": 0, "finish_reasons": ("stop",)}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MANTLE_CLAUDE_ROUTES)
|
||||
def test_mantle_claude_completion_streams_a_wrapped_json_schema_answer_as_the_bare_object(model):
|
||||
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_wrapped_json_tool_stream)))
|
||||
|
||||
chunks = tuple(litellm.completion(**_json_schema_stream_kwargs(model), client=client))
|
||||
|
||||
assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MANTLE_CLAUDE_ROUTES)
|
||||
async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answer_as_the_bare_object(model):
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(_wrapped_json_tool_stream))
|
||||
|
||||
stream = await litellm.acompletion(**_json_schema_stream_kwargs(model), client=client)
|
||||
chunks = tuple([chunk async for chunk in stream])
|
||||
|
||||
assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER
|
||||
|
|
|
|||
|
|
@ -103,9 +103,12 @@ def _sync_gateway_client(
|
|||
return HTTPHandler(client=httpx.Client(transport=_gateway_transport(requests, respond)))
|
||||
|
||||
|
||||
def _async_gateway_client(requests: list[dict]) -> AsyncHTTPHandler:
|
||||
def _async_gateway_client(
|
||||
requests: list[dict],
|
||||
respond: GatewayResponder = _fake_claude_platform_gateway,
|
||||
) -> AsyncHTTPHandler:
|
||||
handler: Final = AsyncHTTPHandler()
|
||||
handler.client = httpx.AsyncClient(transport=_gateway_transport(requests, _fake_claude_platform_gateway))
|
||||
handler.client = httpx.AsyncClient(transport=_gateway_transport(requests, respond))
|
||||
return handler
|
||||
|
||||
|
||||
|
|
@ -718,6 +721,112 @@ def test_chat_completion_claude_platform_override_kwarg_is_honoured_and_not_sent
|
|||
assert requests[0]["body"] == {**EXPECTED_CHAT_BODY, "context_management": context_management}
|
||||
|
||||
|
||||
_RISK_ANSWER: Final = {"tier": 2, "summary": "ok"}
|
||||
_JSON_SCHEMA_STREAM_KWARGS: Final = {
|
||||
"model": "bedrock/claude_platform/claude-scripted-9",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 10,
|
||||
"workspace_id": "wrkspc_test",
|
||||
"stream": True,
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "Risk",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
|
||||
"required": ["tier", "summary"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
},
|
||||
**API_KEY_KWARGS,
|
||||
}
|
||||
|
||||
|
||||
def _wrapped_json_tool_stream(url: str, headers: Mapping[str, str], data: bytes) -> httpx.Response:
|
||||
wrapper: Final = json.dumps({"$PARAMETER_NAME": _RISK_ANSWER})
|
||||
events: Final = (
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-scripted-9",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_test", "name": "json_tool_call", "input": {}},
|
||||
},
|
||||
*(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": wrapper[start : start + 20]},
|
||||
}
|
||||
for start in range(0, len(wrapper), 20)
|
||||
),
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 9},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode(),
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
|
||||
def _streamed_answer(chunks: tuple) -> dict:
|
||||
deltas: Final = tuple(chunk.choices[0].delta for chunk in chunks)
|
||||
return {
|
||||
"content": json.loads("".join(delta.content or "" for delta in deltas)),
|
||||
"tool_call_deltas": sum(1 for delta in deltas if delta.tool_calls),
|
||||
"finish_reasons": tuple(chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason),
|
||||
}
|
||||
|
||||
|
||||
_BARE_STREAMED_ANSWER: Final = {"content": _RISK_ANSWER, "tool_call_deltas": 0, "finish_reasons": ("stop",)}
|
||||
|
||||
|
||||
def test_chat_completion_claude_platform_streams_a_wrapped_json_schema_answer_as_the_bare_object():
|
||||
import litellm
|
||||
|
||||
chunks: Final = tuple(
|
||||
litellm.completion(
|
||||
**_JSON_SCHEMA_STREAM_KWARGS,
|
||||
client=_sync_gateway_client([], respond=_wrapped_json_tool_stream),
|
||||
)
|
||||
)
|
||||
|
||||
assert _streamed_answer(chunks) == _BARE_STREAMED_ANSWER
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chat_completion_claude_platform_streams_a_wrapped_json_schema_answer_as_the_bare_object():
|
||||
import litellm
|
||||
|
||||
stream: Final = await litellm.acompletion(
|
||||
**_JSON_SCHEMA_STREAM_KWARGS,
|
||||
client=_async_gateway_client([], respond=_wrapped_json_tool_stream),
|
||||
)
|
||||
chunks: Final = tuple([chunk async for chunk in stream])
|
||||
|
||||
assert _streamed_answer(chunks) == _BARE_STREAMED_ANSWER
|
||||
|
||||
|
||||
def test_fake_claude_platform_gateway_rejects_leaked_internal_fields():
|
||||
leaked = json.dumps({**EXPECTED_NATIVE_BODY, "workspace_id": "wrkspc_test", "aws_region_name": "us-west-2"})
|
||||
response = _fake_claude_platform_gateway(
|
||||
|
|
|
|||
|
|
@ -4228,6 +4228,54 @@ async def test_completion_streams_after_the_async_transform_request():
|
|||
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async"
|
||||
|
||||
|
||||
class _IteratorRecordingConfig(_TransformRecordingConfig):
|
||||
def __init__(self):
|
||||
super().__init__(transform_async=False)
|
||||
self.iterator_requests = []
|
||||
|
||||
def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False):
|
||||
self.iterator_requests.append({"sync_stream": sync_stream, "json_mode": json_mode})
|
||||
return super().get_model_response_iterator(streaming_response, sync_stream, json_mode)
|
||||
|
||||
|
||||
def _json_mode_stream(config, **transport):
|
||||
return BaseLLMHTTPHandler().completion(
|
||||
model="stub-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base="https://llm.example/v1/chat",
|
||||
custom_llm_provider="openai",
|
||||
model_response=ModelResponse(),
|
||||
encoding=None,
|
||||
logging_obj=_streaming_logging_obj(),
|
||||
optional_params={"json_mode": True},
|
||||
timeout=10.0,
|
||||
litellm_params={},
|
||||
stream=True,
|
||||
provider_config=config,
|
||||
**transport,
|
||||
)
|
||||
|
||||
|
||||
def _streamed_text(chunks):
|
||||
return "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
|
||||
|
||||
|
||||
async def test_json_mode_stream_reaches_the_iterator_of_a_config_without_a_request_hook_only_when_sync():
|
||||
config = _IteratorRecordingConfig()
|
||||
sync_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_sse_echoing_transformed_by)))
|
||||
async_client = AsyncHTTPHandler()
|
||||
async_client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by))
|
||||
|
||||
sync_chunks = list(_json_mode_stream(config, acompletion=False, client=sync_client))
|
||||
async_chunks = [chunk async for chunk in await _json_mode_stream(config, acompletion=True, client=async_client)]
|
||||
|
||||
assert (_streamed_text(sync_chunks), _streamed_text(async_chunks)) == ("sync", "sync")
|
||||
assert config.iterator_requests == [
|
||||
{"sync_stream": True, "json_mode": True},
|
||||
{"sync_stream": False, "json_mode": False},
|
||||
]
|
||||
|
||||
|
||||
class _ImageEditRecordingConfig(BaseImageEditConfig):
|
||||
def __init__(self):
|
||||
self.transform_calls = []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue