From cc12dbc5fc863f649a0d9d25d1b5fed71089e8aa Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sat, 10 Oct 2026 23:53:38 +0530 Subject: [PATCH 1/3] keep web search calls replayable and out of the way of the answer --- .../openai/openai_response_target.py | 51 ++++---- .../unit/backend/test_conversation_editor.py | 5 +- .../target/test_openai_response_adapters.py | 12 +- .../target/test_openai_response_target.py | 109 +++++++++++++++++- .../target/test_tool_call_history.py | 12 +- 5 files changed, 156 insertions(+), 33 deletions(-) diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 3c56aedcfd..70681e4003 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -89,6 +89,10 @@ class MessagePieceType(str, Enum): MCP_APPROVAL_REQUEST = "mcp_approval_request" +# Fields of a web_search_call input item (all required); the output item has the same ones. +_WEB_SEARCH_CALL_INPUT_FIELDS: tuple[str, ...] = ("id", "type", "status", "action") + + class _ResponseToolCallContent(BaseModel): model_config = ConfigDict(extra="allow") type: str = Field(min_length=1, pattern=r"\S") @@ -347,14 +351,14 @@ def _serialize_function_call(self, piece: MessagePiece) -> "ResponseFunctionTool "arguments": call.function.arguments, } - def _serialize_tool_call(self, piece: MessagePiece) -> dict[str, Any]: + def _serialize_tool_call(self, piece: MessagePiece) -> dict[str, Any] | None: stored = _ResponseToolCallContent.model_validate_json(piece.converted_value).model_dump(exclude_unset=True) if stored.get("type") == "web_search_call": - return { - "type": stored["type"], - "call_id": stored.get("call_id"), - "query": stored.get("query"), - } + # The input item takes exactly the output item's fields, and all of them are required. Older rows + # kept only type and id, which the API rejects, so those can't be replayed. + if any(stored.get(key) is None for key in _WEB_SEARCH_CALL_INPUT_FIELDS): + return None + return {key: stored[key] for key in _WEB_SEARCH_CALL_INPUT_FIELDS} filtered = {"type": stored["type"]} filtered.update({key: stored[key] for key in ("call_id", "query", "name", "arguments") if key in stored}) return filtered @@ -377,7 +381,14 @@ async def _serialize_piece_async(self, *, piece: MessagePiece, message_index: in if data_type == "function_call": return _SerializedPiece(item=dict(self._serialize_function_call(piece)), placement="top_level") if data_type == "tool_call": - return _SerializedPiece(item=self._serialize_tool_call(piece), placement="top_level") + tool_call = self._serialize_tool_call(piece) + if tool_call is None: + logger.warning( + "Skipping web_search_call in message index %d: it lacks the id, status or action the API needs.", + message_index, + ) + return None + return _SerializedPiece(item=tool_call, placement="top_level") if data_type == "function_call_output": return _SerializedPiece(item=dict(self._serialize_function_call_output(piece)), placement="top_level") raise ValueError(f"Unsupported data type '{data_type}' in message index {message_index}") @@ -610,7 +621,14 @@ async def _construct_message_from_response_async(self, response: Response, reque # reasoning commonly precedes the actual message in provider output, so # retain it for memory/debugging after the actionable response pieces. # This must stay ahead of the metadata writes below, which target the first piece. - extracted_response_pieces.sort(key=lambda piece: piece.converted_value_data_type == "reasoning") + # Hosted tool records (e.g. web_search_call) also precede the answer; keep them after it too. Replay + # order is unaffected because tool records are always emitted before the message item. + extracted_response_pieces.sort( + key=lambda piece: ( + piece.converted_value_data_type == "reasoning", + piece.converted_value_data_type == "tool_call", + ) + ) # Capture token usage and the stop reason in the first piece's metadata. This also runs on # the truncated path: usage is populated on token-limit responses and is most valuable @@ -834,20 +852,9 @@ def _parse_response_output_section( piece_type = "function_call" elif section_type == MessagePieceType.WEB_SEARCH_CALL: - # Forward web_search_call with only API-expected fields - # Note: web search may have different field structure than function calls - web_search_data = { - "type": "web_search_call", - } - # Add optional fields if they exist - if hasattr(section, "call_id") and section.call_id: - web_search_data["call_id"] = section.call_id - if hasattr(section, "query") and section.query: - web_search_data["query"] = section.query - if hasattr(section, "id") and section.id: - web_search_data["id"] = section.id - - piece_value = json.dumps(web_search_data, separators=(",", ":")) + # Keep the whole item (id, status and the search action). The API needs all of it to accept the + # call back as input on the next turn. + piece_value = json.dumps(section.model_dump(mode="json", exclude_none=True), separators=(",", ":")) piece_type = "tool_call" elif section_type == "custom_tool_call": diff --git a/tests/unit/backend/test_conversation_editor.py b/tests/unit/backend/test_conversation_editor.py index 68732c6a24..99927d54d6 100644 --- a/tests/unit/backend/test_conversation_editor.py +++ b/tests/unit/backend/test_conversation_editor.py @@ -476,7 +476,10 @@ async def test_provider_tool_extensions_round_trip_without_execution_async( ) -> None: request = draft() request.target_registry_name = "responses" - payload = '{"type":"web_search_call","call_id":"web-1","query":"query","extension":{"keep":true}}' + payload = ( + '{"type":"web_search_call","id":"ws-1","status":"completed",' + '"action":{"type":"search","query":"query"},"extension":{"keep":true}}' + ) request.messages = [ ConversationMessageRequest( role="simulated_assistant", diff --git a/tests/unit/prompt_target/target/test_openai_response_adapters.py b/tests/unit/prompt_target/target/test_openai_response_adapters.py index 2b1f081157..ab59e060da 100644 --- a/tests/unit/prompt_target/target/test_openai_response_adapters.py +++ b/tests/unit/prompt_target/target/test_openai_response_adapters.py @@ -48,12 +48,20 @@ def test_tool_preflight_and_serialization_reject_invalid_provider_fields(payload @pytest.mark.usefixtures("patch_central_database") def test_tool_preflight_accepts_empty_history_and_preserves_provider_extensions() -> None: target = object.__new__(OpenAIResponseTarget) - payload = '{"type":"web_search_call","call_id":"web-1","query":"query","extension":{"key":"value"}}' + payload = ( + '{"type":"web_search_call","id":"ws-1","status":"completed",' + '"action":{"type":"search","query":"query"},"extension":{"key":"value"}}' + ) piece = MessagePiece(role="simulated_assistant", original_value_data_type="tool_call", original_value=payload) target.validate_tool_history([]) target.validate_tool_history([Message(message_pieces=[piece])]) assert piece.converted_value == payload - assert target._serialize_tool_call(piece) == {"type": "web_search_call", "call_id": "web-1", "query": "query"} + assert target._serialize_tool_call(piece) == { + "id": "ws-1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "query"}, + } @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/prompt_target/target/test_openai_response_target.py b/tests/unit/prompt_target/target/test_openai_response_target.py index 03c68aa5ac..0d6f33a9c1 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -7,12 +7,18 @@ from collections.abc import MutableSequence from tempfile import NamedTemporaryFile from types import SimpleNamespace -from typing import Any +from typing import Any, get_type_hints from unittest.mock import AsyncMock, MagicMock, patch import pytest from openai import BadRequestError, RateLimitError -from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal, ResponseOutputText +from openai.types.responses import ( + Response, + ResponseFunctionWebSearchParam, + ResponseOutputMessage, + ResponseOutputRefusal, + ResponseOutputText, +) from unit.mocks import ( get_audio_message_piece, get_image_message_piece, @@ -53,8 +59,6 @@ def create_mock_response(response_dict: dict = None) -> MagicMock: Returns: A mock object that simulates the OpenAI SDK response with Pydantic-style attribute access. """ - from openai.types.responses import Response - if response_dict is None: response_dict = openai_response_json_dict() @@ -970,7 +974,15 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac ), MessagePiece( role="assistant", - original_value=json.dumps({"type": "web_search_call", "call_id": "web-1", "id": "drop-id"}), + original_value=json.dumps( + { + "type": "web_search_call", + "id": "ws-1", + "status": "completed", + "action": {"type": "search", "query": "query"}, + "extra": "drop", + } + ), original_value_data_type="tool_call", ), MessagePiece( @@ -1036,7 +1048,12 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac "name": "lookup", "arguments": '{"value":1}', }, - {"type": "web_search_call", "call_id": "web-1", "query": None}, + { + "id": "ws-1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "query"}, + }, { "type": "provider_tool_call", "call_id": "tool-1", @@ -1055,6 +1072,86 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac ] +def _web_search_response() -> Response: + return Response.model_validate( + { + "id": "resp-1", + "object": "response", + "created_at": 1, + "model": "gpt-4.1", + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [{"type": "web_search_preview"}], + "status": "completed", + "output": [ + { + "type": "web_search_call", + "id": "ws-1", + "status": "completed", + "action": {"type": "search", "query": "pyrit release"}, + }, + { + "type": "message", + "id": "msg-1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "PyRIT 0.13 is out.", "annotations": []}], + }, + ], + } + ) + + +async def test_web_search_answer_is_the_response_value_async( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece) + + assert reply.get_value() == "PyRIT 0.13 is out." + assert [piece.converted_value_data_type for piece in reply.message_pieces] == ["text", "tool_call"] + + +async def test_web_search_call_replays_as_a_valid_input_item_async( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece) + follow_up = Message.from_prompt(prompt="Anything else?", role="user") + + items = await target._build_input_for_multi_modal_async([dummy_text_message_piece.to_message(), reply, follow_up]) + + web_search = { + "id": "ws-1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "pyrit release"}, + } + assert items[1] == web_search + assert set(web_search) == set(get_type_hints(ResponseFunctionWebSearchParam)) + assert items[2] == {"role": "assistant", "content": [{"type": "output_text", "text": "PyRIT 0.13 is out."}]} + + +async def test_web_search_call_stored_without_its_action_is_skipped_on_replay_async( + target: OpenAIResponseTarget, caplog: pytest.LogCaptureFixture +): + reply = Message( + message_pieces=[ + MessagePiece(role="assistant", original_value="answer", conversation_id="c-1"), + MessagePiece( + role="assistant", + original_value='{"type":"web_search_call","id":"ws-1"}', + original_value_data_type="tool_call", + conversation_id="c-1", + ), + ] + ) + + with caplog.at_level(logging.WARNING): + items = await target._build_input_for_multi_modal_async([reply]) + + assert items == [{"role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}] + assert "Skipping web_search_call" in caplog.text + + @pytest.mark.parametrize("data_type", ["function_call", "tool_call", "function_call_output"]) async def test_build_input_for_multi_modal_async_validates_malformed_artifact_async( target: OpenAIResponseTarget, data_type: PromptDataType diff --git a/tests/unit/prompt_target/target/test_tool_call_history.py b/tests/unit/prompt_target/target/test_tool_call_history.py index c1e23381d3..1dc03a1e4b 100644 --- a/tests/unit/prompt_target/target/test_tool_call_history.py +++ b/tests/unit/prompt_target/target/test_tool_call_history.py @@ -345,7 +345,10 @@ def test_responses_preflight_and_serializer_reject_invalid_tool_fields(payload: def test_responses_preflight_retains_provider_extensions_and_uses_converted_values() -> None: target = _target(OpenAIResponseTarget) assert isinstance(target, OpenAIResponseTarget) - payload = '{"type":"web_search_call","call_id":"web-1","query":"edited","extension":{"key":"value"}}' + payload = ( + '{"type":"web_search_call","id":"ws-1","status":"completed",' + '"action":{"type":"search","query":"edited"},"extension":{"key":"value"}}' + ) piece = MessagePiece( role="simulated_assistant", original_value_data_type="tool_call", @@ -354,7 +357,12 @@ def test_responses_preflight_retains_provider_extensions_and_uses_converted_valu ) target.validate_tool_history([piece.to_message()]) assert piece.converted_value == payload - assert target._serialize_tool_call(piece) == {"type": "web_search_call", "call_id": "web-1", "query": "edited"} + assert target._serialize_tool_call(piece) == { + "id": "ws-1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "edited"}, + } @pytest.mark.parametrize("target_type", [OpenAIChatTarget, LiteLLMChatTarget]) From 5469c499b8316b773ee3d581560ea89353835525 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sun, 11 Oct 2026 10:43:10 +0530 Subject: [PATCH 2/3] tighten tests: async names, typed args, public paths --- .../target/test_openai_response_target.py | 71 ++++++++++++++++++- 1 file changed, 68 insertions(+), 3 deletions(-) diff --git a/tests/unit/prompt_target/target/test_openai_response_target.py b/tests/unit/prompt_target/target/test_openai_response_target.py index 0d6f33a9c1..68378e483b 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -4,6 +4,7 @@ import json import logging import os +import uuid from collections.abc import MutableSequence from tempfile import NamedTemporaryFile from types import SimpleNamespace @@ -43,6 +44,7 @@ ToolExecutionMetadata, flatten_to_message_pieces, ) +from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import OpenAIResponseTarget, PromptTarget from pyrit.prompt_target.openai.openai_response_target import _ToolDispatchResult, token_usage_from_responses from pyrit.score import SelfAskRefusalScorer, TrueFalseInverterScorer @@ -1072,7 +1074,39 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac ] -def _web_search_response() -> Response: +_WEB_SEARCH_ACTIONS = [ + {"type": "search", "query": "pyrit release"}, + {"type": "open_page", "url": "https://example.com/pyrit"}, + {"type": "find_in_page", "pattern": "release", "url": "https://example.com/pyrit"}, +] + + +def _answer_response(*, text: str, extra_output: list[dict[str, Any]] | None = None) -> Response: + return Response.model_validate( + { + "id": "resp-2", + "object": "response", + "created_at": 1, + "model": "gpt-4.1", + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [{"type": "web_search_preview"}], + "status": "completed", + "output": [ + *(extra_output or []), + { + "type": "message", + "id": "msg-2", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + }, + ], + } + ) + + +def _web_search_response(*, action: dict[str, Any] | None = None) -> Response: return Response.model_validate( { "id": "resp-1", @@ -1088,7 +1122,7 @@ def _web_search_response() -> Response: "type": "web_search_call", "id": "ws-1", "status": "completed", - "action": {"type": "search", "query": "pyrit release"}, + "action": action or _WEB_SEARCH_ACTIONS[0], }, { "type": "message", @@ -1104,7 +1138,7 @@ def _web_search_response() -> Response: async def test_web_search_answer_is_the_response_value_async( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece -): +) -> None: reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece) assert reply.get_value() == "PyRIT 0.13 is out." @@ -1130,6 +1164,37 @@ async def test_web_search_call_replays_as_a_valid_input_item_async( assert items[2] == {"role": "assistant", "content": [{"type": "output_text", "text": "PyRIT 0.13 is out."}]} +@pytest.mark.parametrize("action", _WEB_SEARCH_ACTIONS) +async def test_web_search_conversation_continues_through_send_prompt_async( + *, target: OpenAIResponseTarget, sqlite_instance: MemoryInterface, action: dict[str, Any] +) -> None: + create = AsyncMock(side_effect=[_web_search_response(action=action), _answer_response(text="Nothing else is new.")]) + target._async_client.responses.create = create # type: ignore[method-assign] + normalizer = PromptNormalizer() + conversation_id = str(uuid.uuid4()) + + first = await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="What's new in PyRIT?", role="user"), + target=target, + conversation_id=conversation_id, + ) + second = await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="Anything else?", role="user"), + target=target, + conversation_id=conversation_id, + ) + + web_search = {"id": "ws-1", "type": "web_search_call", "status": "completed", "action": action} + assert first.get_value() == "PyRIT 0.13 is out." + assert second.get_value() == "Nothing else is new." + replayed = create.await_args_list[1].kwargs["input"] + assert replayed[1] == web_search + assert replayed[2]["role"] == "assistant" + stored = await sqlite_instance.get_message_pieces_async(conversation_id=conversation_id) + [tool_call] = [piece for piece in stored if piece.converted_value_data_type == "tool_call"] + assert json.loads(tool_call.converted_value) == web_search + + async def test_web_search_call_stored_without_its_action_is_skipped_on_replay_async( target: OpenAIResponseTarget, caplog: pytest.LogCaptureFixture ): From 924d65b3d42a177efb433c97dfbde5be63f3dd05 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sun, 11 Oct 2026 10:44:03 +0530 Subject: [PATCH 3/3] type the web search test returns --- .../unit/prompt_target/target/test_openai_response_target.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/prompt_target/target/test_openai_response_target.py b/tests/unit/prompt_target/target/test_openai_response_target.py index 68378e483b..31db14a96a 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -1147,7 +1147,7 @@ async def test_web_search_answer_is_the_response_value_async( async def test_web_search_call_replays_as_a_valid_input_item_async( target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece -): +) -> None: reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece) follow_up = Message.from_prompt(prompt="Anything else?", role="user") @@ -1197,7 +1197,7 @@ async def test_web_search_conversation_continues_through_send_prompt_async( async def test_web_search_call_stored_without_its_action_is_skipped_on_replay_async( target: OpenAIResponseTarget, caplog: pytest.LogCaptureFixture -): +) -> None: reply = Message( message_pieces=[ MessagePiece(role="assistant", original_value="answer", conversation_id="c-1"),