Merge remote-tracking branch 'origin/litellm_responses_tool_search_anthropic' into litellm_chat_completions_tool_search_anthropic

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 (
+            elif user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function":
-                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:
                 # 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]

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:
+def _tool_calls_render(message: object, answered: frozenset[str]) -> 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, ...]:
     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 ()
+    search_ids: Final = _unresolved_tool_search_ids(message)
-
-
-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])
     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 not rendered.isdisjoint(indices)),
-            not is_user
-            for is_user, indices in blocks[block_index + 1 :]
-            if any(_renders(messages, index) for index in 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)
+    block_index: Final = _rendered_block(neighbour, rendered, blocks)
-    if block_index is None or not _system_may_follow(block_index, messages, 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)

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 (
+        for (index, output_item), rewrite in (
-            (output_item, rewrites_by_call_id[call_id])
+            (indexed_item, rewrites_by_call_id[call_id])
-            for output_item, call_id in zip(tool_call_items, call_ids)
+            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(
+                    done_item = stream_item_field(event, "item")
-                        stream_item_field(event, "item"),
+                    delivered_item = self._write_tool_call_item(
-                        rewrites_by_call_id[call_id].name,
+                        done_item, rewrites_by_call_id[call_id].name, rewrites_by_call_id[call_id].arguments
-                        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 (
+        for (index, output_item), rewrite in (
-            (output_item, _tool_call_rewrite(before, after))
+            (indexed_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 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

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:
+def test_anthropic_messages_pt_keeps_a_replayed_tool_search_call_the_caller_answered():
-    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():
     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"},

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": [_TOOL_SEARCH_USE]},
-            "role": "assistant",
-            "content": None,
-            "tool_calls": [
-                {
-                    "id": "srvtoolu_01Search",
-                    "type": "function",
-                    "function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'},
-                }
-            ],
-        },
     ],
     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):

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