FEAT: Capture API response stop reason on MessagePiece metadata

pyrit/models/__init__.py

@@ -111,6 +111,7 @@
     COMMON_JSON_SCHEMAS,
     JSON_SCHEMA_METADATA_KEY,
     SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY,
+    TOKEN_USAGE_METADATA_PREFIX,
     CapabilityName,
     JsonResponseConfig,
     JsonSchemaDefinition,
@@ -222,6 +223,7 @@
     "TargetCapabilities",
     "TargetIdentifier",
     "TextDataTypeSerializer",
+    "TOKEN_USAGE_METADATA_PREFIX",
     "TokenUsage",
     "ToolCall",
     "UnvalidatedScore",

pyrit/models/target/__init__.py

@@ -28,7 +28,12 @@
     unregister_common_json_schema,
 )
 from pyrit.models.target.target_capabilities import CapabilityName, TargetCapabilities
-from pyrit.models.target.token_usage import TokenUsage, read_usage_int, read_usage_value
+from pyrit.models.target.token_usage import (
+    TOKEN_USAGE_METADATA_PREFIX,
+    TokenUsage,
+    read_usage_int,
+    read_usage_value,
+)
 
 __all__ = [
     "COMMON_JSON_SCHEMAS",
@@ -37,6 +42,7 @@
     "JsonResponseConfig",
     "JsonSchemaDefinition",
     "SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY",
+    "TOKEN_USAGE_METADATA_PREFIX",
     "TargetCapabilities",
     "TokenUsage",
     "get_common_json_schema",

pyrit/models/target/token_usage.py

@@ -7,8 +7,11 @@
 from dataclasses import dataclass, field
 from typing import Any
 
-#: Prefix for all token-usage keys stored in a MessagePiece's ``prompt_metadata``.
-_METADATA_PREFIX = "token_usage_"
+#: Prefix for all token-usage keys stored in a MessagePiece's ``prompt_metadata``. Public because
+#: the targets that report usage reserve the whole prefix for the provider: they clear it before
+#: writing what the provider reported, so a caller-supplied count is not read back as if the API
+#: had returned it.
+TOKEN_USAGE_METADATA_PREFIX = "token_usage_"
 
 #: Metadata key suffixes that map to first-class ``TokenUsage`` fields. Every other integer
 #: ``token_usage_*`` key round-trips through ``extra``. ``cost`` is not listed because it is a
@@ -108,7 +111,9 @@ def from_metadata(cls, metadata: dict[str, Any]) -> TokenUsage | None:
             TokenUsage | None: The reconstructed usage, or None if no token-usage keys exist.
         """
         stripped = {
-            key[len(_METADATA_PREFIX) :]: value for key, value in metadata.items() if key.startswith(_METADATA_PREFIX)
+            key[len(TOKEN_USAGE_METADATA_PREFIX) :]: value
+            for key, value in metadata.items()
+            if key.startswith(TOKEN_USAGE_METADATA_PREFIX)
         }
         if not stripped:
             return None
@@ -144,15 +149,15 @@ def to_metadata(self) -> dict[str, int]:
         """
         out: dict[str, int] = {}
         if self.input_tokens is not None:
-            out[_METADATA_PREFIX + "input_tokens"] = self.input_tokens
+            out[TOKEN_USAGE_METADATA_PREFIX + "input_tokens"] = self.input_tokens
         if self.output_tokens is not None:
-            out[_METADATA_PREFIX + "output_tokens"] = self.output_tokens
+            out[TOKEN_USAGE_METADATA_PREFIX + "output_tokens"] = self.output_tokens
         if self.total_tokens is not None:
-            out[_METADATA_PREFIX + "total_tokens"] = self.total_tokens
+            out[TOKEN_USAGE_METADATA_PREFIX + "total_tokens"] = self.total_tokens
         if self.reasoning_tokens is not None:
-            out[_METADATA_PREFIX + "reasoning_tokens"] = self.reasoning_tokens
+            out[TOKEN_USAGE_METADATA_PREFIX + "reasoning_tokens"] = self.reasoning_tokens
         if self.cached_tokens is not None:
-            out[_METADATA_PREFIX + "cached_tokens"] = self.cached_tokens
+            out[TOKEN_USAGE_METADATA_PREFIX + "cached_tokens"] = self.cached_tokens
         for name, value in self.extra.items():
-            out[_METADATA_PREFIX + name] = value
+            out[TOKEN_USAGE_METADATA_PREFIX + name] = value
         return out

pyrit/prompt_target/common/chat_completions_response_parser.py

@@ -31,6 +31,7 @@
     read_usage_int,
     read_usage_value,
 )
