mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
Merge remote-tracking branch 'origin/litellm_responses_tool_search_anthropic' into litellm_chat_completions_tool_search_anthropic
This commit is contained in:
mateo-berri
commit
9150723c84
6 changed files with 363 additions and 122 deletions
|
|
@ -25,7 +25,6 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
|
|||
anthropic_system_messages,
|
||||
find_anthropic_server_tool_result,
|
||||
is_anthropic_tool_search_use,
|
||||
tool_search_ids_dropped_with_their_answers,
|
||||
)
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
|
||||
|
|
@ -2475,7 +2474,6 @@ def anthropic_messages_pt(
|
|||
is_bedrock_invoke = model.lower().startswith("invoke/")
|
||||
is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False
|
||||
force_base64 = is_bedrock_invoke or is_vertex_ai
|
||||
dropped_tool_search_ids: Final = tool_search_ids_dropped_with_their_answers(messages)
|
||||
|
||||
msg_i = 0
|
||||
while msg_i < len(messages):
|
||||
|
|
@ -2576,9 +2574,7 @@ def anthropic_messages_pt(
|
|||
|
||||
user_content.append(_anthropic_content_text_element)
|
||||
|
||||
elif (
|
||||
user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function"
|
||||
) and user_message_types_block.get("tool_call_id") not in dropped_tool_search_ids:
|
||||
elif user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function":
|
||||
# OpenAI's tool message content will always be a string
|
||||
user_content.append(
|
||||
convert_to_anthropic_tool_result(user_message_types_block, force_base64=force_base64)
|
||||
|
|
@ -2593,7 +2589,7 @@ def anthropic_messages_pt(
|
|||
unique_tool_ids: set[str] = set()
|
||||
|
||||
assistant_content: list[AnthropicMessagesAssistantMessageValues] = []
|
||||
answered_tool_ids = answered_tool_call_ids(messages, msg_i) - dropped_tool_search_ids
|
||||
answered_tool_ids = answered_tool_call_ids(messages, msg_i)
|
||||
## MERGE CONSECUTIVE ASSISTANT CONTENT ##
|
||||
while msg_i < len(messages) and messages[msg_i]["role"] == "assistant":
|
||||
assistant_content_block: ChatCompletionAssistantMessage = messages[msg_i]
|
||||
|
|
|
|||
|
|
@ -30,7 +30,8 @@ Anthropic wire shape is built later by ``anthropic_messages_pt``.
|
|||
"""
|
||||
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from itertools import chain, dropwhile, groupby, islice, takewhile
|
||||
from itertools import chain, dropwhile, groupby, takewhile
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias, TypeVar
|
||||
|
||||
from litellm.types.llms.anthropic import (
|
||||
|
|
@ -55,6 +56,7 @@ CONVERTED_SYSTEM_NOTE: Final = (
|
|||
|
||||
_USER_TYPE_ROLES: Final = frozenset({"user", "tool", "function"})
|
||||
_TOOL_ROLES: Final = frozenset({"tool", "function"})
|
||||
_NO_IDS: Final[frozenset[str]] = frozenset()
|
||||
_RENDERED_PART_TYPES: Final = frozenset({"text", "image_url", "document", "file"})
|
||||
_RENDERED_ASSISTANT_PART_TYPES: Final = frozenset({"text", "server_tool_use"})
|
||||
_THINKING_BLOCK_TYPES: Final = frozenset({"thinking", "redacted_thinking"})
|
||||
|
|
@ -313,12 +315,46 @@ def _tool_result_ids(answers: Iterable[object]) -> Iterator[str]:
|
|||
|
||||
|
||||
def answered_tool_call_ids(messages: Sequence[object], assistant_index: int) -> frozenset[str]:
|
||||
from_assistant: Final = islice(messages, assistant_index, None)
|
||||
from_assistant: Final = (messages[index] for index in range(assistant_index, len(messages)))
|
||||
following: Final = dropwhile(lambda message: message_field(message, "role") == "assistant", from_assistant)
|
||||
answers: Final = takewhile(lambda message: message_field(message, "role") in _TOOL_ROLES, following)
|
||||
return frozenset(_tool_result_ids(answers))
|
||||
|
||||
|
||||
def _role_group(message: object) -> str:
|
||||
role: Final = message_field(message, "role")
|
||||
if role == "assistant":
|
||||
return "assistant"
|
||||
return "tool" if role in _TOOL_ROLES else "other"
|
||||
|
||||
|
||||
def _answered_by_assistant_index(
|
||||
groups: Sequence[tuple[str, tuple[tuple[int, object], ...]]], group_index: int
|
||||
) -> Iterator[tuple[int, frozenset[str]]]:
|
||||
following: Final = groups[group_index + 1] if group_index + 1 < len(groups) else ("other", ())
|
||||
answered: Final = (
|
||||
frozenset(_tool_result_ids(message for _, message in following[1])) if following[0] == "tool" else _NO_IDS
|
||||
)
|
||||
return ((index, answered) for index, _ in groups[group_index][1])
|
||||
|
||||
|
||||
def _answered_ids_by_index(messages: Sequence[object]) -> Mapping[int, frozenset[str]]:
|
||||
"""``answered_tool_call_ids`` for every assistant message, computed in one pass over the messages."""
|
||||
groups: Final = tuple(
|
||||
(role_group, tuple(group))
|
||||
for role_group, group in groupby(enumerate(messages), key=lambda pair: _role_group(pair[1]))
|
||||
)
|
||||
return MappingProxyType(
|
||||
dict(
|
||||
chain.from_iterable(
|
||||
_answered_by_assistant_index(groups, group_index)
|
||||
for group_index, (role_group, _) in enumerate(groups)
|
||||
if role_group == "assistant"
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _unresolved_tool_search_id(tool_call: object, message: object) -> str | None:
|
||||
tool_id: Final = message_field(tool_call, "id")
|
||||
tool_name: Final = message_field(message_field(tool_call, "function"), "name")
|
||||
|
|
@ -343,42 +379,15 @@ def _unresolved_tool_search_ids(message: object) -> tuple[str, ...]:
|
|||
)
|
||||
|
||||
|
||||
def _is_client_tool_call(tool_call: object) -> bool:
|
||||
tool_id: Final = message_field(tool_call, "id")
|
||||
return isinstance(tool_id, str) and not tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
|
||||
|
||||
|
||||
def _tool_search_ids_beside_client_calls(message: object) -> tuple[str, ...]:
|
||||
def _tool_calls_render(message: object, answered: frozenset[str]) -> bool:
|
||||
tool_calls: Final = parts_of(message_field(message, "tool_calls"))
|
||||
return _unresolved_tool_search_ids(message) if any(map(_is_client_tool_call, tool_calls)) else ()
|
||||
|
||||
|
||||
def tool_search_ids_dropped_with_their_answers(messages: Iterable[object]) -> frozenset[str]:
|
||||
"""Tool search uses ``anthropic_messages_pt`` leaves out together with the caller's answer to them.
|
||||
|
||||
The search already ran on Anthropic's side, so a caller's ``tool`` message for it
|
||||
is never the real result. Dropping the pair is safe only beside a client tool call,
|
||||
whose result keeps a user turn after the assistant turn: Claude rejects a
|
||||
conversation that ends on an assistant turn.
|
||||
"""
|
||||
return frozenset(
|
||||
chain.from_iterable(
|
||||
_tool_search_ids_beside_client_calls(message)
|
||||
for message in messages
|
||||
if message_field(message, "role") == "assistant"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _tool_calls_render(messages: Sequence[object], index: int) -> bool:
|
||||
tool_calls: Final = parts_of(message_field(messages[index], "tool_calls"))
|
||||
search_ids: Final = _unresolved_tool_search_ids(messages[index])
|
||||
search_ids: Final = _unresolved_tool_search_ids(message)
|
||||
if len(search_ids) < len(tool_calls):
|
||||
return True
|
||||
return bool(search_ids) and not answered_tool_call_ids(messages, index).isdisjoint(search_ids)
|
||||
return bool(search_ids) and not answered.isdisjoint(search_ids)
|
||||
|
||||
|
||||
def _assistant_renders(messages: Sequence[object], index: int) -> bool:
|
||||
def _assistant_renders(message: object, answered: frozenset[str]) -> bool:
|
||||
"""Whether ``anthropic_messages_pt`` puts a block on the wire for this assistant message.
|
||||
|
||||
String content (the converter pads an empty one with a placeholder), a text part,
|
||||
|
|
@ -388,7 +397,6 @@ def _assistant_renders(messages: Sequence[object], index: int) -> bool:
|
|||
message with none of them, such as ``content: None`` or an empty list, vanishes
|
||||
from the wire.
|
||||
"""
|
||||
message: Final = messages[index]
|
||||
content: Final = message_field(message, "content")
|
||||
if isinstance(content, str):
|
||||
return True
|
||||
|
|
@ -396,13 +404,13 @@ def _assistant_renders(messages: Sequence[object], index: int) -> bool:
|
|||
return (
|
||||
any(_assistant_part_renders(part) for part in parts)
|
||||
or _separate_thinking_blocks_render(message, parts)
|
||||
or _tool_calls_render(messages, index)
|
||||
or _tool_calls_render(message, answered)
|
||||
or bool(message_field(message, "function_call"))
|
||||
or bool(message_field(message_field(message, "provider_specific_fields"), "compaction_blocks"))
|
||||
)
|
||||
|
||||
|
||||
def _renders(messages: Sequence[object], index: int) -> bool:
|
||||
def _renders(message: object, answered: frozenset[str]) -> bool:
|
||||
"""Whether ``anthropic_messages_pt`` puts a block on the wire for this message.
|
||||
|
||||
A tool message always becomes a ``tool_result`` and a user message with string
|
||||
|
|
@ -410,31 +418,35 @@ def _renders(messages: Sequence[object], index: int) -> bool:
|
|||
renders only through parts of a type the converter emits; ``None``, an empty list,
|
||||
and a list of other parts vanish. Assistant messages follow ``_assistant_renders``.
|
||||
"""
|
||||
message: Final = messages[index]
|
||||
role: Final = message_field(message, "role")
|
||||
if role in _TOOL_ROLES:
|
||||
return True
|
||||
if role == "assistant":
|
||||
return _assistant_renders(messages, index)
|
||||
return _assistant_renders(message, answered)
|
||||
content: Final = message_field(message, "content")
|
||||
return isinstance(content, str) or any(
|
||||
message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)
|
||||
)
|
||||
|
||||
|
||||
def _rendered_indices(messages: Sequence[object]) -> frozenset[int]:
|
||||
answered: Final = _answered_ids_by_index(messages)
|
||||
return frozenset(index for index, message in enumerate(messages) if _renders(message, answered.get(index, _NO_IDS)))
|
||||
|
||||
|
||||
def _rendered_block(
|
||||
message_index: int,
|
||||
messages: Sequence[AllMessageValues],
|
||||
rendered: frozenset[int],
|
||||
blocks: Sequence[tuple[bool, tuple[int, ...]]],
|
||||
) -> int | None:
|
||||
block_index: Final = _block_containing(message_index, blocks)
|
||||
_, indices = blocks[block_index]
|
||||
return block_index if any(_renders(messages, index) for index in indices) else None
|
||||
return block_index if not rendered.isdisjoint(indices) else None
|
||||
|
||||
|
||||
def _system_may_follow(
|
||||
block_index: int,
|
||||
messages: Sequence[AllMessageValues],
|
||||
rendered: frozenset[int],
|
||||
blocks: Sequence[tuple[bool, tuple[int, ...]]],
|
||||
) -> bool:
|
||||
"""Whether a system message behind this block precedes an assistant turn or ends the array on the wire.
|
||||
|
|
@ -443,11 +455,7 @@ def _system_may_follow(
|
|||
first later block that puts anything on the wire is an assistant block.
|
||||
"""
|
||||
return next(
|
||||
(
|
||||
not is_user
|
||||
for is_user, indices in blocks[block_index + 1 :]
|
||||
if any(_renders(messages, index) for index in indices)
|
||||
),
|
||||
(not is_user for is_user, indices in blocks[block_index + 1 :] if not rendered.isdisjoint(indices)),
|
||||
True,
|
||||
)
|
||||
|
||||
|
|
@ -455,6 +463,7 @@ def _system_may_follow(
|
|||
def _anchor_block(
|
||||
run: Sequence[int],
|
||||
messages: Sequence[AllMessageValues],
|
||||
rendered: frozenset[int],
|
||||
blocks: Sequence[tuple[bool, tuple[int, ...]]],
|
||||
) -> int | None:
|
||||
"""The user-type block a system run must follow, or ``None`` when it converts in place.
|
||||
|
|
@ -472,8 +481,8 @@ def _anchor_block(
|
|||
neighbour: Final = previous if _is_user_type(messages[previous]) else run[-1] + 1
|
||||
if neighbour >= len(messages) or not _is_user_type(messages[neighbour]):
|
||||
return None
|
||||
block_index: Final = _rendered_block(neighbour, messages, blocks)
|
||||
if block_index is None or not _system_may_follow(block_index, messages, blocks):
|
||||
block_index: Final = _rendered_block(neighbour, rendered, blocks)
|
||||
if block_index is None or not _system_may_follow(block_index, rendered, blocks):
|
||||
return None
|
||||
return block_index
|
||||
|
||||
|
|
@ -488,7 +497,8 @@ def _placed_for_flagged_model(messages: Sequence[AllMessageValues]) -> tuple[All
|
|||
move it. Runs that share a user turn merge into one system message.
|
||||
"""
|
||||
blocks: Final = _user_type_blocks(messages)
|
||||
anchors: Final = tuple((run, _anchor_block(run, messages, blocks)) for run in _system_runs(messages))
|
||||
rendered: Final = _rendered_indices(messages)
|
||||
anchors: Final = tuple((run, _anchor_block(run, messages, rendered, blocks)) for run in _system_runs(messages))
|
||||
|
||||
def messages_of(run: tuple[int, ...]) -> tuple[AllMessageValues, ...]:
|
||||
return tuple(messages[index] for index in run)
|
||||
|
|
|
|||
|
|
@ -29,9 +29,10 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
|
|||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Mapping, MutableSequence, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, chain, repeat
|
||||
from types import MappingProxyType
|
||||
|
|
@ -230,6 +231,10 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
|
|||
|
||||
|
||||
_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
|
||||
_TOOL_SEARCH_CALL_TYPE: Final = "tool_search_call"
|
||||
_OUTPUT_TOOL_CALL_ITEM_TYPES: Final = _TOOL_CALL_ITEM_TYPES | {_TOOL_SEARCH_CALL_TYPE}
|
||||
_SCANNED_TOOL_SEARCH_NAME: Final = "tool_search"
|
||||
_TOOL_SEARCH_ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"function_call": "arguments", "custom_tool_call": "input"}
|
||||
)
|
||||
|
|
@ -285,7 +290,7 @@ def _is_tool_call_item(item: object) -> bool:
|
|||
|
||||
|
||||
def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
|
||||
if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
|
||||
if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES:
|
||||
return None
|
||||
if isinstance(item, Mapping):
|
||||
return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects
|
||||
|
|
@ -296,6 +301,40 @@ def _is_tool_call_output_item(item: object) -> bool:
|
|||
return _tool_call_output_item_mapping(item) is not None
|
||||
|
||||
|
||||
def _indexed_tool_call_output_items(outputs: Sequence[object]) -> tuple[tuple[int, object], ...]:
|
||||
return tuple((index, item) for index, item in enumerate(outputs) if _is_tool_call_output_item(item))
|
||||
|
||||
|
||||
def _scanned_tool_call_item(item: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""A ``tool_search_call`` reaches guardrails as a ``tool_search`` call whose arguments are JSON text."""
|
||||
if item.get("type") != _TOOL_SEARCH_CALL_TYPE:
|
||||
return item
|
||||
return {**item, "name": _SCANNED_TOOL_SEARCH_NAME, "arguments": json.dumps(item.get("arguments") or {})}
|
||||
|
||||
|
||||
def _tool_search_rewrite(payload: str) -> Mapping[str, object]:
|
||||
try:
|
||||
return {"arguments": _TOOL_SEARCH_ARGUMENTS_ADAPTER.validate_json(payload)}
|
||||
except ValidationError:
|
||||
return {"arguments": {}, "status": "incomplete"}
|
||||
|
||||
|
||||
def _rewritten_tool_search_call(item: object, payload: str | None) -> object:
|
||||
if payload is None:
|
||||
return item
|
||||
update: Final = _tool_search_rewrite(payload)
|
||||
if isinstance(item, BaseModel):
|
||||
return item.model_copy(update=update)
|
||||
if isinstance(item, Mapping):
|
||||
return {**item, **update}
|
||||
return item
|
||||
|
||||
|
||||
def _deliver_into_output(outputs: Sequence[object], index: int, item: object, delivered: object) -> None:
|
||||
if delivered is not item and isinstance(outputs, MutableSequence):
|
||||
outputs[index] = delivered # rebind-ok: a frozen output item takes a rewrite by being swapped in the list
|
||||
|
||||
|
||||
def _released_tool_call_payload(responses_so_far: Sequence[object], item_id: object) -> str | None:
|
||||
events: Final = tuple(event for event in responses_so_far if stream_item_field(event, "item_id") == item_id)
|
||||
finished: Final = tuple(
|
||||
|
|
@ -317,7 +356,9 @@ def _released_tool_call_payload(responses_so_far: Sequence[object], item_id: obj
|
|||
|
||||
|
||||
def _with_released_payload(item: Mapping[str, object], responses_so_far: Sequence[object]) -> Mapping[str, object]:
|
||||
field: Final = _TOOL_CALL_PAYLOAD_FIELDS[str(item.get("type"))]
|
||||
field: Final = _TOOL_CALL_PAYLOAD_FIELDS.get(str(item.get("type")))
|
||||
if field is None:
|
||||
return item
|
||||
payload: Final = _released_tool_call_payload(responses_so_far, item.get("id"))
|
||||
return item if payload is None else {**item, field: payload}
|
||||
|
||||
|
|
@ -921,7 +962,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
task_mappings=task_mappings,
|
||||
)
|
||||
self._write_tool_call_rewrites_to_output(
|
||||
tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
|
||||
outputs=response_output,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=post_guardrail_tool_calls,
|
||||
)
|
||||
|
|
@ -1156,7 +1197,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_event_field(event: object, field: str, value: str) -> None:
|
||||
def _write_event_field(event: object, field: str, value: object) -> None:
|
||||
if isinstance(event, dict):
|
||||
event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place
|
||||
else:
|
||||
|
|
@ -1234,7 +1275,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
pipeline executor discards it and releases the original events."""
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
|
||||
indexed_tool_call_items: Final = _indexed_tool_call_output_items(outputs)
|
||||
tool_call_items: Final = tuple(output_item for _, output_item in indexed_tool_call_items)
|
||||
call_ids: Final = tuple(
|
||||
call_id
|
||||
for output_item in tool_call_items
|
||||
|
|
@ -1268,12 +1310,14 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(guardrail_name, undeliverable_reason)
|
||||
for output_item, rewrite in (
|
||||
(output_item, rewrites_by_call_id[call_id])
|
||||
for output_item, call_id in zip(tool_call_items, call_ids)
|
||||
for (index, output_item), rewrite in (
|
||||
(indexed_item, rewrites_by_call_id[call_id])
|
||||
for indexed_item, call_id in zip(indexed_tool_call_items, call_ids)
|
||||
if call_id in rewrites_by_call_id
|
||||
):
|
||||
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
_deliver_into_output(
|
||||
outputs, index, output_item, self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
)
|
||||
delta_replacements: Final = MappingProxyType(
|
||||
{call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
|
||||
)
|
||||
|
|
@ -1292,28 +1336,34 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
|
||||
)
|
||||
case "response.output_item.done":
|
||||
self._write_tool_call_item(
|
||||
stream_item_field(event, "item"),
|
||||
rewrites_by_call_id[call_id].name,
|
||||
rewrites_by_call_id[call_id].arguments,
|
||||
done_item = stream_item_field(event, "item")
|
||||
delivered_item = self._write_tool_call_item(
|
||||
done_item, rewrites_by_call_id[call_id].name, rewrites_by_call_id[call_id].arguments
|
||||
)
|
||||
if delivered_item is not done_item:
|
||||
self._write_event_field(event, "item", delivered_item)
|
||||
case _:
|
||||
pass
|
||||
|
||||
def _write_tool_call_rewrites_to_output(
|
||||
self,
|
||||
tool_call_items: Sequence[object],
|
||||
outputs: Sequence[object],
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
) -> None:
|
||||
if len(tool_call_items) != len(post_guardrail_tool_calls):
|
||||
indexed_tool_call_items: Final = _indexed_tool_call_output_items(outputs)
|
||||
if len(indexed_tool_call_items) != len(post_guardrail_tool_calls):
|
||||
return
|
||||
for output_item, rewrite in (
|
||||
(output_item, _tool_call_rewrite(before, after))
|
||||
for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
for (index, output_item), rewrite in (
|
||||
(indexed_item, _tool_call_rewrite(before, after))
|
||||
for indexed_item, before, after in zip(
|
||||
indexed_tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls
|
||||
)
|
||||
if after != before
|
||||
):
|
||||
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
_deliver_into_output(
|
||||
outputs, index, output_item, self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
|
||||
|
|
@ -1343,18 +1393,28 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
item: Final = stream_item_field(event, "item")
|
||||
call_id: Final = stream_item_field(item, "call_id")
|
||||
return (
|
||||
call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
|
||||
call_id
|
||||
if stream_item_field(item, "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES and isinstance(call_id, str)
|
||||
else None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
|
||||
def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> object:
|
||||
"""Apply a rewrite and return the item that carries it.
|
||||
|
||||
A ``tool_search_call`` keeps its name and takes the rewritten arguments as a JSON
|
||||
object on a copy, since its item may be frozen. Other tool calls are edited in place.
|
||||
"""
|
||||
item_type: Final = stream_item_field(item, "type")
|
||||
if item is None:
|
||||
return
|
||||
return item
|
||||
if item_type == _TOOL_SEARCH_CALL_TYPE:
|
||||
return _rewritten_tool_search_call(item, payload)
|
||||
if name is not None:
|
||||
OpenAIResponsesHandler._write_event_field(item, "name", name)
|
||||
item_type: Final = stream_item_field(item, "type")
|
||||
if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
|
||||
OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
|
||||
return item
|
||||
|
||||
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
|
||||
"""
|
||||
|
|
@ -1405,7 +1465,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
|
||||
or (
|
||||
stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
|
||||
and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES
|
||||
and stream_item_field(stream_item_field(event, "item"), "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES
|
||||
)
|
||||
for event in responses_so_far
|
||||
)
|
||||
|
|
@ -1414,7 +1474,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
def _completed_response_scan_key(response: object) -> StreamingScanKey:
|
||||
output_items: Final = stream_item_items(response, "output")
|
||||
message_items: Final = tuple(
|
||||
item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
|
||||
item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES
|
||||
)
|
||||
return StreamingScanKey(
|
||||
texts=tuple(
|
||||
|
|
@ -1426,7 +1486,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
tool_calls=tuple(
|
||||
stream_item_fingerprint(item)
|
||||
for item in output_items
|
||||
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
|
||||
if stream_item_field(item, "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES
|
||||
),
|
||||
stream_ended=True,
|
||||
)
|
||||
|
|
@ -1540,7 +1600,9 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
tool_call_item: Final = _tool_call_output_item_mapping(output_item)
|
||||
if tool_call_item is not None:
|
||||
if tool_calls_to_check is not None:
|
||||
tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
|
||||
tool_calls_to_check.append(
|
||||
tool_call_dict_from_output_item(_scanned_tool_call_item(tool_call_item), output_idx)
|
||||
)
|
||||
return
|
||||
|
||||
# Handle both GenericResponseOutputItem and dict
|
||||
|
|
|
|||
|
|
@ -7144,22 +7144,18 @@ def test_anthropic_messages_pt_drops_replayed_tool_search_call_without_its_resul
|
|||
]
|
||||
|
||||
|
||||
def _tool_search_call() -> dict:
|
||||
return {
|
||||
"id": "srvtoolu_01Search",
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'},
|
||||
}
|
||||
|
||||
|
||||
def test_anthropic_messages_pt_drops_an_answered_tool_search_call_and_its_answer_beside_a_client_call():
|
||||
def test_anthropic_messages_pt_keeps_a_replayed_tool_search_call_the_caller_answered():
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
_tool_search_call(),
|
||||
{
|
||||
"id": "srvtoolu_01Search",
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'},
|
||||
},
|
||||
{
|
||||
"id": "toolu_01Weather",
|
||||
"type": "function",
|
||||
|
|
@ -7167,32 +7163,22 @@ def test_anthropic_messages_pt_drops_an_answered_tool_search_call_and_its_answer
|
|||
},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "Error: unknown tool tool_search_tool_regex"},
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": '[{"name": "get_weather"}]'},
|
||||
{"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C and cloudy"},
|
||||
]
|
||||
|
||||
result: Final = anthropic_messages_pt(messages=messages, model="claude-sonnet-5-5", llm_provider="anthropic")
|
||||
|
||||
assert [(block["type"], block["id"]) for block in result[1]["content"]] == [("tool_use", "toolu_01Weather")]
|
||||
assert [(block["type"], block["id"]) for block in result[1]["content"]] == [
|
||||
("tool_use", "srvtoolu_01Search"),
|
||||
("tool_use", "toolu_01Weather"),
|
||||
]
|
||||
assert [block["tool_use_id"] for block in result[2]["content"] if block["type"] == "tool_result"] == [
|
||||
"toolu_01Weather"
|
||||
"srvtoolu_01Search",
|
||||
"toolu_01Weather",
|
||||
]
|
||||
|
||||
|
||||
def test_anthropic_messages_pt_keeps_an_answered_tool_search_call_that_is_the_turns_only_call():
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{"role": "assistant", "content": "Let me look for a weather tool.", "tool_calls": [_tool_search_call()]},
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "Error: unknown tool tool_search_tool_regex"},
|
||||
]
|
||||
|
||||
result: Final = anthropic_messages_pt(messages=messages, model="claude-sonnet-5-5", llm_provider="anthropic")
|
||||
|
||||
assert [message["role"] for message in result] == ["user", "assistant", "user"]
|
||||
assert [block["type"] for block in result[1]["content"]] == ["text", "tool_use"]
|
||||
assert [block["tool_use_id"] for block in result[2]["content"]] == ["srvtoolu_01Search"]
|
||||
|
||||
|
||||
def _merged_assistant_run_cost(run_length: int) -> float:
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "Plan the trip"},
|
||||
|
|
|
|||
|
|
@ -5,6 +5,10 @@ Anthropic, Vertex, Azure AI and Bedrock Invoke transformation tests; these pin
|
|||
the pure placement rules on the OpenAI-format message list.
|
||||
"""
|
||||
|
||||
import time
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
|
@ -255,6 +259,13 @@ def test_flagged_placement_keeps_a_system_behind_a_user_turn_merged_with_an_empt
|
|||
assert _roles(placed) == ["user", "user", "system", "assistant"]
|
||||
|
||||
|
||||
_TOOL_SEARCH_USE = {
|
||||
"id": "srvtoolu_01Search",
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'},
|
||||
}
|
||||
|
||||
|
||||
EMPTY_ASSISTANT = pytest.mark.parametrize(
|
||||
"empty_assistant",
|
||||
[
|
||||
|
|
@ -281,17 +292,7 @@ EMPTY_ASSISTANT = pytest.mark.parametrize(
|
|||
"thinking_blocks": [{"type": "redacted_thinking", "data": "x"}],
|
||||
},
|
||||
litellm.Message(role="assistant", content=None),
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "srvtoolu_01Search",
|
||||
"type": "function",
|
||||
"function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_TOOL_SEARCH_USE]},
|
||||
],
|
||||
ids=[
|
||||
"none",
|
||||
|
|
@ -307,6 +308,50 @@ EMPTY_ASSISTANT = pytest.mark.parametrize(
|
|||
)
|
||||
|
||||
|
||||
def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time():
|
||||
search_uses: Final = tuple(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{**_TOOL_SEARCH_USE, "id": f"srvtoolu_{index:06d}"}],
|
||||
}
|
||||
for index in range(3000)
|
||||
)
|
||||
reminders: Final = tuple(
|
||||
chain.from_iterable(
|
||||
({"role": "system", "content": f"reminder {index}"}, {"role": "user", "content": f"q{index}"})
|
||||
for index in range(10)
|
||||
)
|
||||
)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "start"},
|
||||
*reminders,
|
||||
*search_uses,
|
||||
{"role": "user", "content": "end"},
|
||||
]
|
||||
|
||||
started: Final = time.perf_counter()
|
||||
placed: Final = place_mid_conversation_system(messages, supports_mid_conversation_system=True)
|
||||
elapsed: Final = time.perf_counter() - started
|
||||
|
||||
assert set(_roles(placed)[:21]) == {"user"}
|
||||
assert elapsed < 1.0
|
||||
|
||||
|
||||
def test_flagged_placement_keeps_a_system_before_a_tool_search_use_the_caller_answered():
|
||||
placed = place_mid_conversation_system(
|
||||
[
|
||||
{"role": "user", "content": "q1"},
|
||||
{"role": "system", "content": "reminder"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_TOOL_SEARCH_USE]},
|
||||
{"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "[]"},
|
||||
],
|
||||
supports_mid_conversation_system=True,
|
||||
)
|
||||
|
||||
assert _roles(placed) == ["user", "system", "assistant", "tool"]
|
||||
|
||||
|
||||
@EMPTY_ASSISTANT
|
||||
def test_flagged_placement_converts_a_run_when_the_assistant_turn_after_its_anchor_sends_nothing(empty_assistant):
|
||||
placed = place_mid_conversation_system(
|
||||
|
|
@ -355,6 +400,14 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th
|
|||
content="",
|
||||
tool_calls=[{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}],
|
||||
),
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
_TOOL_SEARCH_USE,
|
||||
{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}},
|
||||
],
|
||||
},
|
||||
],
|
||||
ids=[
|
||||
"tool-calls",
|
||||
|
|
@ -363,6 +416,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th
|
|||
"signed-thinking-part",
|
||||
"function-call",
|
||||
"pydantic-tool-calls",
|
||||
"unanswered-tool-search-use-beside-a-client-call",
|
||||
],
|
||||
)
|
||||
def test_flagged_placement_keeps_a_system_before_an_assistant_turn_that_renders_without_text(assistant_turn):
|
||||
|
|
|
|||
|
|
@ -38,7 +38,13 @@ from litellm.responses.litellm_completion_transformation.transformation import (
|
|||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText
|
||||
from litellm.responses.litellm_completion_transformation.tool_search import build_tool_search_call
|
||||
from litellm.types.responses.main import (
|
||||
CustomToolCallOutputItem,
|
||||
GenericResponseOutputItem,
|
||||
OutputText,
|
||||
OutputToolSearchCall,
|
||||
)
|
||||
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
|
|
@ -138,6 +144,40 @@ class PersimmonMaskingGuardrail(CustomGuardrail):
|
|||
return {**inputs, "tool_calls": tool_calls}
|
||||
|
||||
|
||||
class RecordingPersimmonMaskingGuardrail(PersimmonMaskingGuardrail):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.seen_tool_calls: list[tuple[str, str]] = []
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.seen_tool_calls.extend(
|
||||
(tool_call["function"]["name"], tool_call["function"]["arguments"])
|
||||
for tool_call in inputs.get("tool_calls", [])
|
||||
)
|
||||
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
|
||||
|
||||
|
||||
class ArgumentsRedactingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
tool_calls = [
|
||||
{**tool_call, "function": {**tool_call["function"], "arguments": "[REDACTED]"}}
|
||||
for tool_call in inputs.get("tool_calls", [])
|
||||
]
|
||||
return {**inputs, "tool_calls": tool_calls}
|
||||
|
||||
|
||||
class FlatShapeGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
@ -1064,6 +1104,49 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
|
|||
assert result["output"][0]["input"] == "echo [MASKED]"
|
||||
assert "name" not in result["output"][0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_scans_and_rewrites_a_tool_search_call(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
guardrail = RecordingPersimmonMaskingGuardrail(guardrail_name="mask")
|
||||
response = ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=1,
|
||||
model="claude-sonnet-5-5",
|
||||
object="response",
|
||||
status="completed",
|
||||
output=[build_tool_search_call("srvtoolu_1", '{"query": "persimmon"}', "completed")],
|
||||
)
|
||||
|
||||
result = await handler.process_output_response(response, guardrail)
|
||||
|
||||
assert guardrail.seen_tool_calls == [("tool_search", '{"query": "persimmon"}')]
|
||||
search_call = result.output[0]
|
||||
assert isinstance(search_call, OutputToolSearchCall)
|
||||
assert search_call.arguments == {"query": "[MASKED]"}
|
||||
assert search_call.call_id == "srvtoolu_1"
|
||||
assert "name" not in search_call.model_dump()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_never_returns_original_search_arguments_after_a_non_json_rewrite(self):
|
||||
response = ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=1,
|
||||
model="claude-sonnet-5-5",
|
||||
object="response",
|
||||
status="completed",
|
||||
output=[build_tool_search_call("srvtoolu_1", '{"query": "persimmon"}', "completed")],
|
||||
)
|
||||
|
||||
result = await OpenAIResponsesHandler().process_output_response(
|
||||
response, ArgumentsRedactingGuardrail(guardrail_name="redact")
|
||||
)
|
||||
|
||||
search_call = result.output[0]
|
||||
assert isinstance(search_call, OutputToolSearchCall)
|
||||
assert search_call.arguments == {}
|
||||
assert search_call.status == "incomplete"
|
||||
assert "persimmon" not in result.model_dump_json()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_with_tool_calls(self):
|
||||
"""Test processing output response containing function tool calls"""
|
||||
|
|
@ -1857,6 +1940,56 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
|
|||
assert completed_event.response.output[0].input == "echo [MASKED]"
|
||||
assert completed_event.response.output[0].name == "exec"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_a_tool_search_call(self):
|
||||
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
|
||||
from litellm.types.llms.openai import (
|
||||
OutputItemAddedEvent,
|
||||
OutputItemDoneEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
def search_call(arguments: str, status: Literal["in_progress", "completed"]) -> OutputToolSearchCall:
|
||||
return build_tool_search_call("srvtoolu_1", arguments, status)
|
||||
|
||||
handler = OpenAIResponsesHandler()
|
||||
added = OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=0,
|
||||
item=BaseLiteLLMOpenAIResponseObject(**search_call("", "in_progress").model_dump()),
|
||||
)
|
||||
done = OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=0,
|
||||
sequence_number=2,
|
||||
item=BaseLiteLLMOpenAIResponseObject(**search_call('{"query": "persimmon"}', "completed").model_dump()),
|
||||
)
|
||||
completed = ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=1,
|
||||
model="claude-sonnet-5-5",
|
||||
object="response",
|
||||
status="completed",
|
||||
output=[search_call('{"query": "persimmon"}', "completed")],
|
||||
),
|
||||
)
|
||||
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=[added, done, completed],
|
||||
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
|
||||
litellm_logging_obj=None,
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert added.item.model_dump()["arguments"] == {}
|
||||
assert done.item.model_dump()["arguments"] == {"query": "[MASKED]"}
|
||||
assert isinstance(completed.response.output[0], OutputToolSearchCall)
|
||||
assert completed.response.output[0].arguments == {"query": "[MASKED]"}
|
||||
assert "name" not in done.item.model_dump()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue