FIX: warn instead of raising on reasoning-model token truncation (Chat and Responses targets)

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",
 ]

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
     # ------------------------------------------------------------------ #

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",
 ]

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

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,34 +315,43 @@ 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).
 
     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,

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."
+    )