MAINT: Consolidate OpenAI response handling

pyrit/prompt_target/openai/_response_adapter.py

@@ -0,0 +1,229 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import logging
+from typing import Any, Generic, TypeVar
+
+from openai.types.chat import ChatCompletion
+from openai.types.responses import Response, ResponseOutputText
+
+from pyrit.exceptions import EmptyResponseException, PyritException
+from pyrit.models import MessagePiece, TokenUsage, read_usage_int, read_usage_value
+from pyrit.prompt_target.common.chat_completions_response_parser import (
+    capture_token_usage,
+    capture_usage_and_finish_reason,
+    extract_partial_content,
+    get_finish_reason,
+    is_content_filter_response,
+    validate_chat_completion_response,
+)
+from pyrit.prompt_target.common.utils import (
+    set_response_metadata,
+    set_token_usage_metadata,
+    warn_truncated_response,
+)
+from pyrit.prompt_target.openai.openai_error_handling import _is_content_filter_error
+
+logger = logging.getLogger(__name__)
+
+ResponseT = TypeVar("ResponseT", contravariant=True)
+
+
+class OpenAIResponseAdapter(Generic[ResponseT]):
+    """Base response-format behavior used by ``OpenAITarget``."""
+
+    def is_content_filtered(self, *, response: ResponseT) -> bool:
+        """Return whether the response was blocked by a content filter."""
+        return False
+
+    def extract_partial_content(self, *, response: ResponseT) -> str | None:
+        """
+        Extract content emitted before a content filter stopped generation.
+
+        Args:
+            response (ResponseT): The provider response.
+
+        Returns:
+            str | None: Partial content, if available.
+        """
+        return None
+
+    def capture_metadata(self, *, response: ResponseT, pieces: list[MessagePiece]) -> None:
+        """Copy provider response metadata to PyRIT response pieces."""
+
+    def validate(self, *, response: ResponseT, is_truncated: bool) -> None:
+        """Validate the provider response."""
+
+    def is_truncated(self, *, response: ResponseT) -> bool:
+        """Return whether generation stopped at the output-token limit."""
+        return False
+
+
+class ChatCompletionsResponseAdapter(OpenAIResponseAdapter[ChatCompletion]):
+    """Response behavior for the OpenAI Chat Completions wire format."""
+
+    def is_content_filtered(self, *, response: ChatCompletion) -> bool:
+        """Return whether ``finish_reason`` reports a content filter."""
+        return is_content_filter_response(response)
+
+    def extract_partial_content(self, *, response: ChatCompletion) -> str | None:
+        """
+        Extract text emitted before a content filter stopped generation.
+
+        Args:
+            response (ChatCompletion): The provider response.
+
+        Returns:
+            str | None: Partial text, if present.
+        """
+        return extract_partial_content(response)
+
+    def capture_metadata(self, *, response: ChatCompletion, pieces: list[MessagePiece]) -> None:
+        """Capture token usage and the first choice's finish reason."""
+        capture_usage_and_finish_reason(pieces=pieces, response=response)
+
+    def validate(self, *, response: ChatCompletion, is_truncated: bool) -> None:
+        """Validate the response while accepting token-limit truncation."""
+        if is_truncated:
+            warn_truncated_response(signal="finish_reason='length'", limit_parameter="max_completion_tokens")
+            return
+        validate_chat_completion_response(response=response)
+
+    def is_truncated(self, *, response: ChatCompletion) -> bool:
+        """Return whether ``finish_reason`` reports token-limit truncation."""
+        return get_finish_reason(response=response) == "length"
+
+
+class CompletionsResponseAdapter(OpenAIResponseAdapter[Any]):
+    """Response behavior for the legacy OpenAI Completions wire format."""
+
+    def capture_metadata(self, *, response: Any, pieces: list[MessagePiece]) -> None:
+        """Capture call-level usage and each choice's finish reason."""
+        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], finish_reason=getattr(choice, "finish_reason", None))
+
+
+class ResponsesResponseAdapter(OpenAIResponseAdapter[Response]):
+    """Response behavior for the OpenAI Responses API wire format."""
+
+    def is_content_filtered(self, *, response: Response) -> bool:
+        """Return whether the response reports content filtering."""
+        error = getattr(response, "error", None)
+        if error is not None and _is_content_filter_error(response.model_dump()):
+            return True
+
+        if getattr(response, "status", None) != "incomplete":
+            return False
+        incomplete_details = getattr(response, "incomplete_details", None)
+        return incomplete_details is not None and incomplete_details.reason == "content_filter"
+
+    def extract_partial_content(self, *, response: Response) -> str | None:
+        """
+        Extract text from completed message sections in a filtered response.
+
+        Args:
+            response (Response): The provider response.
+
+        Returns:
+            str | None: Partial text, if present.
+        """
+        try:
+            parts = [
+                content_item.text
+                for section in response.output or []
+                if getattr(section, "type", None) == "message" and getattr(section, "status", None) == "completed"
+                for content_item in getattr(section, "content", None) or []
+                if isinstance(content_item, ResponseOutputText) and content_item.text
+            ]
+        except (AttributeError, IndexError, TypeError):
+            return None
+        return "\n".join(parts) if parts else None
+
+    def capture_metadata(self, *, response: Response, pieces: list[MessagePiece]) -> None:
+        """Capture token usage, response status, and incomplete reason."""
+        if not pieces:
+            return
+
+        usage = getattr(response, "usage", None)
+        parsed_usage = token_usage_from_responses(usage) if usage is not None else None
+        set_token_usage_metadata(pieces=pieces, usage=parsed_usage)
+
+        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, status=status, incomplete_reason=incomplete_reason)
+
+    def validate(self, *, response: Response, is_truncated: bool) -> None:
+        """
+        Validate the response while accepting token-limit truncation.
+
+        Args:
+            response (Response): The provider response.
+            is_truncated (bool): Whether the target classified the response as token-limit truncation.
+
+        Raises:
+            PyritException: If the provider reports an error or unexpected status.
+            EmptyResponseException: If a completed response has no output.
+        """
+        if response.error is not None and response.error.code != "content_filter":
+            raise PyritException(message=f"Response error: {response.error.code} - {response.error.message}")
+
+        if is_truncated:
+            warn_truncated_response(
+                signal="status='incomplete', reason='max_output_tokens'",
+                limit_parameter="max_output_tokens",
+            )
+            return
+
+        if response.status != "completed":
+            raise PyritException(message=f"Unexpected status: {response.status}")
+
+        if not response.output:
+            logger.error("The response returned no valid output.")
+            raise EmptyResponseException(message="The response returned an empty response.")
+
+    def is_truncated(self, *, response: Response) -> bool:
+        """Return whether the response stopped at ``max_output_tokens``."""
+        if response.status != "incomplete":
+            return False
+        incomplete_details = response.incomplete_details
+        reason = incomplete_details.reason if incomplete_details else None
+        return reason == "max_output_tokens"
+
+
+def token_usage_from_responses(usage: Any) -> TokenUsage:
+    """
+    Build a ``TokenUsage`` from a Responses API ``usage`` payload.
+
+    Args:
+        usage (Any): The Responses API usage object.
+
+    Returns:
+        TokenUsage: The parsed token usage.
+    """
+    input_details = read_usage_value(source=usage, name="input_tokens_details")
+    output_details = read_usage_value(source=usage, name="output_tokens_details")
+
+    input_tokens = read_usage_int(source=usage, name="input_tokens")
+    output_tokens = read_usage_int(source=usage, name="output_tokens")
+    total_tokens = read_usage_int(source=usage, name="total_tokens")
+    if total_tokens is None and input_tokens is not None and output_tokens is not None:
+        total_tokens = input_tokens + output_tokens
+
+    extra: dict[str, int] = {}
+    cache_write_tokens = read_usage_int(source=input_details, name="cache_write_tokens")
+    if cache_write_tokens is not None:
+        extra["cache_write_tokens"] = cache_write_tokens
+
+    return TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+        reasoning_tokens=read_usage_int(source=output_details, name="reasoning_tokens"),
+        cached_tokens=read_usage_int(source=input_details, name="cached_tokens"),
+        extra=extra,
+    )

