diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index ce428d9fca..1f07548d1e 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -117,6 +117,8 @@ TargetCapabilities, TokenUsage, get_common_json_schema, + read_usage_int, + read_usage_value, register_common_json_schema, unregister_common_json_schema, ) @@ -223,6 +225,8 @@ "TokenUsage", "ToolCall", "UnvalidatedScore", + "read_usage_int", + "read_usage_value", "validate_registry_name", "RetryEvent", ] diff --git a/pyrit/models/messages/message_piece.py b/pyrit/models/messages/message_piece.py index 9ab340011f..799171a3cc 100644 --- a/pyrit/models/messages/message_piece.py +++ b/pyrit/models/messages/message_piece.py @@ -44,6 +44,7 @@ class MessagePiece(BaseModel): """ STRUCTURED_REFUSAL_METADATA_KEY: ClassVar[str] = "structured_refusal" + TRUNCATED_METADATA_KEY: ClassVar[str] = "truncated" model_config = ConfigDict( arbitrary_types_allowed=True, @@ -193,6 +194,20 @@ def structured_refusal(self) -> str | None: refusal = self.prompt_metadata.get(self.STRUCTURED_REFUSAL_METADATA_KEY) return refusal if isinstance(refusal, str) and refusal else None + def mark_as_truncated(self) -> None: + """ + Record that the target cut this response off at its output-token limit. + + A truncated piece may still carry a partial answer with ``response_error == "none"``, so + without this flag a consumer cannot tell a complete answer from a clipped one. + """ + self.prompt_metadata[self.TRUNCATED_METADATA_KEY] = True + + @property + def is_truncated(self) -> bool: + """Whether the target cut this response off at its output-token limit.""" + return bool(self.prompt_metadata.get(self.TRUNCATED_METADATA_KEY)) + # ------------------------------------------------------------------ # # Adversarial placeholder support # ------------------------------------------------------------------ # diff --git a/pyrit/models/target/__init__.py b/pyrit/models/target/__init__.py index d43320dda2..7e3455bfaf 100644 --- a/pyrit/models/target/__init__.py +++ b/pyrit/models/target/__init__.py @@ -28,7 +28,7 @@ unregister_common_json_schema, ) from pyrit.models.target.target_capabilities import CapabilityName, TargetCapabilities -from pyrit.models.target.token_usage import TokenUsage +from pyrit.models.target.token_usage import TokenUsage, read_usage_int, read_usage_value __all__ = [ "COMMON_JSON_SCHEMAS", @@ -40,6 +40,8 @@ "TargetCapabilities", "TokenUsage", "get_common_json_schema", + "read_usage_int", + "read_usage_value", "register_common_json_schema", "unregister_common_json_schema", ] diff --git a/pyrit/models/target/token_usage.py b/pyrit/models/target/token_usage.py index 32bc39ce1e..08374788e0 100644 --- a/pyrit/models/target/token_usage.py +++ b/pyrit/models/target/token_usage.py @@ -3,6 +3,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field from typing import Any @@ -15,6 +16,47 @@ _CORE_SUFFIXES = frozenset({"input_tokens", "output_tokens", "total_tokens", "reasoning_tokens", "cached_tokens"}) +def read_usage_value(*, source: Any, name: str) -> Any: + """ + Read ``name`` from a provider usage payload, which may be a mapping or an attribute object. + + Provider SDKs surface usage either as a typed object (OpenAI/LiteLLM ``Usage``) or as a + ``model_dump``'d mapping, so both access styles are supported. Use this to reach nested + breakdown objects (for example ``prompt_tokens_details``) before reading counts out of them + with ``read_usage_int``. + + Args: + source (Any): The usage object or nested details object (may be None). + name (str): The field name to read. + + Returns: + Any: The field value, or None when absent. + """ + if isinstance(source, Mapping): + return source.get(name) + return getattr(source, name, None) + + +def read_usage_int(*, source: Any, name: str) -> int | None: + """ + Read ``name`` from a provider usage payload and return it only when it is an integer count. + + Which field name holds which count is wire-format specific and therefore the caller's + concern; this helper only owns the read and the int guard, so a partial usage payload + contributes just the counts the provider actually reports. Booleans are rejected even though + ``bool`` is a subclass of ``int``. + + Args: + source (Any): The usage object or nested details object (may be None). + name (str): The field name to read. + + Returns: + int | None: The integer value, or None when absent or not an integer. + """ + value = read_usage_value(source=source, name=name) + return value if isinstance(value, int) and not isinstance(value, bool) else None + + @dataclass(frozen=True) class TokenUsage: """ @@ -31,7 +73,9 @@ class TokenUsage: This is a pure value object: it holds counts and (de)serializes them to metadata. Turning a provider ``usage`` payload into a ``TokenUsage`` is the responsibility of the target/parser that knows which wire format it received (for example, the Chat Completions parser in - ``pyrit.prompt_target.common.chat_completions_response_parser``). + ``pyrit.prompt_target.common.chat_completions_response_parser``). Only the format-agnostic part + of that read -- mapping-or-attribute access and the integer guard -- is shared here, via + ``read_usage_value`` and ``read_usage_int``. Neither cost nor the responding model name is modeled here: cost is a currency amount (tracked separately under ``token_usage_cost``) and the model identity is already recorded on the diff --git a/pyrit/prompt_target/common/chat_completions_response_parser.py b/pyrit/prompt_target/common/chat_completions_response_parser.py index 8709fa831d..16c7cb5256 100644 --- a/pyrit/prompt_target/common/chat_completions_response_parser.py +++ b/pyrit/prompt_target/common/chat_completions_response_parser.py @@ -13,9 +13,10 @@ import base64 import json import logging -from collections.abc import Mapping from typing import Any +from openai.types.chat import ChatCompletion + from pyrit.exceptions import ( EmptyResponseException, PyritException, @@ -27,6 +28,8 @@ MessagePiece, TokenUsage, construct_response_from_request, + read_usage_int, + read_usage_value, ) logger = logging.getLogger(__name__) @@ -37,6 +40,21 @@ DEFAULT_VALID_FINISH_REASONS: frozenset[str] = frozenset({"stop", "length", "content_filter", "tool_calls"}) +def get_finish_reason(*, response: ChatCompletion) -> str | None: + """ + Extract the first choice's ``finish_reason`` from a Chat Completions response. + + Args: + response (ChatCompletion): The Chat Completions response object. + + Returns: + str | None: The first choice's ``finish_reason``, or None when there are no choices. + """ + if not response.choices: + return None + return response.choices[0].finish_reason + + def detect_response_content(message: Any) -> tuple[bool, bool, bool]: """ Detect which content types are present in a Chat Completions ``message``. @@ -284,44 +302,6 @@ def capture_token_usage(*, pieces: list[MessagePiece], response: Any) -> None: pieces[0].prompt_metadata.update(token_usage.to_metadata()) -def _read(source: Any, name: str) -> Any: - """ - Read ``name`` from ``source``, which may be a mapping or an attribute object. - - Args: - source (Any): The usage object (may be None). - name (str): The field name to read. - - Returns: - Any: The field value, or None when absent. - """ - if isinstance(source, Mapping): - return source.get(name) - return getattr(source, name, None) - - -def _usage_field(source: Any, *names: str) -> int | None: - """ - Return the first int-valued field among ``names`` on ``source``, else None. - - ``source`` may be either a mapping (for example, a ``model_dump``'d usage payload) or an - attribute object (the OpenAI/LiteLLM SDK ``Usage`` type), so both access styles are supported. - Booleans are rejected even though ``bool`` is a subclass of ``int``. - - Args: - source (Any): The usage object or nested details object (may be None). - names (str): Candidate field names, tried in order. - - Returns: - int | None: The first integer value found, or None. - """ - for name in names: - value = _read(source, name) - if isinstance(value, int) and not isinstance(value, bool): - return value - return None - - def token_usage_from_chat_completion(usage: Any) -> TokenUsage: """ Build a ``TokenUsage`` from a Chat Completions ``usage`` payload (OpenAI or LiteLLM). @@ -335,7 +315,8 @@ def token_usage_from_chat_completion(usage: Any) -> TokenUsage: This parser is specific to the Chat Completions wire format. The Responses API reports usage under different names (``input_tokens`` / ``output_tokens``); a target that speaks that format - should parse it in its own module rather than overloading this function. + should parse it in its own module rather than overloading this function, reusing the shared + ``read_usage_value`` / ``read_usage_int`` reads. Args: usage (Any): The Chat Completions usage object (attribute object or mapping). @@ -343,26 +324,34 @@ def token_usage_from_chat_completion(usage: Any) -> TokenUsage: Returns: TokenUsage: The parsed token usage. """ - input_tokens = _usage_field(usage, "prompt_tokens") - output_tokens = _usage_field(usage, "completion_tokens") - total_tokens = _usage_field(usage, "total_tokens") + input_tokens = read_usage_int(source=usage, name="prompt_tokens") + output_tokens = read_usage_int(source=usage, name="completion_tokens") + total_tokens = read_usage_int(source=usage, name="total_tokens") if total_tokens is None and input_tokens is not None and output_tokens is not None: total_tokens = input_tokens + output_tokens - prompt_details = _read(usage, "prompt_tokens_details") - completion_details = _read(usage, "completion_tokens_details") + prompt_details = read_usage_value(source=usage, name="prompt_tokens_details") + completion_details = read_usage_value(source=usage, name="completion_tokens_details") - cached_tokens = _usage_field(prompt_details, "cached_tokens") + cached_tokens = read_usage_int(source=prompt_details, name="cached_tokens") if cached_tokens is None: - cached_tokens = _usage_field(usage, "cache_read_input_tokens") - reasoning_tokens = _usage_field(completion_details, "reasoning_tokens") + cached_tokens = read_usage_int(source=usage, name="cache_read_input_tokens") + reasoning_tokens = read_usage_int(source=completion_details, name="reasoning_tokens") extra: dict[str, int] = {} - _add_extra(extra, "input_audio_tokens", _usage_field(prompt_details, "audio_tokens")) - _add_extra(extra, "cache_write_tokens", _usage_field(usage, "cache_creation_input_tokens")) - _add_extra(extra, "output_audio_tokens", _usage_field(completion_details, "audio_tokens")) - _add_extra(extra, "accepted_prediction_tokens", _usage_field(completion_details, "accepted_prediction_tokens")) - _add_extra(extra, "rejected_prediction_tokens", _usage_field(completion_details, "rejected_prediction_tokens")) + _add_extra(extra, "input_audio_tokens", read_usage_int(source=prompt_details, name="audio_tokens")) + _add_extra(extra, "cache_write_tokens", read_usage_int(source=usage, name="cache_creation_input_tokens")) + _add_extra(extra, "output_audio_tokens", read_usage_int(source=completion_details, name="audio_tokens")) + _add_extra( + extra, + "accepted_prediction_tokens", + read_usage_int(source=completion_details, name="accepted_prediction_tokens"), + ) + _add_extra( + extra, + "rejected_prediction_tokens", + read_usage_int(source=completion_details, name="rejected_prediction_tokens"), + ) return TokenUsage( input_tokens=input_tokens, diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index 6b6a74a813..92fdffff70 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/pyrit/prompt_target/common/utils.py @@ -2,10 +2,14 @@ # Licensed under the MIT license. import asyncio +import logging from collections.abc import Callable from typing import Any from pyrit.exceptions import PyritException +from pyrit.models import Message, MessagePiece, construct_response_from_request + +logger = logging.getLogger(__name__) def validate_temperature(temperature: float | None) -> None: @@ -57,3 +61,47 @@ async def set_max_rpm_async(*args: Any, **kwargs: Any) -> Any: return await func(*args, **kwargs) return set_max_rpm_async + + +def build_empty_truncated_response(*, request: MessagePiece) -> Message: + """ + Build a graceful empty response for a token-limit-truncated model response. + + A response truncated at the token limit (Chat Completions ``finish_reason == "length"`` or the + Responses API ``status == "incomplete"`` with ``reason == "max_output_tokens"``) may legitimately + contain no visible content. Callers gate this on their own truncation check (for example a + target's ``_is_truncated_response``); returning an empty ``error="empty"`` text response lets the + run continue instead of raising. + + Args: + request (MessagePiece): The originating request piece. + + Returns: + Message: An empty text response marked with ``error="empty"``. + """ + return construct_response_from_request( + request=request, + response_text_pieces=[""], + response_type="text", + error="empty", + ) + + +def warn_truncated_response(*, signal: str, limit_parameter: str) -> None: + """ + Log the shared warning for a response cut off at the output-token limit. + + Every API shape signals truncation differently but the advice is identical, so the wording + lives here to keep targets from drifting apart. + + Args: + signal (str): How the API reported the truncation, quoted into the message (for example + ``"finish_reason='length'"``). + limit_parameter (str): The request parameter to raise (for example ``"max_output_tokens"``). + """ + logger.warning( + f"The response was truncated because it reached the token limit ({signal}). Reasoning models " + f"consume tokens on hidden reasoning in addition to the visible answer, so a low " + f"{limit_parameter} can truncate or empty the response. Increase {limit_parameter} if you " + "expected complete content." + ) diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 2d21e5a1b5..3d9f00a2e8 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/pyrit/prompt_target/openai/openai_chat_target.py @@ -5,6 +5,8 @@ from collections.abc import MutableSequence from typing import Any +from openai.types.chat import ChatCompletion + from pyrit.common import forward_init_parameters from pyrit.exceptions import ( EmptyResponseException, @@ -28,6 +30,7 @@ capture_token_usage, detect_response_content, extract_partial_content, + get_finish_reason, is_content_filter_response, save_audio_response_async, validate_chat_completion_response, @@ -35,9 +38,11 @@ from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import ( + build_empty_truncated_response, limit_requests_per_minute, validate_temperature, validate_top_p, + warn_truncated_response, ) from pyrit.prompt_target.openai.openai_chat_audio_config import OpenAIChatAudioConfig from pyrit.prompt_target.openai.openai_target import OpenAITarget @@ -271,7 +276,7 @@ def _extract_partial_content(self, response: Any) -> str | None: """ return extract_partial_content(response) - def _validate_response(self, response: Any, request: MessagePiece) -> Message | None: + def _validate_response(self, response: ChatCompletion, request: MessagePiece) -> None: """ Validate a Chat Completions API response for errors. @@ -280,19 +285,48 @@ def _validate_response(self, response: Any, request: MessagePiece) -> Message | - Invalid finish_reason - At least one valid response type (text content, audio, or tool_calls) + A ``finish_reason == "length"`` (token-limit truncation) response is treated as valid, with a + warning, so that ``_construct_message_from_response_async`` can preserve any partial content + or fall back to a graceful empty response. Genuinely empty responses (no truncation) are + raised so the retry logic can attempt to get a complete response. Content filter responses + are handled separately by ``_check_content_filter``. + Args: response: The ChatCompletion response from OpenAI SDK. request: The original request MessagePiece. - Returns: - None if valid, does not return Message for content filter (handled by _check_content_filter). - Raises: PyritException: For unexpected response structures or finish reasons. - EmptyResponseException: When the API returns an empty response. + EmptyResponseException: When the API returns an empty response that was not caused by + token-limit truncation. """ + # Token-limit truncation is handled before the shared validator, which would otherwise raise + # EmptyResponseException on a validly truncated but empty response. Reasoning models can spend + # the whole budget on hidden reasoning before emitting a visible answer, and a low limit may be + # deliberate, so warn instead of raising and let construction preserve any partial content or + # fall back to a graceful empty response. + if self._is_truncated_response(response): + warn_truncated_response(signal="finish_reason='length'", limit_parameter="max_completion_tokens") + return + + # Genuinely empty responses (no truncation) raise so the retry logic can attempt to get a + # complete response. validate_chat_completion_response(response=response) - return None + + def _is_truncated_response(self, response: ChatCompletion) -> bool: + """ + Return True if the response was cut off by the token limit. + + The Chat Completions API signals token-limit truncation via ``finish_reason == "length"`` + on the first choice. + + Args: + response: A ChatCompletion response from the OpenAI SDK. + + Returns: + bool: True if the response was truncated at the token limit, False otherwise. + """ + return get_finish_reason(response=response) == "length" def _detect_response_content(self, message: Any) -> tuple[bool, bool, bool]: """ @@ -337,7 +371,7 @@ def _should_skip_sending_audio( prefer_transcript_for_history=prefer_transcript_for_history, ) - async def _construct_message_from_response_async(self, response: Any, request: MessagePiece) -> Message: + async def _construct_message_from_response_async(self, response: ChatCompletion, request: MessagePiece) -> Message: """ Construct a Message from a ChatCompletion response. @@ -354,16 +388,30 @@ async def _construct_message_from_response_async(self, response: Any, request: M Message: Constructed message with one or more MessagePiece entries. Raises: - EmptyResponseException: If the response contains no content, audio, or tool calls. + EmptyResponseException: If a non-truncated response contains no content, audio, or tool + calls. A truncated (``finish_reason == "length"``) response with no content instead + yields a graceful empty piece so the run continues. Truncated responses are flagged + via ``MessagePiece.mark_as_truncated`` on the first piece. """ audio_format = self._audio_response_config.audio_format if self._audio_response_config else "wav" + truncated = self._is_truncated_response(response) pieces = await build_response_pieces_async(response=response, request=request, audio_format=audio_format) if not pieces: + # A truncated (finish_reason == "length") response may legitimately produce no content; + # return a graceful empty piece so the run continues. Validation already raised for + # genuinely empty (non-truncated) responses. + if truncated: + empty_message = build_empty_truncated_response(request=request) + capture_token_usage(pieces=empty_message.message_pieces, response=response) + empty_message.message_pieces[0].mark_as_truncated() + return empty_message raise EmptyResponseException(message="Failed to extract any response content.") # Capture token usage from the API response and store in the first piece's metadata capture_token_usage(pieces=pieces, response=response) + if truncated: + pieces[0].mark_as_truncated() return Message(message_pieces=pieces) diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 2b14e465fd..62e384a739 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -12,7 +12,7 @@ cast, ) -from openai.types.responses import ResponseOutputRefusal, ResponseOutputText +from openai.types.responses import Response, ResponseOutputRefusal, ResponseOutputText from openai.types.shared import ReasoningEffort from pyrit.common import forward_init_parameters @@ -29,13 +29,18 @@ MessagePiece, PromptDataType, PromptResponseError, + TokenUsage, + read_usage_int, + read_usage_value, ) from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.utils import ( + build_empty_truncated_response, limit_requests_per_minute, validate_temperature, validate_top_p, + warn_truncated_response, ) from pyrit.prompt_target.openai.openai_error_handling import _is_content_filter_error from pyrit.prompt_target.openai.openai_target import OpenAITarget @@ -67,6 +72,47 @@ class MessagePieceType(str, Enum): MCP_APPROVAL_REQUEST = "mcp_approval_request" +def token_usage_from_responses(usage: Any) -> TokenUsage: + """ + Build a ``TokenUsage`` from a Responses API ``usage`` payload. + + The Responses API reports usage under different names than Chat Completions -- top-level + ``input_tokens`` / ``output_tokens`` / ``total_tokens`` with ``input_tokens_details`` and + ``output_tokens_details`` breakdowns -- so the field names are resolved here rather than by + ``token_usage_from_chat_completion``. Both parsers share the format-agnostic reads + (``read_usage_value`` / ``read_usage_int``), so a partial usage payload contributes only the + counts the provider actually reports. ``total_tokens`` is derived when the provider omits it. + + Args: + usage (Any): The Responses API usage object. + + Returns: + TokenUsage: The parsed token usage. + """ + input_details = read_usage_value(source=usage, name="input_tokens_details") + output_details = read_usage_value(source=usage, name="output_tokens_details") + + input_tokens = read_usage_int(source=usage, name="input_tokens") + output_tokens = read_usage_int(source=usage, name="output_tokens") + total_tokens = read_usage_int(source=usage, name="total_tokens") + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + + extra: dict[str, int] = {} + cache_write_tokens = read_usage_int(source=input_details, name="cache_write_tokens") + if cache_write_tokens is not None: + extra["cache_write_tokens"] = cache_write_tokens + + return TokenUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + reasoning_tokens=read_usage_int(source=output_details, name="reasoning_tokens"), + cached_tokens=read_usage_int(source=input_details, name="cached_tokens"), + extra=extra, + ) + + class OpenAIResponseTarget(OpenAITarget): """ Enables communication with endpoints that support the OpenAI Response API. @@ -516,31 +562,47 @@ def _extract_partial_content(self, response: Any) -> str | None: except (AttributeError, IndexError, TypeError): return None - def _validate_response(self, response: Any, request: MessagePiece) -> Message | None: + def _validate_response(self, response: Response, request: MessagePiece) -> None: """ Validate a Response API response for errors. Checks for: - Error responses (excluding content filtering which is checked separately) + - Truncation at the token limit (``max_output_tokens``), which is warned about, not raised - Invalid status - Empty output + Truncation is treated as valid, with a warning, so that + ``_construct_message_from_response_async`` can preserve any completed output (reasoning, + partial text) or fall back to a graceful empty response. Genuinely empty responses (no + truncation) are raised so the retry logic can attempt to get a complete response. Content + filter responses are handled separately by ``_check_content_filter``. + Args: response: The Response object from the OpenAI SDK. request: The original request MessagePiece. - Returns: - None if valid, does not return Message for content filter (handled by _check_content_filter). - Raises: PyritException: For unexpected response structures or errors. - EmptyResponseException: When the API returns no valid output. + EmptyResponseException: When the API returns no valid output (and was not truncated). """ # Check for error response - error is a ResponseError object or None # (content_filter is handled by _check_content_filter) if response.error is not None and response.error.code != "content_filter": raise PyritException(message=f"Response error: {response.error.code} - {response.error.message}") + # Truncation: the model hit max_output_tokens. Mirroring OpenAIChatTarget's handling of + # finish_reason == "length", warn instead of raising so the run continues -- reasoning models + # can spend the whole budget on hidden reasoning before emitting a visible answer, and a low + # limit may be a deliberate configuration. Construction preserves any completed output + # (reasoning, partial text) and falls back to a graceful empty response. + if self._is_truncated_response(response): + warn_truncated_response( + signal="status='incomplete', reason='max_output_tokens'", + limit_parameter="max_output_tokens", + ) + return + # Check status - should be "completed" for successful responses if response.status != "completed": raise PyritException(message=f"Unexpected status: {response.status}") @@ -550,36 +612,92 @@ def _validate_response(self, response: Any, request: MessagePiece) -> Message | logger.error("The response returned no valid output.") raise EmptyResponseException(message="The response returned an empty response.") - return None + def _is_truncated_response(self, response: Response) -> bool: + """ + Return True if the response was cut off by the ``max_output_tokens`` limit. + + The Responses API signals truncation via ``status == "incomplete"`` with + ``incomplete_details.reason == "max_output_tokens"`` (``content_filter`` is handled + separately by ``_check_content_filter``). + + Args: + response: A Response object from the OpenAI SDK. - async def _construct_message_from_response_async(self, response: Any, request: MessagePiece) -> Message: + Returns: + bool: True if the response was truncated at the token limit, False otherwise. + """ + if response.status != "incomplete": + return False + incomplete_details = response.incomplete_details + reason = incomplete_details.reason if incomplete_details else None + return reason == "max_output_tokens" + + async def _construct_message_from_response_async(self, response: Response, request: MessagePiece) -> Message: """ Construct a Message from a Response API response. + For a truncated response (see ``_is_truncated_response``), empty output sections are + tolerated, partial tool/function calls are skipped so an incomplete call cannot re-enter the + agentic loop, and a graceful empty text piece is appended when no visible response was + produced. Reasoning, any partial text, and structured refusals are always preserved. + Args: response: The Response object from OpenAI SDK. request: The original request MessagePiece. Returns: - Message: Constructed message with extracted content from output sections. + Message: Constructed message with extracted content from output sections. Token-usage + counts from ``response.usage`` are recorded in the first piece's ``prompt_metadata``. + Truncated responses are flagged via ``MessagePiece.mark_as_truncated`` on the first + piece. """ - # Extract and parse message pieces from validated output sections + truncated = self._is_truncated_response(response) + + # Extract and parse message pieces from validated output sections. A truncated response + # skips the empty-output guard in _validate_response, so ``output`` is falsy-guarded here to + # keep the graceful-empty fallback working even if the section list is missing. extracted_response_pieces: list[MessagePiece] = [] - for section in response.output: + has_visible_response = False + for section in response.output or []: piece = self._parse_response_output_section( section=section, message_piece=request, error=None, # error is already handled in validation + tolerate_empty=truncated, ) if piece is None: continue + # On truncation, drop partial tool/function calls so an incomplete call cannot + # re-enter the agentic loop. Everything else (reasoning, partial text, structured + # refusals) is preserved. + if truncated and piece.original_value_data_type in ("function_call", "tool_call"): + continue extracted_response_pieces.append(piece) + # Reasoning is the one output the caller cannot read as an answer, so anything else + # with a value counts as a visible response and suppresses the empty fallback below. + if piece.original_value and piece.original_value_data_type != "reasoning": + has_visible_response = True + + if truncated and not has_visible_response: + empty_piece = build_empty_truncated_response(request=request).message_pieces[0] + extracted_response_pieces.append(empty_piece) # Consumers use the first piece as the semantic response. Responses API # 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") + # Capture token usage in the first piece's metadata. This also runs on the truncated path: + # usage is populated on token-limit responses and is most valuable there, since the whole + # budget may have been spent on hidden reasoning with no visible answer. + usage = getattr(response, "usage", None) + if usage is not None and extracted_response_pieces: + extracted_response_pieces[0].prompt_metadata.update(token_usage_from_responses(usage).to_metadata()) + + if truncated and extracted_response_pieces: + extracted_response_pieces[0].mark_as_truncated() + return Message(message_pieces=extracted_response_pieces) @limit_requests_per_minute @@ -658,7 +776,8 @@ def _parse_response_message_content( content: list[ResponseOutputText | ResponseOutputRefusal], message_piece: MessagePiece, error: PromptResponseError | None, - ) -> MessagePiece: + tolerate_empty: bool = False, + ) -> MessagePiece | None: """ Parse a Responses API message content union into a PyRIT message piece. @@ -666,16 +785,22 @@ def _parse_response_message_content( content (list[ResponseOutputText | ResponseOutputRefusal]): Typed message content. message_piece (MessagePiece): The original request piece. error (PromptResponseError | None): Any response error classification. + tolerate_empty (bool): When True, empty content returns None instead of raising + EmptyResponseException. Used when constructing a truncated response. Returns: - MessagePiece: A text piece or blocked-error refusal piece. + MessagePiece | None: A text piece or blocked-error refusal piece, or None when the + content is empty and tolerate_empty is True. Raises: - EmptyResponseException: If the message content has no usable value. + EmptyResponseException: If the message content has no usable value (and tolerate_empty + is False). PyritException: If the SDK returns an unsupported message content model. """ if not content: - raise EmptyResponseException(message="The chat returned an empty message section.") + if tolerate_empty: + return None + raise EmptyResponseException(message="The response returned an empty message section.") unsupported = [ content_item @@ -710,7 +835,9 @@ def _parse_response_message_content( piece_value = "\n".join(text_parts) if not piece_value: - raise EmptyResponseException(message="The chat returned an empty response.") + if tolerate_empty: + return None + raise EmptyResponseException(message="The response returned an empty response.") return MessagePiece( role="assistant", original_value=piece_value, @@ -720,7 +847,12 @@ def _parse_response_message_content( ) def _parse_response_output_section( - self, *, section: Any, message_piece: MessagePiece, error: PromptResponseError | None + self, + *, + section: Any, + message_piece: MessagePiece, + error: PromptResponseError | None, + tolerate_empty: bool = False, ) -> MessagePiece | None: """ Parse model output sections, forwarding tool-calls for the agentic loop. @@ -729,12 +861,15 @@ def _parse_response_output_section( section: The section object from OpenAI SDK (Pydantic model). message_piece: The original message piece. error: Any error information from OpenAI. + tolerate_empty: When True, empty sections return None instead of raising + EmptyResponseException. Used when constructing a truncated response. Returns: A MessagePiece for this section, or None to skip. Raises: - EmptyResponseException: If the section content is empty or invalid. + EmptyResponseException: If the section content is empty or invalid (and tolerate_empty + is False). PyritException: If a message section contains an unsupported content model. ValueError: If the section type is unsupported. """ @@ -747,6 +882,7 @@ def _parse_response_output_section( content=section.content, message_piece=message_piece, error=error, + tolerate_empty=tolerate_empty, ) if section_type == MessagePieceType.REASONING: @@ -800,7 +936,9 @@ def _parse_response_output_section( raise ValueError(msg) piece_value = section.input if len(piece_value) == 0: - raise EmptyResponseException(message="The chat returned an empty message section.") + if tolerate_empty: + return None + raise EmptyResponseException(message="The response returned an empty message section.") else: # Other possible types are not yet handled in PyRIT @@ -808,7 +946,9 @@ def _parse_response_output_section( # Handle empty response if not piece_value: - raise EmptyResponseException(message="The chat returned an empty response.") + if tolerate_empty: + return None + raise EmptyResponseException(message="The response returned an empty response.") return MessagePiece( role="assistant", diff --git a/pyrit/prompt_target/openai/openai_target.py b/pyrit/prompt_target/openai/openai_target.py index b943c44a54..c6d0c1961b 100644 --- a/pyrit/prompt_target/openai/openai_target.py +++ b/pyrit/prompt_target/openai/openai_target.py @@ -432,10 +432,8 @@ async def _handle_openai_request_async( if self._check_content_filter(response): return self._handle_content_filter_response(response, request_piece) - # Validate response via subclass implementation - error_message = self._validate_response(response, request_piece) - if error_message: - return error_message + # Validate response via subclass implementation (raises on invalid responses) + self._validate_response(response, request_piece) # Construct and return Message from validated response return await self._construct_message_from_response_async(response, request_piece) @@ -588,24 +586,40 @@ def _extract_partial_content(self, response: Any) -> str | None: """ return None - def _validate_response(self, response: Any, request: MessagePiece) -> Message | None: + def _validate_response(self, response: Any, request: MessagePiece) -> None: """ - Validate the response and return error Message if needed. + Validate the response, raising if it is invalid. - Override this method in subclasses that need custom response validation. - Default implementation returns None (no validation errors). + Override this method in subclasses that need custom response validation. Validation only + inspects the response; constructing the resulting Message is the responsibility of + ``_construct_message_from_response_async``. The default implementation is a no-op. Args: response: The response object from OpenAI SDK. request: The original request MessagePiece. - Returns: - Message | None: Error Message if validation fails, None otherwise. - Raises: Various exceptions for validation failures. """ - return None + + def _is_truncated_response(self, response: Any) -> bool: + """ + Return True if the response was cut off by the output-token limit. + + Every API shape signals truncation differently (Chat Completions + ``finish_reason == "length"``, Responses ``status == "incomplete"`` with + ``reason == "max_output_tokens"``), so subclasses that can detect it override this. A + truncated response is valid but incomplete: ``_validate_response`` warns instead of + raising, and ``_construct_message_from_response_async`` preserves whatever the model + produced. The base implementation reports no truncation. + + Args: + response: The response object from OpenAI SDK. + + Returns: + bool: True if the response was truncated at the token limit, False otherwise. + """ + return False @abstractmethod def _set_openai_env_configuration_vars(self) -> None: diff --git a/tests/unit/models/test_message_piece.py b/tests/unit/models/test_message_piece.py index 40758bb01e..e532d663b9 100644 --- a/tests/unit/models/test_message_piece.py +++ b/tests/unit/models/test_message_piece.py @@ -1120,3 +1120,30 @@ def test_unknown_kwarg_raises(self) -> None: with pytest.raises(Exception) as exc_info: MessagePiece(role="user", original_value="hello", typo_field="oops") assert "typo_field" in str(exc_info.value) or "Extra" in str(exc_info.value) + + +class TestTruncationFlag: + def test_is_truncated_defaults_to_false(self) -> None: + piece = MessagePiece(role="assistant", original_value="hello") + assert piece.is_truncated is False + + def test_mark_as_truncated_sets_flag(self) -> None: + piece = MessagePiece(role="assistant", original_value="partial answer") + piece.mark_as_truncated() + assert piece.is_truncated is True + assert piece.prompt_metadata[MessagePiece.TRUNCATED_METADATA_KEY] is True + + def test_mark_as_truncated_preserves_existing_metadata(self) -> None: + piece = MessagePiece( + role="assistant", original_value="partial", prompt_metadata={"token_usage_output_tokens": 5} + ) + piece.mark_as_truncated() + assert piece.prompt_metadata["token_usage_output_tokens"] == 5 + assert piece.is_truncated is True + + def test_truncated_piece_can_still_report_no_error(self) -> None: + """A truncated partial answer is not an error, so is_truncated is the only signal.""" + piece = MessagePiece(role="assistant", original_value="partial answer") + piece.mark_as_truncated() + assert piece.has_error() is False + assert piece.is_truncated is True diff --git a/tests/unit/models/test_token_usage.py b/tests/unit/models/test_token_usage.py index 439c92efd2..55f24091da 100644 --- a/tests/unit/models/test_token_usage.py +++ b/tests/unit/models/test_token_usage.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from pyrit.models import TokenUsage +from types import SimpleNamespace + +from pyrit.models import TokenUsage, read_usage_int, read_usage_value def test_to_metadata_uses_input_output_key_names_and_omits_none(): @@ -64,3 +66,34 @@ def test_from_metadata_ignores_cost_and_unrelated_keys(): def test_from_metadata_returns_none_without_token_usage_keys(): assert TokenUsage.from_metadata({"partial_content": "x"}) is None + + +def test_read_usage_value_reads_attribute_objects(): + usage = SimpleNamespace(input_tokens=11, input_tokens_details=SimpleNamespace(cached_tokens=3)) + assert read_usage_value(source=usage, name="input_tokens") == 11 + assert read_usage_value(source=usage, name="input_tokens_details").cached_tokens == 3 + + +def test_read_usage_value_reads_mappings(): + usage = {"input_tokens": 11, "input_tokens_details": {"cached_tokens": 3}} + assert read_usage_value(source=usage, name="input_tokens") == 11 + assert read_usage_value(source=usage, name="input_tokens_details") == {"cached_tokens": 3} + + +def test_read_usage_value_returns_none_for_missing_and_none_source(): + assert read_usage_value(source=SimpleNamespace(), name="input_tokens") is None + assert read_usage_value(source=None, name="input_tokens") is None + assert read_usage_value(source={}, name="input_tokens") is None + + +def test_read_usage_int_guards_non_integer_values(): + usage = SimpleNamespace(input_tokens=11, output_tokens=None, total_tokens="30", cached_tokens=True) + assert read_usage_int(source=usage, name="input_tokens") == 11 + assert read_usage_int(source=usage, name="output_tokens") is None + assert read_usage_int(source=usage, name="total_tokens") is None + assert read_usage_int(source=usage, name="cached_tokens") is None + + +def test_read_usage_int_reads_mappings_and_missing_sources(): + assert read_usage_int(source={"input_tokens": 7}, name="input_tokens") == 7 + assert read_usage_int(source=None, name="input_tokens") is None diff --git a/tests/unit/prompt_target/target/test_azure_openai_completion_target.py b/tests/unit/prompt_target/target/test_azure_openai_completion_target.py index 236c967474..a43f2e4fb7 100644 --- a/tests/unit/prompt_target/target/test_azure_openai_completion_target.py +++ b/tests/unit/prompt_target/target/test_azure_openai_completion_target.py @@ -110,3 +110,11 @@ def test_azure_invalid_endpoint_raises(): endpoint="", api_key="xxxxx", ) + + +async def test_completion_target_does_not_detect_truncation(azure_completion_target: OpenAICompletionTarget): + """A target that does not implement truncation detection inherits the base opt-out.""" + response = MagicMock() + response.choices = [MagicMock(finish_reason="length")] + + assert azure_completion_target._is_truncated_response(response) is False diff --git a/tests/unit/prompt_target/target/test_chat_completions_helpers.py b/tests/unit/prompt_target/target/test_chat_completions_helpers.py index 00373fed04..abf1d093aa 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -26,6 +26,7 @@ build_response_pieces_async, capture_token_usage, extract_partial_content, + get_finish_reason, is_content_filter_response, save_audio_response_async, token_usage_from_chat_completion, @@ -134,6 +135,16 @@ def test_validate_response_accepts_valid(): validate_chat_completion_response(response=_mock_response(finish_reason=reason)) +def test_get_finish_reason_returns_first_choice_reason(): + assert get_finish_reason(response=_mock_response(finish_reason="length")) == "length" + + +def test_get_finish_reason_returns_none_without_choices(): + resp = MagicMock() + resp.choices = [] + assert get_finish_reason(response=resp) is None + + def test_capture_token_usage_populates_metadata(): resp = _mock_response("ok") resp.usage.prompt_tokens = 3 diff --git a/tests/unit/prompt_target/target/test_openai_chat_target.py b/tests/unit/prompt_target/target/test_openai_chat_target.py index 534ad11552..2bc2bf76dc 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -1079,11 +1079,78 @@ def test_validate_response_success_stop(target: OpenAIChatTarget, dummy_text_mes assert result is None -def test_validate_response_success_length(target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece): - """Test _validate_response passes for valid length response.""" +def test_validate_response_success_length( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece, caplog: pytest.LogCaptureFixture +): + """Test _validate_response passes for a truncated response that still has content, and warns.""" mock_response = create_mock_completion(content="Hello", finish_reason="length") - result = target._validate_response(mock_response, dummy_text_message_piece) + with caplog.at_level(logging.WARNING): + result = target._validate_response(mock_response, dummy_text_message_piece) + assert result is None + assert "finish_reason='length'" in caplog.text + + +def test_validate_response_length_empty_does_not_raise( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece, caplog: pytest.LogCaptureFixture +): + """Test _validate_response treats a truncated-but-empty response as valid (warns, does not raise).""" + mock_response = create_mock_completion(content="", finish_reason="length") + with caplog.at_level(logging.WARNING): + result = target._validate_response(mock_response, dummy_text_message_piece) assert result is None + assert "finish_reason='length'" in caplog.text + + +def test_is_truncated_response_detects_length_finish_reason(target: OpenAIChatTarget): + """_is_truncated_response is True only when the completion stopped on the token limit.""" + assert target._is_truncated_response(create_mock_completion(content="", finish_reason="length")) is True + assert target._is_truncated_response(create_mock_completion(content="hi", finish_reason="stop")) is False + + +async def test_construct_message_length_empty_returns_graceful_empty_and_captures_usage( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece +): + """Test construction returns a graceful empty piece with usage captured for truncated empty responses.""" + mock_response = create_mock_completion(content="", finish_reason="length") + mock_response.usage = MagicMock() + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 20 + mock_response.usage.total_tokens = 30 + mock_response.usage.completion_tokens_details.reasoning_tokens = 100 + + result = await target._construct_message_from_response_async(mock_response, dummy_text_message_piece) + + assert isinstance(result, Message) + piece = result.message_pieces[0] + assert piece.original_value == "" + assert piece.response_error == "empty" + assert piece.prompt_metadata["token_usage_input_tokens"] == 10 + assert piece.prompt_metadata["token_usage_output_tokens"] == 20 + assert piece.prompt_metadata["token_usage_total_tokens"] == 30 + assert piece.prompt_metadata["token_usage_reasoning_tokens"] == 100 + assert piece.is_truncated is True + + +async def test_construct_message_length_with_content_sets_truncated_metadata( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece +): + """Test construction preserves partial content and marks token-limit truncation.""" + mock_response = create_mock_completion(content="Partial answer", finish_reason="length") + + result = await target._construct_message_from_response_async(mock_response, dummy_text_message_piece) + + piece = result.message_pieces[0] + assert piece.original_value == "Partial answer" + assert piece.is_truncated is True + + +async def test_construct_message_empty_non_truncated_raises( + target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece +): + """Test construction raises for a genuinely empty (non-truncated) response so retries can kick in.""" + mock_response = create_mock_completion(content="", finish_reason="stop") + with pytest.raises(EmptyResponseException, match="Failed to extract any response content"): + await target._construct_message_from_response_async(mock_response, dummy_text_message_piece) def test_validate_response_no_choices(target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece): @@ -2084,6 +2151,7 @@ async def test_construct_message_from_response_captures_token_usage( assert piece.prompt_metadata["token_usage_total_tokens"] == 30 assert piece.prompt_metadata["token_usage_cached_tokens"] == 5 assert piece.prompt_metadata["token_usage_reasoning_tokens"] == 7 + assert piece.is_truncated is False async def test_construct_message_from_response_no_usage_no_metadata( 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 30bba963b3..defa9987d7 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -2,9 +2,11 @@ # Licensed under the MIT license. import json +import logging import os from collections.abc import MutableSequence from tempfile import NamedTemporaryFile +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +30,7 @@ from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import AttackOutcome, JsonResponseConfig, Message, MessagePiece, flatten_to_message_pieces from pyrit.prompt_target import OpenAIResponseTarget, PromptTarget +from pyrit.prompt_target.openai.openai_response_target import token_usage_from_responses from pyrit.score import SelfAskRefusalScorer, TrueFalseInverterScorer @@ -54,6 +57,7 @@ def create_mock_response(response_dict: dict = None) -> MagicMock: # Set attributes based on response_dict to match OpenAI SDK Response type mock_response.error = response_dict.get("error") # Should be None for successful responses mock_response.status = response_dict.get("status") # Should be "completed" for successful responses + mock_response.usage = response_dict.get("usage") # Optional usage payload (None when absent) # Mock the output sections with Pydantic-style attribute access if "output" in response_dict: @@ -1253,9 +1257,196 @@ def test_validate_response_empty_output(target: OpenAIResponseTarget, dummy_text target._validate_response(mock_response, dummy_text_message_piece) +def _make_reasoning_section() -> MagicMock: + section = MagicMock() + section.type = "reasoning" + section.model_dump.return_value = {"type": "reasoning", "summary": []} + return section + + +def _make_message_section(text: str) -> MagicMock: + section = MagicMock() + section.type = "message" + section.content = [ResponseOutputText(annotations=[], text=text, type="output_text")] + return section + + +def _make_empty_message_section() -> MagicMock: + section = MagicMock() + section.type = "message" + section.content = [] + return section + + +def _make_truncated_response(output: list | None) -> MagicMock: + mock_response = MagicMock() + mock_response.error = None + mock_response.status = "incomplete" + incomplete_details = MagicMock() + incomplete_details.reason = "max_output_tokens" + mock_response.incomplete_details = incomplete_details + mock_response.output = output + return mock_response + + +def test_is_truncated_response_detects_max_output_tokens(target: OpenAIResponseTarget): + """_is_truncated_response is True only for incomplete status with a max_output_tokens reason.""" + truncated = _make_truncated_response(output=[]) + assert target._is_truncated_response(truncated) is True + + content_filtered = _make_truncated_response(output=[]) + content_filtered.incomplete_details.reason = "content_filter" + assert target._is_truncated_response(content_filtered) is False + + completed = MagicMock() + completed.status = "completed" + assert target._is_truncated_response(completed) is False + + +def test_validate_response_truncated_warns_and_does_not_raise( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece, caplog: pytest.LogCaptureFixture +): + """Truncation is treated as valid: _validate_response warns and returns None (does not raise).""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_empty_message_section()]) + + with caplog.at_level(logging.WARNING): + result = target._validate_response(response, dummy_text_message_piece) + + assert result is None + assert "max_output_tokens" in caplog.text + + +async def test_construct_message_truncated_keeps_reasoning_and_empty_text( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Truncated response with reasoning but empty text: keep reasoning, add a graceful empty text piece.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_empty_message_section()]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + reasoning_pieces = [p for p in result.message_pieces if p.original_value_data_type == "reasoning"] + text_pieces = [p for p in result.message_pieces if p.original_value_data_type == "text"] + assert len(reasoning_pieces) == 1 + assert len(text_pieces) == 1 + assert text_pieces[0].original_value == "" + assert text_pieces[0].response_error == "empty" + assert result.message_pieces[0].is_truncated is True + + +async def test_construct_message_truncated_keeps_partial_text( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Truncated response with partial visible text keeps it (error=none), no empty piece added.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_message_section("Partial answer")]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + text_pieces = [p for p in result.message_pieces if p.original_value_data_type == "text"] + assert len(text_pieces) == 1 + assert text_pieces[0].original_value == "Partial answer" + assert text_pieces[0].response_error == "none" + assert result.message_pieces[0].is_truncated is True + + +async def test_construct_message_truncated_records_metadata_on_primary_piece_not_reasoning( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Truncation/usage metadata lands on the primary piece, even though reasoning is emitted first.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_message_section("Partial answer")]) + response.usage = _make_usage() + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + primary = result.message_pieces[0] + assert primary.converted_value_data_type == "text" + assert primary.is_truncated is True + assert primary.prompt_metadata["token_usage_reasoning_tokens"] == 7 + assert result.message_pieces[-1].converted_value_data_type == "reasoning" + assert result.message_pieces[-1].is_truncated is False + + +async def test_construct_message_truncated_tolerates_empty_typed_content( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Typed message content that is present but empty is tolerated on the truncated path.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_message_section("")]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + text_pieces = [p for p in result.message_pieces if p.original_value_data_type == "text"] + assert len(text_pieces) == 1 + assert text_pieces[0].original_value == "" + assert text_pieces[0].response_error == "empty" + + +async def test_construct_message_truncated_keeps_structured_refusal( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """A structured refusal is preserved when the response is also truncated.""" + refusal = "I cannot assist with that request." + refusal_section = MagicMock() + refusal_section.type = "message" + refusal_section.content = [ResponseOutputRefusal(refusal=refusal, type="refusal")] + response = _make_truncated_response(output=[refusal_section]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + assert len(result.message_pieces) == 1 + assert result.message_pieces[0].structured_refusal == refusal + assert result.message_pieces[0].is_truncated is True + + +async def test_construct_message_truncated_empty_output_returns_graceful_empty( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Truncated response with no output yields a single graceful empty text piece (does not raise).""" + response = _make_truncated_response(output=[]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + assert len(result.message_pieces) == 1 + assert result.message_pieces[0].original_value == "" + assert result.message_pieces[0].response_error == "empty" + assert result.message_pieces[0].is_truncated is True + + +async def test_construct_message_truncated_none_output_returns_graceful_empty( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Truncated response whose output is None still yields a graceful empty piece (does not raise).""" + response = _make_truncated_response(output=None) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + assert len(result.message_pieces) == 1 + assert result.message_pieces[0].original_value == "" + assert result.message_pieces[0].response_error == "empty" + assert result.message_pieces[0].is_truncated is True + + +async def test_construct_message_truncated_skips_partial_tool_call( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """A partial function_call in a truncated response is skipped so it cannot re-enter the agentic loop.""" + func_section = MagicMock() + func_section.type = "function_call" + func_section.call_id = "call_1" + func_section.name = "do_thing" + func_section.arguments = "{}" + response = _make_truncated_response(output=[_make_reasoning_section(), func_section]) + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + data_types = [p.original_value_data_type for p in result.message_pieces] + assert "function_call" not in data_types + assert "reasoning" in data_types + assert any(p.original_value_data_type == "text" and p.response_error == "empty" for p in result.message_pieces) + + async def test_construct_message_from_response(target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece): """Test _construct_message_from_response parses output sections.""" mock_response = MagicMock() + mock_response.status = "completed" mock_response.output = [{"type": "message", "content": [{"type": "text", "text": "Hello from Response API"}]}] # Mock the _parse_response_output_section method @@ -1272,9 +1463,99 @@ async def test_construct_message_from_response(target: OpenAIResponseTarget, dum assert isinstance(result, Message) assert len(result.message_pieces) == 1 + assert result.message_pieces[0].is_truncated is False mock_parse.assert_called_once() +def _make_usage( + *, + input_tokens: int | None = 11, + output_tokens: int | None = 22, + total_tokens: int | None = 33, + reasoning_tokens: int | None = 7, + cached_tokens: int | None = 3, + cache_write_tokens: int | None = 2, +) -> SimpleNamespace: + return SimpleNamespace( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + input_tokens_details=SimpleNamespace(cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens), + output_tokens_details=SimpleNamespace(reasoning_tokens=reasoning_tokens), + ) + + +def test_token_usage_from_responses_maps_fields(): + """token_usage_from_responses maps the Responses usage shape onto TokenUsage.""" + result = token_usage_from_responses(_make_usage()) + + assert result.input_tokens == 11 + assert result.output_tokens == 22 + assert result.total_tokens == 33 + assert result.reasoning_tokens == 7 + assert result.cached_tokens == 3 + assert result.extra == {"cache_write_tokens": 2} + + +def test_token_usage_from_responses_ignores_missing_and_non_int(): + """Missing details objects and non-integer counts are dropped rather than stored as zero.""" + usage = SimpleNamespace(input_tokens=5, output_tokens=None, input_tokens_details=None, output_tokens_details=None) + + result = token_usage_from_responses(usage) + + assert result.input_tokens == 5 + assert result.output_tokens is None + assert result.total_tokens is None + assert result.reasoning_tokens is None + assert result.cached_tokens is None + assert result.extra == {} + + +def test_token_usage_from_responses_derives_total_when_omitted(): + """A provider that reports only input/output counts still gets a total, as in Chat Completions.""" + usage = SimpleNamespace(input_tokens=5, output_tokens=6, input_tokens_details=None, output_tokens_details=None) + + result = token_usage_from_responses(usage) + + assert result.total_tokens == 11 + + +async def test_construct_message_captures_token_usage( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """A completed response records token-usage counts in the first piece's metadata.""" + response = MagicMock() + response.status = "completed" + response.output = [_make_message_section("Answer")] + response.usage = _make_usage() + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + metadata = result.message_pieces[0].prompt_metadata + assert metadata["token_usage_input_tokens"] == 11 + assert metadata["token_usage_output_tokens"] == 22 + assert metadata["token_usage_total_tokens"] == 33 + assert metadata["token_usage_reasoning_tokens"] == 7 + assert metadata["token_usage_cached_tokens"] == 3 + assert metadata["token_usage_cache_write_tokens"] == 2 + + +async def test_construct_message_truncated_captures_token_usage( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Usage is captured on the truncated path too, alongside the truncated marker.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_empty_message_section()]) + response.usage = _make_usage() + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + piece = result.message_pieces[0] + assert piece.is_truncated is True + assert piece.prompt_metadata["token_usage_input_tokens"] == 11 + assert piece.prompt_metadata["token_usage_output_tokens"] == 22 + assert piece.prompt_metadata["token_usage_reasoning_tokens"] == 7 + + async def test_handle_openai_request_output_text(target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece): output_message = ResponseOutputMessage( id="text-message", diff --git a/tests/unit/prompt_target/test_target_utils.py b/tests/unit/prompt_target/test_target_utils.py index 72c8b64987..1fd83f34c7 100644 --- a/tests/unit/prompt_target/test_target_utils.py +++ b/tests/unit/prompt_target/test_target_utils.py @@ -1,18 +1,26 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest from pyrit.exceptions import PyritException +from pyrit.models import MessagePiece from pyrit.prompt_target.common.utils import ( + build_empty_truncated_response, limit_requests_per_minute, validate_temperature, validate_top_p, + warn_truncated_response, ) +def _request_piece(text: str = "ask") -> MessagePiece: + return MessagePiece(role="user", conversation_id="c", original_value=text, original_value_data_type="text") + + def test_validate_temperature_none(): validate_temperature(None) @@ -102,3 +110,39 @@ async def test_limit_requests_per_minute_zero_rpm(): result = await decorated(mock_self, message="test") mock_sleep.assert_not_called() assert result == "response" + + +def test_build_empty_truncated_response_returns_empty_message(): + request = _request_piece("ask") + result = build_empty_truncated_response(request=request) + + assert result is not None + assert len(result.message_pieces) == 1 + assert result.message_pieces[0].converted_value == "" + assert result.message_pieces[0].converted_value_data_type == "text" + assert result.message_pieces[0].response_error == "empty" + + +def test_warn_truncated_response_names_the_signal_and_limit(caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING): + warn_truncated_response(signal="finish_reason='length'", limit_parameter="max_completion_tokens") + + assert "finish_reason='length'" in caplog.text + assert caplog.text.count("max_completion_tokens") == 2 + + +def test_warn_truncated_response_wording_is_shared_across_api_shapes(caplog: pytest.LogCaptureFixture): + """Only the signal and limit parameter differ between targets; the shared advice must not drift.""" + advice = "Reasoning models consume tokens on hidden reasoning in addition to the visible answer" + + with caplog.at_level(logging.WARNING): + warn_truncated_response(signal="finish_reason='length'", limit_parameter="max_completion_tokens") + warn_truncated_response( + signal="status='incomplete', reason='max_output_tokens'", limit_parameter="max_output_tokens" + ) + + chat_message, responses_message = (record.getMessage() for record in caplog.records) + assert advice in chat_message + assert advice in responses_message + assert "max_output_tokens" in responses_message + assert "max_output_tokens" not in chat_message