+from pyrit.prompt_target.common.utils import set_response_metadata, set_token_usage_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -286,20 +287,45 @@ def capture_token_usage(*, pieces: list[MessagePiece], response: Any) -> None:
     Copy token-usage numbers from ``response.usage`` into the first piece's metadata.
 
     Parses the Chat Completions ``usage`` payload (see ``token_usage_from_chat_completion``) and
-    writes the resulting counts onto the first piece. Only fields the provider actually reports are
-    written; missing counts are omitted rather than stored as a misleading zero. No-op when the
-    response has no usage data or there are no pieces.
+    writes the resulting counts onto the first piece via ``set_token_usage_metadata``, which also
+    clears any stale caller-supplied counts. Only fields the provider actually reports are written;
+    missing counts are omitted rather than stored as a misleading zero.
 
     Args:
         pieces (list[MessagePiece]): The constructed response pieces.
         response (Any): The Chat Completions response object.
     """
     usage = getattr(response, "usage", None)
-    if not usage or not pieces:
+    set_token_usage_metadata(pieces=pieces, usage=token_usage_from_chat_completion(usage) if usage else None)
+
+
+def capture_usage_and_finish_reason(*, pieces: list[MessagePiece], response: Any) -> None:
+    """
+    Copy the provider's response metadata into the response pieces' metadata.
+
+    Captures token usage (see ``capture_token_usage``) together with the first choice's
+    ``finish_reason``, which records why generation stopped — ``stop``, ``length`` (token limit),
+    ``content_filter`` or ``tool_calls``. Both come from the same response object, so capturing them
+    together keeps every call site consistent: success, truncation, and content filter all record
+    the same set of keys. Applies to every ``usage``-plus-``choices`` response shape: Chat
+    Completions, Completions, and the LiteLLM responses modeled on them.
+
+    Args:
+        pieces (list[MessagePiece]): The constructed response pieces.
+        response (Any): The response object. Objects that carry no ``usage`` or no ``choices`` —
+            such as the synthetic response used when the SDK raises on a content filter — are
+            tolerated and simply yield no metadata.
+    """
+    if not pieces:
         return
 
-    token_usage = token_usage_from_chat_completion(usage)
-    pieces[0].prompt_metadata.update(token_usage.to_metadata())
+    capture_token_usage(pieces=pieces, response=response)
+
+    # Synthetic content-filter responses have no ``choices`` attribute at all, so probe for it
+    # rather than relying on ``get_finish_reason``'s empty-list guard.
+    choices = getattr(response, "choices", None)
+    finish_reason = get_finish_reason(response=response) if choices else None
+    set_response_metadata(pieces=pieces, values={"finish_reason": finish_reason})
 
 
 def token_usage_from_chat_completion(usage: Any) -> TokenUsage:

pyrit/prompt_target/common/utils.py

@@ -3,11 +3,17 @@
 
 import asyncio
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Any
 
 from pyrit.exceptions import PyritException
-from pyrit.models import Message, MessagePiece, construct_response_from_request
+from pyrit.models import (
+    TOKEN_USAGE_METADATA_PREFIX,
+    Message,
+    MessagePiece,
+    TokenUsage,
+    construct_response_from_request,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -87,6 +93,70 @@ def build_empty_truncated_response(*, request: MessagePiece) -> Message:
     )
 
 
+#: ``prompt_metadata`` keys that record why a provider stopped generating. Every API shape names this
+#: differently, so the union is reserved rather than any single key: whichever target handles the
+#: response clears all of them and writes back only what its own provider reported.
+RESERVED_RESPONSE_METADATA_KEYS: frozenset[str] = frozenset({"finish_reason", "status", "incomplete_reason"})
+
+
+def set_response_metadata(*, pieces: list[MessagePiece], values: Mapping[str, Any]) -> None:
+    """
+    Record provider-reported, response-level metadata on the first response piece.
+
+    ``prompt_metadata`` is caller-controlled, and ``construct_response_from_request`` merges the
+    request's entries into every response piece, so a caller-supplied value could otherwise be
+    mistaken for the provider's. ``RESERVED_RESPONSE_METADATA_KEYS`` are therefore reserved for the
+    provider: all of them are cleared from every piece first, then the reported ones are set on the
+    first piece. Clearing the whole set in one pass — rather than one key per call — is what makes
+    the reservation hold, since a target only writes the subset its own API reports. Response-level
+    metadata lives on the first piece, matching where ``capture_token_usage`` writes token counts.
+
+    Args:
+        pieces (list[MessagePiece]): The constructed response pieces.
+        values (Mapping[str, Any]): The provider-reported values, keyed by reserved metadata key.
+            ``prompt_metadata`` is persisted as JSON, so anything that is not a non-empty string —
+            including a missing field read off a loosely-typed response object — is treated as "not
+            reported" and leaves that key unset.
+    """
+    if not pieces:
+        return
+
+    for piece in pieces:
+        for reserved_key in RESERVED_RESPONSE_METADATA_KEYS:
+            piece.prompt_metadata.pop(reserved_key, None)
+
+    for key, value in values.items():
+        if isinstance(value, str) and value:
+            pieces[0].prompt_metadata[key] = value
+
+
+def set_token_usage_metadata(*, pieces: list[MessagePiece], usage: TokenUsage | None) -> None:
+    """
+    Record the provider's token counts on the first response piece.
+
+    The whole ``token_usage_`` prefix is reserved for the provider for the same reason
+    ``RESERVED_RESPONSE_METADATA_KEYS`` are, with one extra consequence: the public
+    ``TokenUsage.from_metadata`` would read a caller-supplied count back as if the API had reported
+    it. Clearing the prefix is what makes "no usage reported" distinguishable from "the caller
+    guessed", which matters most on the paths that carry no usage at all, such as a content-filtered
+    response.
+
+    Args:
+        pieces (list[MessagePiece]): The constructed response pieces.
+        usage (TokenUsage | None): The provider's parsed counts, or None when the response reported
+            no usage. Either way the stale keys are cleared first; only reported counts are written.
+    """
+    if not pieces:
+        return
+
+    for piece in pieces:
+        for key in [k for k in piece.prompt_metadata if k.startswith(TOKEN_USAGE_METADATA_PREFIX)]:
+            del piece.prompt_metadata[key]
+
+    if usage is not None:
+        pieces[0].prompt_metadata.update(usage.to_metadata())
+
+
 def warn_truncated_response(*, signal: str, limit_parameter: str) -> None:
     """
     Log the shared warning for a response cut off at the output-token limit.

