diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ad46121e6d9..6ae3bc8c300 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -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] diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 3fa4bd25b1a..45fea8119b7 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -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) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 8815e1165ca..cd196e0c09a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -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 diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index e0abc14c6bd..14fea46516e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -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"}, diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 912b4687e62..8f86e11270b 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -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): diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 7a9d36f330c..17740f2a39e 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -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