pyrit/prompt_target/openai/openai_chat_target.py

@@ -27,13 +27,8 @@
 )
 from pyrit.prompt_target.common.chat_completions_response_parser import (
     build_response_pieces_async,
-    capture_usage_and_finish_reason,
     detect_response_content,
-    extract_partial_content,
-    get_finish_reason,
-    is_content_filter_response,
     save_audio_response_async,
-    validate_chat_completion_response,
 )
 from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.target_configuration import TargetConfiguration
@@ -42,8 +37,8 @@
     limit_requests_per_minute,
     validate_temperature,
     validate_top_p,
-    warn_truncated_response,
 )
+from pyrit.prompt_target.openai._response_adapter import ChatCompletionsResponseAdapter
 from pyrit.prompt_target.openai.openai_chat_audio_config import OpenAIChatAudioConfig
 from pyrit.prompt_target.openai.openai_target import OpenAITarget
 
@@ -93,6 +88,7 @@ class OpenAIChatTarget(OpenAITarget):
             ),
         )
     )
+    _response_adapter = ChatCompletionsResponseAdapter()
 
     @forward_init_parameters
     def __init__(
@@ -249,96 +245,6 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
         )
         return [response]
 
-    def _check_content_filter(self, response: Any) -> bool:
-        """
-        Check if a Chat Completions API response has finish_reason=content_filter.
-
-        Args:
-            response: A ChatCompletion object from the OpenAI SDK.
-
-        Returns:
-            True if content was filtered, False otherwise.
-        """
-        return is_content_filter_response(response)
-
-    def _extract_partial_content(self, response: Any) -> str | None:
-        """
-        Extract partial content from a Chat Completions response with finish_reason=content_filter.
-
-        When Azure Content Safety triggers mid-generation, the model may have produced partial
-        text in ``response.choices[0].message.content`` before being cut off.
-
-        Args:
-            response: A ChatCompletion object from the OpenAI SDK.
-
-        Returns:
-            The partial text content, or None if no content was generated.
-        """
-        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.
-
-        Checks for:
-        - Missing choices
-        - Invalid finish_reason
-        - At least one valid response type (text content, audio, or tool_calls)
-
-        A ``finish_reason == "length"`` (token-limit truncation) response is treated as valid, with a
-        warning, so that ``_construct_message_from_response_async`` can preserve any partial content
-        or fall back to a graceful empty response. Genuinely empty responses (no truncation) are
-        raised so the retry logic can attempt to get a complete response. Content filter responses
-        are handled separately by ``_check_content_filter``.
-
-        Args:
-            response: The ChatCompletion response from OpenAI SDK.
-            request: The original request MessagePiece.
-
-        Raises:
-            PyritException: For unexpected response structures or finish reasons.
-            EmptyResponseException: When the API returns an empty response that was not caused by
-                token-limit truncation.
-        """
-        # Token-limit truncation is handled before the shared validator, which would otherwise raise
-        # EmptyResponseException on a validly truncated but empty response. Reasoning models can spend
-        # the whole budget on hidden reasoning before emitting a visible answer, and a low limit may be
-        # deliberate, so warn instead of raising and let construction preserve any partial content or
-        # fall back to a graceful empty response.
-        if self._is_truncated_response(response):
-            warn_truncated_response(signal="finish_reason='length'", limit_parameter="max_completion_tokens")
-            return
-
-        # Genuinely empty responses (no truncation) raise so the retry logic can attempt to get a
-        # complete response.
-        validate_chat_completion_response(response=response)
-
-    def _is_truncated_response(self, response: ChatCompletion) -> bool:
-        """
-        Return True if the response was cut off by the token limit.
-
-        The Chat Completions API signals token-limit truncation via ``finish_reason == "length"``
-        on the first choice.
-
-        Args:
-            response: A ChatCompletion response from the OpenAI SDK.
-
-        Returns:
-            bool: True if the response was truncated at the token limit, False otherwise.
-        """
-        return get_finish_reason(response=response) == "length"
-
     def _detect_response_content(self, message: Any) -> tuple[bool, bool, bool]:
         """
         Detect what content types are present in a ChatCompletion message.