From f6a684bb2d9a57d88b0e6f557db6453f0781af38 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 5 Aug 2026 20:17:15 +0000 Subject: [PATCH 1/3] FEAT: Capture API response stop reason on MessagePiece metadata Targets already recorded token usage but discarded why generation stopped, and dropped usage entirely on content-filtered responses. Capture the provider's stop reason alongside token usage: finish_reason for Chat Completions, Completions and LiteLLM; status plus incomplete_reason for the Responses API. A base no-op hook on OpenAITarget is called from _handle_content_filter_response, so a filtered response now records the tokens it consumed instead of returning bare metadata. These keys are reserved for the provider. construct_response_from_request merges the request's metadata into every response piece, so all of them are cleared before a capture writes back the subset its own API reports. --- pyrit/models/__init__.py | 2 + pyrit/models/target/__init__.py | 8 +- pyrit/models/target/token_usage.py | 23 ++- .../chat_completions_response_parser.py | 38 +++- pyrit/prompt_target/common/utils.py | 74 +++++++- pyrit/prompt_target/litellm_chat_target.py | 21 ++- .../openai/openai_chat_target.py | 19 +- .../openai/openai_completion_target.py | 32 +++- .../openai/openai_response_target.py | 38 +++- pyrit/prompt_target/openai/openai_target.py | 23 +++ .../test_openai_chat_target_integration.py | 5 + .../test_interface_prompts.py | 27 +++ .../test_azure_openai_completion_target.py | 82 +++++++++ .../target/test_chat_completions_helpers.py | 168 +++++++++++++++++- .../target/test_litellm_chat_target.py | 24 +++ .../target/test_openai_chat_target.py | 77 ++++++++ .../target/test_openai_response_target.py | 69 +++++++ .../target/test_prompt_target.py | 3 +- 18 files changed, 690 insertions(+), 43 deletions(-) diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 1f07548d1e..289187f44c 100644 --- a/pyrit/models/__init__.py +++ b/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", diff --git a/pyrit/models/target/__init__.py b/pyrit/models/target/__init__.py index 7e3455bfaf..554154bdd8 100644 --- a/pyrit/models/target/__init__.py +++ b/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", diff --git a/pyrit/models/target/token_usage.py b/pyrit/models/target/token_usage.py index 08374788e0..19c3f755f3 100644 --- a/pyrit/models/target/token_usage.py +++ b/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 diff --git a/pyrit/prompt_target/common/chat_completions_response_parser.py b/pyrit/prompt_target/common/chat_completions_response_parser.py index 16c7cb5256..a1358ad206 100644 --- a/pyrit/prompt_target/common/chat_completions_response_parser.py +++ b/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: diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index 92fdffff70..ac51efec28 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/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. diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index 070ede86c8..f5d1767d6a 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/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) diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 3d9f00a2e8..ff82e3bc78 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/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() diff --git a/pyrit/prompt_target/openai/openai_completion_target.py b/pyrit/prompt_target/openai/openai_completion_target.py index af34dcc270..81fb6d5ca0 100644 --- a/pyrit/prompt_target/openai/openai_completion_target.py +++ b/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 diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 62e384a739..14f34c6a0e 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -38,6 +38,8 @@ from pyrit.prompt_target.common.utils import ( build_empty_truncated_response, limit_requests_per_minute, + set_response_metadata, + set_token_usage_metadata, validate_temperature, validate_top_p, warn_truncated_response, @@ -562,6 +564,31 @@ def _extract_partial_content(self, response: Any) -> str | None: except (AttributeError, IndexError, TypeError): return None + def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece]) -> None: + """ + Record token usage, ``status`` and ``incomplete_reason`` from a Responses API response. + + The Responses API reports why generation stopped as ``status`` plus, when the status is + ``incomplete``, ``incomplete_details.reason`` (for example ``max_output_tokens`` or + ``content_filter``). Together they are this format's equivalent of Chat Completions' + ``finish_reason``. + + Args: + response: A Response 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. + """ + if not pieces: + return + + usage = getattr(response, "usage", None) + set_token_usage_metadata(pieces=pieces, usage=token_usage_from_responses(usage) if usage is not None else None) + + status = getattr(response, "status", None) + incomplete_details = getattr(response, "incomplete_details", None) + incomplete_reason = getattr(incomplete_details, "reason", None) if incomplete_details else None + set_response_metadata(pieces=pieces, values={"status": status, "incomplete_reason": incomplete_reason}) + def _validate_response(self, response: Response, request: MessagePiece) -> None: """ Validate a Response API response for errors. @@ -688,12 +715,11 @@ async def _construct_message_from_response_async(self, response: Response, reque # 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()) + # Capture token usage and the stop reason in the first piece's metadata. This also runs on + # the truncated path: usage is populated on token-limit responses and is most valuable + # there, since the whole budget may have been spent on hidden reasoning with no visible + # answer. + self._capture_response_metadata(response=response, pieces=extracted_response_pieces) if truncated and extracted_response_pieces: extracted_response_pieces[0].mark_as_truncated() diff --git a/pyrit/prompt_target/openai/openai_target.py b/pyrit/prompt_target/openai/openai_target.py index c6d0c1961b..aea513a065 100644 --- a/pyrit/prompt_target/openai/openai_target.py +++ b/pyrit/prompt_target/openai/openai_target.py @@ -547,6 +547,9 @@ def _handle_content_filter_response(self, response: Any, request: MessagePiece) it is attached to each response piece as ``prompt_metadata["partial_content"]`` so that scorers with ``score_blocked_content=True`` can evaluate it. + Provider-reported metadata (token usage, stop reason) is captured via + ``_capture_response_metadata`` so a filtered response records what it consumed. + Args: response: The response object from OpenAI SDK. request: The original request message piece. @@ -569,6 +572,8 @@ def _handle_content_filter_response(self, response: Any, request: MessagePiece) for piece in error_message.message_pieces: piece.prompt_metadata["partial_content"] = partial_content + self._capture_response_metadata(response=response, pieces=error_message.message_pieces) + return error_message def _extract_partial_content(self, response: Any) -> str | None: @@ -586,6 +591,24 @@ def _extract_partial_content(self, response: Any) -> str | None: """ return None + def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece]) -> None: + """ + Record provider-reported response metadata (token usage, stop reason) onto the pieces. + + Override this in subclasses to read API-specific response structures. The base + implementation is a no-op. + + Subclasses call this on their success path and the base class calls it on the + content-filter path, so a single override covers both. A filtered response still reports + the tokens it consumed, which would otherwise be lost precisely where a red teamer most + wants to see it. + + Args: + response: The response object from OpenAI SDK. May be a synthetic stand-in that carries + no usage or completion data, so implementations must tolerate missing attributes. + pieces (list[MessagePiece]): The constructed response pieces. + """ + def _validate_response(self, response: Any, request: MessagePiece) -> None: """ Validate the response, raising if it is invalid. diff --git a/tests/integration/targets/test_openai_chat_target_integration.py b/tests/integration/targets/test_openai_chat_target_integration.py index 4e631b5ad4..d86944c7c4 100644 --- a/tests/integration/targets/test_openai_chat_target_integration.py +++ b/tests/integration/targets/test_openai_chat_target_integration.py @@ -18,6 +18,9 @@ from pyrit.common.path import HOME_PATH from pyrit.models import MessagePiece, TokenUsage from pyrit.prompt_target import OpenAIChatAudioConfig, OpenAIChatTarget, TargetCapabilities, TargetConfiguration +from pyrit.prompt_target.common.chat_completions_response_parser import ( + DEFAULT_VALID_FINISH_REASONS, +) # Path to sample audio file for testing SAMPLE_AUDIO_FILE = HOME_PATH / "assets" / "converted_audio.wav" @@ -234,6 +237,7 @@ async def test_openai_chat_target_token_usage_in_metadata(sqlite_instance, azure 1. Token usage is recoverable via ``TokenUsage.from_metadata`` 2. Token counts are positive integers 3. The total equals input + output + 4. The provider's ``finish_reason`` is captured alongside it """ target = OpenAIChatTarget(**azure_gpt5_chat_args) @@ -258,3 +262,4 @@ async def test_openai_chat_target_token_usage_in_metadata(sqlite_instance, azure assert usage.output_tokens is not None and usage.output_tokens > 0 assert usage.total_tokens is not None and usage.total_tokens > 0 assert usage.total_tokens == usage.input_tokens + usage.output_tokens + assert first_piece.prompt_metadata["finish_reason"] in DEFAULT_VALID_FINISH_REASONS diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index f295d6f168..00b1ad0104 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -31,6 +31,7 @@ SeedPrompt, TargetIdentifier, ) +from pyrit.prompt_target.common.utils import set_response_metadata def _test_scorer_id(name: str = "TestScorer") -> ComponentIdentifier: @@ -948,6 +949,32 @@ def test_get_message_pieces_metadata(sqlite_instance: MemoryInterface): assert "key2" in retrieved_entry.prompt_metadata +@pytest.mark.parametrize( + "key, value", + [("finish_reason", "content_filter"), ("status", "incomplete"), ("incomplete_reason", "max_output_tokens")], +) +def test_get_message_pieces_captured_response_metadata(sqlite_instance: MemoryInterface, key: str, value: str): + """The response metadata captured by the targets must be queryable after a round trip.""" + matching = MessagePiece( + conversation_id=str(uuid4()), + role="assistant", + original_value="blocked", + ) + other = MessagePiece( + conversation_id=str(uuid4()), + role="assistant", + original_value="fine", + ) + set_response_metadata(pieces=[matching], values={key: value}) + set_response_metadata(pieces=[other], values={key: "something_else"}) + sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=matching), PromptMemoryEntry(entry=other)]) + + retrieved = sqlite_instance.get_message_pieces(prompt_metadata={key: value}) + + assert len(retrieved) == 1 + assert retrieved[0].prompt_metadata[key] == value + + def test_get_message_pieces_id(sqlite_instance: MemoryInterface): entries = [ PromptMemoryEntry( 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 a43f2e4fb7..9d3f8b530d 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 @@ -6,6 +6,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from openai.types.completion import Completion +from openai.types.completion_choice import CompletionChoice +from openai.types.completion_usage import CompletionUsage from unit.mocks import get_image_message_piece, get_sample_conversations from pyrit.memory.central_memory import CentralMemory @@ -118,3 +121,82 @@ async def test_completion_target_does_not_detect_truncation(azure_completion_tar response.choices = [MagicMock(finish_reason="length")] assert azure_completion_target._is_truncated_response(response) is False + + +@pytest.mark.parametrize("finish_reason", ["stop", "length", "content_filter"]) +async def test_completion_target_captures_usage_and_finish_reason( + azure_completion_target: OpenAICompletionTarget, + sample_conversations: MutableSequence[MessagePiece], + finish_reason: str, +): + """The Completions API reports the same usage and finish_reason fields as Chat Completions.""" + response = Completion( + id="cmpl-1", + object="text_completion", + created=0, + model="gpt-35-turbo", + choices=[CompletionChoice(finish_reason=finish_reason, index=0, text="hi")], + usage=CompletionUsage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + + message = await azure_completion_target._construct_message_from_response_async( + response=response, request=sample_conversations[0] + ) + + metadata = message.message_pieces[0].prompt_metadata + assert metadata["finish_reason"] == finish_reason + assert metadata["token_usage_input_tokens"] == 11 + assert metadata["token_usage_output_tokens"] == 7 + assert metadata["token_usage_total_tokens"] == 18 + + +async def test_completion_target_captures_finish_reason_per_choice( + azure_completion_target: OpenAICompletionTarget, + sample_conversations: MutableSequence[MessagePiece], +): + """With n>1 each piece is its own choice, so a filter on a later choice must not be hidden.""" + response = Completion( + id="cmpl-1", + object="text_completion", + created=0, + model="gpt-35-turbo", + choices=[ + CompletionChoice(finish_reason="stop", index=0, text="allowed"), + CompletionChoice(finish_reason="content_filter", index=1, text=""), + ], + usage=CompletionUsage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + + message = await azure_completion_target._construct_message_from_response_async( + response=response, request=sample_conversations[0] + ) + + assert [piece.prompt_metadata.get("finish_reason") for piece in message.message_pieces] == [ + "stop", + "content_filter", + ] + # Usage is per call, not per choice, so it stays on the first piece only. + assert message.message_pieces[0].prompt_metadata["token_usage_total_tokens"] == 18 + assert "token_usage_total_tokens" not in message.message_pieces[1].prompt_metadata + + +async def test_completion_target_clears_pieces_without_a_matching_choice( + azure_completion_target: OpenAICompletionTarget, + sample_conversations: MutableSequence[MessagePiece], +): + """Every piece is cleared, so a piece the provider said nothing about reports nothing.""" + pieces = [sample_conversations[0], sample_conversations[1]] + for piece in pieces: + piece.prompt_metadata["finish_reason"] = "caller_supplied" + response = Completion( + id="cmpl-1", + object="text_completion", + created=0, + model="gpt-35-turbo", + choices=[CompletionChoice(finish_reason="stop", index=0, text="allowed")], + usage=None, + ) + + azure_completion_target._capture_response_metadata(response=response, pieces=pieces) + + assert [piece.prompt_metadata.get("finish_reason") for piece in pieces] == ["stop", None] 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 abf1d093aa..0fa6406c7c 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -11,7 +11,7 @@ import pytest from pyrit.exceptions import EmptyResponseException, PyritException -from pyrit.models import JsonResponseConfig, Message, MessagePiece +from pyrit.models import JsonResponseConfig, Message, MessagePiece, TokenUsage from pyrit.prompt_target.common.chat_completions_message_builder import ( build_multimodal_chat_messages_async, build_response_format, @@ -25,6 +25,7 @@ build_content_filter_message, build_response_pieces_async, capture_token_usage, + capture_usage_and_finish_reason, extract_partial_content, get_finish_reason, is_content_filter_response, @@ -32,6 +33,7 @@ token_usage_from_chat_completion, validate_chat_completion_response, ) +from pyrit.prompt_target.common.utils import RESERVED_RESPONSE_METADATA_KEYS, set_response_metadata pytestmark = pytest.mark.usefixtures("patch_central_database") @@ -163,7 +165,7 @@ def test_capture_token_usage_populates_metadata(): assert "token_usage_model_name" not in metadata -def test_capture_token_usage_noop_without_usage(): +def test_capture_token_usage_writes_nothing_without_usage(): resp = _mock_response("ok") resp.usage = None pieces = [_request_piece("ok")] @@ -171,6 +173,168 @@ def test_capture_token_usage_noop_without_usage(): assert "token_usage_total_tokens" not in pieces[0].prompt_metadata +# --------------------------------------------------------------------------- +# response metadata capture: token usage plus the stop reason +# --------------------------------------------------------------------------- + + +class _SyntheticContentFilterResponse: + """Mirrors ``OpenAITarget``'s synthetic stand-in: no ``usage``, no ``choices``.""" + + def model_dump_json(self) -> str: + return "{}" + + +def test_capture_usage_and_finish_reason_captures_usage_and_finish_reason(): + resp = _mock_response("ok", finish_reason="length") + resp.usage.prompt_tokens = 3 + resp.usage.completion_tokens = 4 + resp.usage.total_tokens = 7 + resp.usage.prompt_tokens_details.cached_tokens = 1 + resp.usage.completion_tokens_details.reasoning_tokens = 2 + pieces = [_request_piece("ok")] + + capture_usage_and_finish_reason(pieces=pieces, response=resp) + + metadata = pieces[0].prompt_metadata + assert metadata["finish_reason"] == "length" + assert metadata["token_usage_input_tokens"] == 3 + assert metadata["token_usage_output_tokens"] == 4 + assert metadata["token_usage_reasoning_tokens"] == 2 + + +@pytest.mark.parametrize("finish_reason", ["stop", "length", "content_filter", "tool_calls"]) +def test_capture_usage_and_finish_reason_records_each_finish_reason(finish_reason): + pieces = [_request_piece("ok")] + capture_usage_and_finish_reason(pieces=pieces, response=_mock_response("ok", finish_reason=finish_reason)) + assert pieces[0].prompt_metadata["finish_reason"] == finish_reason + + +def test_capture_usage_and_finish_reason_stores_finish_reason_as_string(): + """``prompt_metadata`` is persisted as JSON and queried as a string.""" + pieces = [_request_piece("ok")] + capture_usage_and_finish_reason(pieces=pieces, response=_mock_response("ok")) + assert isinstance(pieces[0].prompt_metadata["finish_reason"], str) + + +def test_capture_usage_and_finish_reason_writes_only_to_first_piece(): + resp = _mock_response("ok") + resp.usage = None + pieces = [_request_piece("a"), _request_piece("b")] + capture_usage_and_finish_reason(pieces=pieces, response=resp) + assert pieces[0].prompt_metadata["finish_reason"] == "stop" + assert "finish_reason" not in pieces[1].prompt_metadata + + +def test_capture_usage_and_finish_reason_clears_stale_metadata_from_every_piece(): + """Request metadata is merged into every piece, so a stale value must not survive on any of them.""" + stale = {"finish_reason": "caller_supplied", "status": "caller_supplied", "token_usage_input_tokens": 999999} + pieces = [_request_piece("a"), _request_piece("b")] + for piece in pieces: + piece.prompt_metadata.update(stale) + resp = _mock_response("ok", finish_reason="length") + resp.usage = _usage(prompt_tokens=3, completion_tokens=4, total_tokens=7) + + capture_usage_and_finish_reason(pieces=pieces, response=resp) + + assert pieces[0].prompt_metadata["finish_reason"] == "length" + assert pieces[0].prompt_metadata["token_usage_input_tokens"] == 3 + assert "status" not in pieces[0].prompt_metadata + assert not any(key in pieces[1].prompt_metadata for key in stale) + + +def test_capture_usage_and_finish_reason_tolerates_response_without_choices(): + """The SDK-raised content-filter path passes an object with neither usage nor choices.""" + pieces = [_request_piece("ok")] + capture_usage_and_finish_reason(pieces=pieces, response=_SyntheticContentFilterResponse()) + assert pieces[0].prompt_metadata == {} + + +def test_capture_usage_and_finish_reason_noop_without_pieces(): + capture_usage_and_finish_reason(pieces=[], response=_mock_response("ok")) + + +def test_capture_usage_and_finish_reason_clears_caller_supplied_finish_reason(): + """``finish_reason`` is reserved for the provider, so an inherited value must not survive.""" + piece = _request_piece("ok") + piece.prompt_metadata["finish_reason"] = "caller_supplied" + capture_usage_and_finish_reason(pieces=[piece], response=_SyntheticContentFilterResponse()) + assert "finish_reason" not in piece.prompt_metadata + + +def test_capture_usage_and_finish_reason_overwrites_caller_supplied_finish_reason(): + piece = _request_piece("ok") + piece.prompt_metadata["finish_reason"] = "caller_supplied" + capture_usage_and_finish_reason(pieces=[piece], response=_mock_response("ok", finish_reason="length")) + assert piece.prompt_metadata["finish_reason"] == "length" + + +@pytest.mark.parametrize("value", ["", None, 0, MagicMock()]) +def test_set_response_metadata_ignores_unreported_values(value): + """``prompt_metadata`` is JSON-serialized, so anything but a non-empty string is not reported.""" + piece = _request_piece("ok") + set_response_metadata(pieces=[piece], values={"status": value}) + assert "status" not in piece.prompt_metadata + + +@pytest.mark.parametrize("reserved_key", sorted(RESERVED_RESPONSE_METADATA_KEYS)) +def test_set_response_metadata_clears_every_reserved_key(reserved_key): + """A target only writes the keys its own API reports, so all of them must be cleared.""" + piece = _request_piece("ok") + piece.prompt_metadata[reserved_key] = "caller_supplied" + + set_response_metadata(pieces=[piece], values={"finish_reason": "stop"}) + + assert piece.prompt_metadata.get(reserved_key) == ("stop" if reserved_key == "finish_reason" else None) + + +def test_set_response_metadata_keeps_all_reported_values(): + """Clearing runs once up front, so a second reported key must not wipe the first.""" + piece = _request_piece("ok") + + set_response_metadata(pieces=[piece], values={"status": "incomplete", "incomplete_reason": "max_output_tokens"}) + + assert piece.prompt_metadata["status"] == "incomplete" + assert piece.prompt_metadata["incomplete_reason"] == "max_output_tokens" + + +def test_set_response_metadata_leaves_unreserved_caller_metadata_untouched(): + piece = _request_piece("ok") + piece.prompt_metadata["video_id"] = "caller_supplied" + + set_response_metadata(pieces=[piece], values={"finish_reason": "stop"}) + + assert piece.prompt_metadata["video_id"] == "caller_supplied" + + +def test_reserved_response_metadata_keys_are_the_stop_reason_keys(): + """Pinned explicitly: parametrizing over the set lets a dropped key delete its own test case.""" + assert {"finish_reason", "status", "incomplete_reason"} == RESERVED_RESPONSE_METADATA_KEYS + + +def test_capture_token_usage_clears_caller_supplied_counts_when_none_reported(): + """The whole prefix is reserved, so a guess must not read back as what the provider charged.""" + piece = _request_piece("ok") + piece.prompt_metadata.update({"token_usage_input_tokens": 999999, "token_usage_bogus": 777}) + + capture_token_usage(pieces=[piece], response=_SyntheticContentFilterResponse()) + + assert TokenUsage.from_metadata(piece.prompt_metadata) is None + + +def test_capture_token_usage_replaces_stale_counts_the_provider_did_not_report(): + """A reported payload must replace the caller's leftovers, not merge into them.""" + piece = _request_piece("ok") + piece.prompt_metadata["token_usage_reasoning_tokens"] = 999999 + resp = _mock_response("ok") + resp.usage = _usage(prompt_tokens=3, completion_tokens=4, total_tokens=7) + + capture_token_usage(pieces=[piece], response=resp) + + assert "token_usage_reasoning_tokens" not in piece.prompt_metadata + assert piece.prompt_metadata["token_usage_input_tokens"] == 3 + + # --------------------------------------------------------------------------- # token_usage_from_chat_completion (Chat Completions usage parsing) # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_litellm_chat_target.py b/tests/unit/prompt_target/target/test_litellm_chat_target.py index 30cd9245b5..dcf3cf91ec 100644 --- a/tests/unit/prompt_target/target/test_litellm_chat_target.py +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -395,6 +395,15 @@ async def test_send_prompt_captures_token_usage(target, litellm_stub): assert metadata["token_usage_total_tokens"] == 15 +async def test_send_prompt_captures_finish_reason(target, litellm_stub): + """LiteLLM normalizes provider stop reasons (Anthropic ``end_turn`` becomes ``stop``).""" + litellm_stub.acompletion = AsyncMock(return_value=_mock_response("ok", finish_reason="stop")) + + result = await target.send_prompt_async(message=_user_message("hi")) + + assert result[0].message_pieces[0].prompt_metadata["finish_reason"] == "stop" + + async def test_send_prompt_captures_response_cost_from_hidden_params(target, litellm_stub): response = _mock_response("ok") response._hidden_params = {"response_cost": 0.00042} @@ -574,6 +583,21 @@ async def test_content_filter_finish_reason_returns_error_message(target, litell assert piece.prompt_metadata.get("partial_content") == "partial answer" +async def test_content_filter_captures_token_usage_finish_reason_and_cost(target, litellm_stub): + """A blocked response still reports the tokens and spend it consumed.""" + response = _mock_response(content="partial answer", finish_reason="content_filter") + response._hidden_params = {"response_cost": 0.00042} + litellm_stub.acompletion = AsyncMock(return_value=response) + + result = await target.send_prompt_async(message=_user_message("bad prompt")) + + metadata = result[0].message_pieces[0].prompt_metadata + assert metadata["finish_reason"] == "content_filter" + assert metadata["token_usage_input_tokens"] == 10 + assert metadata["token_usage_output_tokens"] == 5 + assert float(metadata["token_usage_cost"]) == pytest.approx(0.00042) + + async def test_content_policy_exception_returns_error_message(target, litellm_stub): exc = litellm_stub.exceptions.ContentPolicyViolationError("content_filter triggered") litellm_stub.acompletion = AsyncMock(side_effect=exc) 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 2bc2bf76dc..b4d2450f55 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -7,6 +7,7 @@ import os from collections.abc import MutableSequence from tempfile import NamedTemporaryFile +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -55,6 +56,17 @@ def create_mock_completion(content: str = "hi", finish_reason: str = "stop"): return mock_completion +def _mock_usage(*, prompt_tokens: int, completion_tokens: int, total_tokens: int) -> SimpleNamespace: + """Build a Chat Completions ``usage`` stand-in with no nested detail breakdowns.""" + return SimpleNamespace( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + prompt_tokens_details=None, + completion_tokens_details=None, + ) + + @pytest.fixture def sample_conversations() -> MutableSequence[MessagePiece]: conversations = get_sample_conversations() @@ -541,6 +553,71 @@ async def test_send_prompt_async_content_filter_200(target: OpenAIChatTarget): assert response[0].message_pieces[0].structured_refusal is None +async def test_send_prompt_async_captures_finish_reason(target: OpenAIChatTarget): + message = Message(message_pieces=[MessagePiece(role="user", conversation_id="c", original_value="hi")]) + mock_completion = create_mock_completion(content="hello", finish_reason="stop") + target._async_client.chat.completions.create = AsyncMock( # type: ignore[method-assign] + return_value=mock_completion + ) + + response = await target.send_prompt_async(message=message) + + assert response[0].message_pieces[0].prompt_metadata["finish_reason"] == "stop" + + +async def test_content_filter_200_captures_token_usage_and_finish_reason(target: OpenAIChatTarget): + """A blocked response still reports what it consumed, which is exactly where it matters most.""" + message = Message(message_pieces=[MessagePiece(role="user", conversation_id="c", original_value="harmful")]) + mock_completion = create_mock_completion(content="partial", finish_reason="content_filter") + mock_completion.usage = _mock_usage(prompt_tokens=12, completion_tokens=5, total_tokens=17) + target._async_client.chat.completions.create = AsyncMock( # type: ignore[method-assign] + return_value=mock_completion + ) + + response = await target.send_prompt_async(message=message) + + piece = response[0].message_pieces[0] + assert piece.response_error == "blocked" + assert piece.prompt_metadata["finish_reason"] == "content_filter" + assert piece.prompt_metadata["token_usage_input_tokens"] == 12 + assert piece.prompt_metadata["token_usage_output_tokens"] == 5 + assert piece.prompt_metadata["token_usage_total_tokens"] == 17 + + +async def test_truncated_empty_response_captures_finish_reason(target: OpenAIChatTarget): + """The graceful empty-truncated fallback must carry metadata too.""" + message = Message(message_pieces=[MessagePiece(role="user", conversation_id="c", original_value="hi")]) + mock_completion = create_mock_completion(content=None, finish_reason="length") + mock_completion.usage = _mock_usage(prompt_tokens=9, completion_tokens=100, total_tokens=109) + target._async_client.chat.completions.create = AsyncMock( # type: ignore[method-assign] + return_value=mock_completion + ) + + response = await target.send_prompt_async(message=message) + + piece = response[0].message_pieces[0] + assert piece.is_truncated is True + assert piece.prompt_metadata["finish_reason"] == "length" + assert piece.prompt_metadata["token_usage_output_tokens"] == 100 + + +async def test_sdk_content_filter_error_omits_finish_reason( + target: OpenAIChatTarget, sample_conversations: MutableSequence[MessagePiece] +): + """The SDK-raised path has no response object to read, so no metadata is invented.""" + message_piece = sample_conversations[0] + message_piece.conversation_id = "test-conv-id" + request = Message(message_pieces=[message_piece]) + + with patch.object(target._async_client.chat.completions, "create", new_callable=AsyncMock) as mock_create: + mock_create.side_effect = ContentFilterFinishReasonError() + response = await target.send_prompt_async(message=request) + + piece = response[0].message_pieces[0] + assert piece.response_error == "blocked" + assert "finish_reason" not in piece.prompt_metadata + + async def test_send_prompt_async_structured_refusal(target: OpenAIChatTarget): refusal = "I cannot assist with that request." completion = ChatCompletion( 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 defa9987d7..106ee4482b 100644 --- a/tests/unit/prompt_target/target/test_openai_response_target.py +++ b/tests/unit/prompt_target/target/test_openai_response_target.py @@ -1556,6 +1556,75 @@ async def test_construct_message_truncated_captures_token_usage( assert piece.prompt_metadata["token_usage_reasoning_tokens"] == 7 +async def test_construct_message_captures_completed_status( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """``status`` is the Responses equivalent of Chat Completions' ``finish_reason``.""" + response = MagicMock() + response.status = "completed" + response.incomplete_details = None + response.output = [_make_message_section("Answer")] + response.usage = None + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + metadata = result.message_pieces[0].prompt_metadata + assert metadata["status"] == "completed" + assert "incomplete_reason" not in metadata + + +async def test_construct_message_truncated_captures_status_and_incomplete_reason( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + response = _make_truncated_response(output=[_make_message_section("Partial answer")]) + response.usage = None + + result = await target._construct_message_from_response_async(response, dummy_text_message_piece) + + metadata = result.message_pieces[0].prompt_metadata + assert metadata["status"] == "incomplete" + assert metadata["incomplete_reason"] == "max_output_tokens" + + +async def test_construct_message_records_status_on_primary_piece_not_reasoning( + target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece +): + """Status metadata must be written after the reasoning sort, like usage.""" + response = _make_truncated_response(output=[_make_reasoning_section(), _make_message_section("Partial answer")]) + response.usage = None + + 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.prompt_metadata["status"] == "incomplete" + assert result.message_pieces[-1].converted_value_data_type == "reasoning" + assert "status" not in result.message_pieces[-1].prompt_metadata + + +async def test_content_filter_captures_usage_status_and_incomplete_reason(target: OpenAIResponseTarget): + """A content-filtered response still reports what it consumed and why it stopped.""" + request = MessagePiece(role="user", conversation_id="c", original_value="harmful") + response = MagicMock() + response.error = None + response.status = "incomplete" + incomplete_details = MagicMock() + incomplete_details.reason = "content_filter" + response.incomplete_details = incomplete_details + response.output = [] + response.usage = _make_usage() + response.model_dump_json.return_value = "{}" + + message = target._handle_content_filter_response(response, request) + + piece = message.message_pieces[0] + assert piece.response_error == "blocked" + assert piece.prompt_metadata["status"] == "incomplete" + assert piece.prompt_metadata["incomplete_reason"] == "content_filter" + assert piece.prompt_metadata["token_usage_input_tokens"] == 11 + assert piece.prompt_metadata["token_usage_output_tokens"] == 22 + + 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/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index e179a78425..3de86263e7 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -275,7 +275,8 @@ async def test_response_preserves_metadata_after_history_squash(): response_piece = response_messages[0].message_pieces[0] assert response_piece.conversation_id == _LINEAGE_CONVERSATION_ID - assert response_piece.prompt_metadata == _LINEAGE_PROMPT_METADATA + # Lineage metadata survives alongside the metadata captured from the API response. + assert response_piece.prompt_metadata == {**_LINEAGE_PROMPT_METADATA, "finish_reason": "stop"} @pytest.mark.usefixtures("patch_central_database") From e739cf9b7f1029b17b2576ae117a4e704b67b669 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 12 Aug 2026 01:52:31 +0000 Subject: [PATCH 2/3] Address review feedback on reserved response metadata Take the reserved keys as explicit keyword parameters instead of an open mapping, so the clearing loop and the writing loop cannot drift apart: an unrecognized key is now a type error at the call site rather than a value that is silently persisted. Add the two requested inline comments in OpenAICompletionTarget explaining why the capture is per piece and why it borrows the Chat Completions usage parser. --- .../chat_completions_response_parser.py | 2 +- pyrit/prompt_target/common/utils.py | 25 +++++++++++---- .../openai/openai_completion_target.py | 4 ++- .../openai/openai_response_target.py | 2 +- .../test_interface_prompts.py | 4 +-- .../target/test_chat_completions_helpers.py | 31 ++++++++++++++++--- 6 files changed, 53 insertions(+), 15 deletions(-) diff --git a/pyrit/prompt_target/common/chat_completions_response_parser.py b/pyrit/prompt_target/common/chat_completions_response_parser.py index a1358ad206..b1a28a0f18 100644 --- a/pyrit/prompt_target/common/chat_completions_response_parser.py +++ b/pyrit/prompt_target/common/chat_completions_response_parser.py @@ -325,7 +325,7 @@ def capture_usage_and_finish_reason(*, pieces: list[MessagePiece], response: Any # 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}) + set_response_metadata(pieces=pieces, finish_reason=finish_reason) def token_usage_from_chat_completion(usage: Any) -> TokenUsage: diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index ac51efec28..13d7c12231 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/pyrit/prompt_target/common/utils.py @@ -3,7 +3,7 @@ import asyncio import logging -from collections.abc import Callable, Mapping +from collections.abc import Callable from typing import Any from pyrit.exceptions import PyritException @@ -99,7 +99,13 @@ def build_empty_truncated_response(*, request: MessagePiece) -> Message: RESERVED_RESPONSE_METADATA_KEYS: frozenset[str] = frozenset({"finish_reason", "status", "incomplete_reason"}) -def set_response_metadata(*, pieces: list[MessagePiece], values: Mapping[str, Any]) -> None: +def set_response_metadata( + *, + pieces: list[MessagePiece], + finish_reason: Any = None, + status: Any = None, + incomplete_reason: Any = None, +) -> None: """ Record provider-reported, response-level metadata on the first response piece. @@ -111,16 +117,23 @@ def set_response_metadata(*, pieces: list[MessagePiece], values: Mapping[str, An 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. + There is one keyword parameter per reserved key so the two cannot drift apart: a key the + clearing loop does not know about is a type error rather than a value that is silently + persisted. ``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. + 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. + finish_reason (Any): The stop reason a Chat Completions or Completions response reported. + status (Any): The status a Responses API response reported. + incomplete_reason (Any): The incomplete detail a Responses API response reported. """ if not pieces: return + values = {"finish_reason": finish_reason, "status": status, "incomplete_reason": incomplete_reason} + for piece in pieces: for reserved_key in RESERVED_RESPONSE_METADATA_KEYS: piece.prompt_metadata.pop(reserved_key, None) diff --git a/pyrit/prompt_target/openai/openai_completion_target.py b/pyrit/prompt_target/openai/openai_completion_target.py index 81fb6d5ca0..8d067da123 100644 --- a/pyrit/prompt_target/openai/openai_completion_target.py +++ b/pyrit/prompt_target/openai/openai_completion_target.py @@ -172,12 +172,14 @@ def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece the SDK raises on a content filter. pieces (list[MessagePiece]): The constructed response pieces. """ + # The Completions and Chat Completions APIs report the same ``usage`` schema, so the parser is shared. capture_token_usage(pieces=pieces, response=response) choices = getattr(response, "choices", None) or [] for index, piece in enumerate(pieces): + # Per piece, not per response: each piece is one choice, so each carries its own stop reason. choice = choices[index] if index < len(choices) else None - set_response_metadata(pieces=[piece], values={"finish_reason": getattr(choice, "finish_reason", None)}) + set_response_metadata(pieces=[piece], finish_reason=getattr(choice, "finish_reason", None)) async def _construct_message_from_response_async(self, response: Any, request: Any) -> Message: """ diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 14f34c6a0e..11f985223a 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -587,7 +587,7 @@ def _capture_response_metadata(self, *, response: Any, pieces: list[MessagePiece status = getattr(response, "status", None) incomplete_details = getattr(response, "incomplete_details", None) incomplete_reason = getattr(incomplete_details, "reason", None) if incomplete_details else None - set_response_metadata(pieces=pieces, values={"status": status, "incomplete_reason": incomplete_reason}) + set_response_metadata(pieces=pieces, status=status, incomplete_reason=incomplete_reason) def _validate_response(self, response: Response, request: MessagePiece) -> None: """ diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index 00b1ad0104..580529ab44 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -965,8 +965,8 @@ def test_get_message_pieces_captured_response_metadata(sqlite_instance: MemoryIn role="assistant", original_value="fine", ) - set_response_metadata(pieces=[matching], values={key: value}) - set_response_metadata(pieces=[other], values={key: "something_else"}) + set_response_metadata(pieces=[matching], **{key: value}) + set_response_metadata(pieces=[other], **{key: "something_else"}) sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=matching), PromptMemoryEntry(entry=other)]) retrieved = sqlite_instance.get_message_pieces(prompt_metadata={key: value}) 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 0fa6406c7c..0b9662072f 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -4,6 +4,7 @@ """Unit tests for the shared OpenAI Chat Completions wire-format helpers.""" import base64 +import inspect import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -273,7 +274,7 @@ def test_capture_usage_and_finish_reason_overwrites_caller_supplied_finish_reaso def test_set_response_metadata_ignores_unreported_values(value): """``prompt_metadata`` is JSON-serialized, so anything but a non-empty string is not reported.""" piece = _request_piece("ok") - set_response_metadata(pieces=[piece], values={"status": value}) + set_response_metadata(pieces=[piece], status=value) assert "status" not in piece.prompt_metadata @@ -283,7 +284,7 @@ def test_set_response_metadata_clears_every_reserved_key(reserved_key): piece = _request_piece("ok") piece.prompt_metadata[reserved_key] = "caller_supplied" - set_response_metadata(pieces=[piece], values={"finish_reason": "stop"}) + set_response_metadata(pieces=[piece], finish_reason="stop") assert piece.prompt_metadata.get(reserved_key) == ("stop" if reserved_key == "finish_reason" else None) @@ -292,7 +293,7 @@ def test_set_response_metadata_keeps_all_reported_values(): """Clearing runs once up front, so a second reported key must not wipe the first.""" piece = _request_piece("ok") - set_response_metadata(pieces=[piece], values={"status": "incomplete", "incomplete_reason": "max_output_tokens"}) + set_response_metadata(pieces=[piece], status="incomplete", incomplete_reason="max_output_tokens") assert piece.prompt_metadata["status"] == "incomplete" assert piece.prompt_metadata["incomplete_reason"] == "max_output_tokens" @@ -302,7 +303,7 @@ def test_set_response_metadata_leaves_unreserved_caller_metadata_untouched(): piece = _request_piece("ok") piece.prompt_metadata["video_id"] = "caller_supplied" - set_response_metadata(pieces=[piece], values={"finish_reason": "stop"}) + set_response_metadata(pieces=[piece], finish_reason="stop") assert piece.prompt_metadata["video_id"] == "caller_supplied" @@ -613,3 +614,25 @@ async def test_build_response_pieces_async_orders_text_audio_tool(): pieces = await build_response_pieces_async(response=resp, request=_request_piece(), audio_format="wav") assert [p.converted_value_data_type for p in pieces] == ["text", "text", "audio_path", "function_call"] + + +def test_set_response_metadata_records_every_reserved_key(): + """Every reserved key must be reachable: named in the signature and written back to the piece.""" + piece = _request_piece("ok") + reported = {key: f"reported_{key}" for key in RESERVED_RESPONSE_METADATA_KEYS} + parameters = inspect.signature(set_response_metadata).parameters + keyword_only = {name for name, p in parameters.items() if p.kind is inspect.Parameter.KEYWORD_ONLY} + + assert keyword_only - {"pieces"} == RESERVED_RESPONSE_METADATA_KEYS + + set_response_metadata(pieces=[piece], **reported) + + assert piece.prompt_metadata == reported + + +def test_set_response_metadata_rejects_an_unreserved_key(): + """An unreserved key is drift or a typo, so it fails at the call site instead of being persisted.""" + piece = _request_piece("ok") + + with pytest.raises(TypeError): + set_response_metadata(pieces=[piece], reasoning_status="not_reserved") From b39485e0a1bab026bf2bd0a5c40dc51b3f045600 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 12 Aug 2026 02:45:08 +0000 Subject: [PATCH 3/3] Scope the reserved-key comment to the paths that actually capture The module comment claimed every target clears the reserved keys, which overstates it: a path that never reaches a provider response captures nothing and so clears nothing. Name that boundary where the contract is written down. Also pin the unreserved-key test to the argument it is really about, and note in _capture_response_cost that it has to run after usage capture, which clears the whole token_usage_ prefix. --- pyrit/prompt_target/common/utils.py | 6 ++++-- pyrit/prompt_target/litellm_chat_target.py | 4 +++- .../prompt_target/target/test_chat_completions_helpers.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index 13d7c12231..a28b164252 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/pyrit/prompt_target/common/utils.py @@ -94,8 +94,10 @@ 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. +#: differently, so the union is reserved rather than any single key: a target that captures response +#: metadata clears all of them and writes back only what its own provider reported. A path that never +#: reaches a provider response, such as the one that handles an HTTP 400, captures nothing and so +#: clears nothing. RESERVED_RESPONSE_METADATA_KEYS: frozenset[str] = frozenset({"finish_reason", "status", "incomplete_reason"}) diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index f5d1767d6a..2a5c18a551 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -480,7 +480,9 @@ def _capture_response_cost(self, *, pieces: list[MessagePiece], response: Any) - model-aware and is unique to LiteLLM (the raw OpenAI SDK response carries no cost), so this mirrors ``capture_token_usage`` and writes ``token_usage_cost`` alongside the token counts. The value is stored as a string to honor the ``prompt_metadata`` value contract, and any - failure is swallowed so cost accounting never breaks the response path. + failure is swallowed so cost accounting never breaks the response path. This has to run + after usage capture, which clears the whole ``token_usage_`` prefix and would otherwise + drop the cost again. Args: pieces (list[MessagePiece]): The constructed response pieces. 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 0b9662072f..5c2c0ffb28 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -634,5 +634,5 @@ def test_set_response_metadata_rejects_an_unreserved_key(): """An unreserved key is drift or a typo, so it fails at the call site instead of being persisted.""" piece = _request_piece("ok") - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="reasoning_status"): set_response_metadata(pieces=[piece], reasoning_status="not_reserved")