pyrit/prompt_target/litellm_chat_target.py

@@ -31,7 +31,7 @@
 from pyrit.prompt_target.common.chat_completions_response_parser import (
     build_content_filter_message,
     build_response_pieces_async,
-    capture_token_usage,
+    capture_usage_and_finish_reason,
     extract_partial_content,
     is_content_filter_response,
     validate_chat_completion_response,
@@ -382,13 +382,16 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
         # exception) so attacks can continue and blocked-content scorers can still score.
         if is_content_filter_response(response):
             logger.warning("Output content filtered by content policy.")
-            return [
-                build_content_filter_message(
-                    response=response,
-                    request=request_piece,
-                    partial_content=extract_partial_content(response),
-                )
-            ]
+            filter_message = build_content_filter_message(
+                response=response,
+                request=request_piece,
+                partial_content=extract_partial_content(response),
+            )
+            # A filtered response still reports the tokens (and cost) it consumed, which would
+            # otherwise be lost precisely where a red teamer most wants to see it.
+            capture_usage_and_finish_reason(pieces=filter_message.message_pieces, response=response)
+            self._capture_response_cost(pieces=filter_message.message_pieces, response=response)
+            return [filter_message]
 
         validate_chat_completion_response(response=response)
         return [await self._construct_message_from_response_async(response=response, request=request_piece)]
@@ -464,7 +467,7 @@ async def _construct_message_from_response_async(self, *, response: Any, request
         pieces = await build_response_pieces_async(response=response, request=request, audio_format=audio_format)
         if not pieces:
             raise EmptyResponseException(message="Failed to extract any response content from LiteLLM.")
-        capture_token_usage(pieces=pieces, response=response)
+        capture_usage_and_finish_reason(pieces=pieces, response=response)
         self._capture_response_cost(pieces=pieces, response=response)
         return Message(message_pieces=pieces)
 

pyrit/prompt_target/openai/openai_chat_target.py

@@ -27,7 +27,7 @@
 )
 from pyrit.prompt_target.common.chat_completions_response_parser import (
     build_response_pieces_async,
-    capture_token_usage,
+    capture_usage_and_finish_reason,
     detect_response_content,
     extract_partial_content,
     get_finish_reason,
@@ -276,6 +276,17 @@ def _extract_partial_content(self, response: Any) -> str | None:
         """
         return extract_partial_content(response)
 
+    def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece]) -> None:
+        """
+        Record token usage and ``finish_reason`` from a Chat Completions response.
+
+        Args:
+            response: A ChatCompletion object from the OpenAI SDK, or the synthetic stand-in used
+                when the SDK raises on a content filter.
+            pieces (list[MessagePiece]): The constructed response pieces.
+        """
+        capture_usage_and_finish_reason(pieces=pieces, response=response)
+
     def _validate_response(self, response: ChatCompletion, request: MessagePiece) -> None:
         """
         Validate a Chat Completions API response for errors.
@@ -403,13 +414,13 @@ async def _construct_message_from_response_async(self, response: ChatCompletion,
             # 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)
+                self._capture_response_metadata(response=response, pieces=empty_message.message_pieces)
                 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)
+        # Capture token usage and the stop reason from the API response into the first piece.
+        self._capture_response_metadata(response=response, pieces=pieces)
         if truncated:
             pieces[0].mark_as_truncated()
 

pyrit/prompt_target/openai/openai_completion_target.py

@@ -8,10 +8,13 @@
 from pyrit.exceptions.exception_classes import (
     pyrit_target_retry,
 )
-from pyrit.models import ComponentIdentifier, Message, construct_response_from_request
+from pyrit.models import ComponentIdentifier, Message, MessagePiece, construct_response_from_request
+from pyrit.prompt_target.common.chat_completions_response_parser import (
+    capture_token_usage,
+)
 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 limit_requests_per_minute
+from pyrit.prompt_target.common.utils import limit_requests_per_minute, set_response_metadata
 from pyrit.prompt_target.openai.openai_target import OpenAITarget
 
 logger = logging.getLogger(__name__)
@@ -155,6 +158,27 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
         )
         return [response]
 
+    def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece]) -> None:
+        """
+        Record token usage and each choice's ``finish_reason`` from a Completion response.
+
+        Usage is per-call and lands on the first piece, as everywhere else. ``finish_reason`` is
+        per-choice, and this is the one target that maps a piece to each choice, so each piece gets
+        its own: reading only ``choices[0]`` would report one generation's stop reason for all of
+        them and hide a content filter that tripped on a later choice.
+
+        Args:
+            response: A Completion object from the OpenAI SDK, or the synthetic stand-in used when
+                the SDK raises on a content filter.
+            pieces (list[MessagePiece]): The constructed response pieces.
+        """
+        capture_token_usage(pieces=pieces, response=response)
+
+        choices = getattr(response, "choices", None) or []
+        for index, piece in enumerate(pieces):
+            choice = choices[index] if index < len(choices) else None
+            set_response_metadata(pieces=[piece], values={"finish_reason": getattr(choice, "finish_reason", None)})
+
     async def _construct_message_from_response_async(self, response: Any, request: Any) -> Message:
         """
         Construct a Message from a Completion response.
@@ -171,4 +195,6 @@ async def _construct_message_from_response_async(self, response: Any, request: A
         # Extract response text from validated choices
         extracted_response = [choice.text for choice in response.choices]
 
-        return construct_response_from_request(request=request, response_text_pieces=extracted_response)
+        message = construct_response_from_request(request=request, response_text_pieces=extracted_response)
+        self._capture_response_metadata(response=response, pieces=message.message_pieces)
+        return message