Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 29 additions & 22 deletions pyrit/prompt_target/openai/openai_response_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -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}")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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":
Expand Down
5 changes: 4 additions & 1 deletion tests/unit/backend/test_conversation_editor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 10 additions & 2 deletions tests/unit/prompt_target/target/test_openai_response_adapters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
174 changes: 168 additions & 6 deletions tests/unit/prompt_target/target/test_openai_response_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,22 @@
import json
import logging
import os
import uuid
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,
Expand All @@ -37,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
Expand All @@ -53,8 +61,6 @@ def create_mock_response(response_dict: dict | None = 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()

Expand Down Expand Up @@ -970,7 +976,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(
Expand Down Expand Up @@ -1036,7 +1050,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",
Expand All @@ -1055,6 +1074,149 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac
]


_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",
"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": action or _WEB_SEARCH_ACTIONS[0],
},
{
"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
) -> 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."
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
) -> 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")

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."}]}


@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
) -> None:
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
Expand Down
12 changes: 10 additions & 2 deletions tests/unit/prompt_target/target/test_tool_call_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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])
Expand Down
Loading