diff --git a/pyproject.toml b/pyproject.toml index bdc563d000..eecbddbb9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -136,6 +136,9 @@ opencv = [ speech = [ "azure-cognitiveservices-speech>=1.46.0", ] +litellm = [ + "litellm>=1.83.0,<2.0.0", +] # all includes all functional dependencies excluding the ones from the "dev" dependency group all = [ @@ -146,6 +149,7 @@ all = [ "flask>=3.1.3", "ipykernel>=6.29.5", "jupyter>=1.1.1", + "litellm>=1.83.0,<2.0.0", "ml-collections>=1.1.0", "ollama>=0.5.1", "opencv-python>=4.11.0.86", diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 1093e1da02..82e484c2cd 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -56,6 +56,7 @@ # These are re-exported from the seeds submodule from pyrit.models.storage_io import AzureBlobStorageIO, DiskStorageIO, StorageIO from pyrit.models.strategy_result import StrategyResult, StrategyResultT +from pyrit.models.token_usage import TokenUsage __all__ = [ "ALLOWED_CHAT_MESSAGE_ROLES", @@ -114,6 +115,7 @@ "StrategyResult", "StrategyResultT", "TextDataTypeSerializer", + "TokenUsage", "UnvalidatedScore", "VideoPathDataTypeSerializer", "RetryEvent", diff --git a/pyrit/models/token_usage.py b/pyrit/models/token_usage.py new file mode 100644 index 0000000000..32bc39ce1e --- /dev/null +++ b/pyrit/models/token_usage.py @@ -0,0 +1,114 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +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_" + +#: 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 +#: currency amount stored as a string and is filtered out by the int guard regardless. +_CORE_SUFFIXES = frozenset({"input_tokens", "output_tokens", "total_tokens", "reasoning_tokens", "cached_tokens"}) + + +@dataclass(frozen=True) +class TokenUsage: + """ + Provider-neutral token accounting for a single model call. + + Field names use the ``input``/``output`` vocabulary (aligned with the OpenAI Responses API, + Anthropic, and Gemini) rather than the Chat Completions ``prompt``/``completion`` terms. The + object is persisted onto a ``MessagePiece``'s ``prompt_metadata`` via ``to_metadata`` using + matching ``token_usage_input_tokens`` / ``token_usage_output_tokens`` key names (one consistent + vocabulary end to end). ``reasoning_tokens`` and ``cached_tokens`` are the two widely-available + sub-breakdowns promoted to fields; any other provider-specific counts (audio, predicted-output, + cache-write) ride along in ``extra``. + + This is a pure value object: it holds counts and (de)serializes them to metadata. Turning a + provider ``usage`` payload into a ``TokenUsage`` is the responsibility of the target/parser that + knows which wire format it received (for example, the Chat Completions parser in + ``pyrit.prompt_target.common.chat_completions_response_parser``). + + Neither cost nor the responding model name is modeled here: cost is a currency amount (tracked + separately under ``token_usage_cost``) and the model identity is already recorded on the + target's identifier. Both would be a category error inside a token-count value object. + + Only fields the provider actually reports are populated; absent values stay None (and are + omitted from ``to_metadata``) rather than being coerced to a misleading zero. + """ + + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + reasoning_tokens: int | None = None + cached_tokens: int | None = None + extra: dict[str, int] = field(default_factory=dict) + + @classmethod + def from_metadata(cls, metadata: dict[str, Any]) -> TokenUsage | None: + """ + Reconstruct a ``TokenUsage`` from a ``MessagePiece``'s ``prompt_metadata``. + + Reads the ``token_usage_input_tokens`` / ``token_usage_output_tokens`` keys written by + ``to_metadata``. Non-core integer ``token_usage_*`` keys are collected into ``extra``; + the string ``token_usage_cost`` key is ignored (cost is tracked separately). + + Args: + metadata (dict[str, Any]): The prompt metadata to read. + + Returns: + 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) + } + if not stripped: + return None + + def _pick(suffix: str) -> int | None: + value = stripped.get(suffix) + return value if isinstance(value, int) and not isinstance(value, bool) else None + + extra = { + key: value + for key, value in stripped.items() + if key not in _CORE_SUFFIXES and isinstance(value, int) and not isinstance(value, bool) + } + return cls( + input_tokens=_pick("input_tokens"), + output_tokens=_pick("output_tokens"), + total_tokens=_pick("total_tokens"), + reasoning_tokens=_pick("reasoning_tokens"), + cached_tokens=_pick("cached_tokens"), + extra=extra, + ) + + def to_metadata(self) -> dict[str, int]: + """ + Serialize to flat ``token_usage_*`` metadata keys, omitting fields that are None. + + Uses the ``input``/``output`` vocabulary for the key names to match the field names (one + consistent naming end to end). ``extra`` counts are written verbatim under the + ``token_usage_`` prefix. + + Returns: + dict[str, int]: The metadata fragment to merge into ``prompt_metadata``. + """ + out: dict[str, int] = {} + if self.input_tokens is not None: + out[_METADATA_PREFIX + "input_tokens"] = self.input_tokens + if self.output_tokens is not None: + out[_METADATA_PREFIX + "output_tokens"] = self.output_tokens + if self.total_tokens is not None: + out[_METADATA_PREFIX + "total_tokens"] = self.total_tokens + if self.reasoning_tokens is not None: + out[_METADATA_PREFIX + "reasoning_tokens"] = self.reasoning_tokens + if self.cached_tokens is not None: + out[_METADATA_PREFIX + "cached_tokens"] = self.cached_tokens + for name, value in self.extra.items(): + out[_METADATA_PREFIX + name] = value + return out diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 82f897c156..55bbb5aaad 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -36,6 +36,7 @@ ) from pyrit.prompt_target.http_target.httpx_api_target import HTTPXAPITarget from pyrit.prompt_target.hugging_face.hugging_face_endpoint_target import HuggingFaceEndpointTarget +from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget from pyrit.prompt_target.openai.openai_chat_audio_config import OpenAIChatAudioConfig from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget from pyrit.prompt_target.openai.openai_completion_target import OpenAICompletionTarget @@ -87,6 +88,7 @@ def __getattr__(name: str) -> object: "HuggingFaceChatTarget", "HuggingFaceEndpointTarget", "limit_requests_per_minute", + "LiteLLMChatTarget", "OpenAICompletionTarget", "OpenAIChatAudioConfig", "OpenAIChatTarget", diff --git a/pyrit/prompt_target/common/chat_completions_message_builder.py b/pyrit/prompt_target/common/chat_completions_message_builder.py new file mode 100644 index 0000000000..149cc1023f --- /dev/null +++ b/pyrit/prompt_target/common/chat_completions_message_builder.py @@ -0,0 +1,267 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Helpers for building request payloads in the OpenAI *Chat Completions* wire format. + +This format (a ``messages`` list of ``{"role", "content"}`` dicts, where ``content`` is +either a string or a list of typed content parts) is an industry standard implemented by +the OpenAI SDK and by LiteLLM's ``acompletion``. Keeping these builders here lets every +target that speaks that format (``OpenAIChatTarget``, ``LiteLLMChatTarget``, ...) share one +implementation instead of re-inventing message construction. +""" + +from collections.abc import MutableSequence +from typing import Any + +from pyrit.common.data_url_converter import convert_local_image_to_data_url +from pyrit.models import ( + ChatMessage, + DataTypeSerializer, + Message, + MessagePiece, + data_serializer_factory, +) +from pyrit.models.json_response_config import _JsonResponseConfig + +# Data types that render as a plain text content part. +_TEXT_DATA_TYPES = ("text", "error") + +# Audio formats accepted by the OpenAI Chat Completions ``input_audio`` content part. +# OpenAI SDK: openai/types/chat/chat_completion_content_part_input_audio_param.py +# defines format: Required[Literal["wav", "mp3"]]. +_INPUT_AUDIO_EXTENSIONS = (".wav", ".mp3") + + +def is_text_only_conversation(conversation: MutableSequence[Message]) -> bool: + """ + Return whether every message in the conversation is a single text (or error) piece. + + When true, the simpler text-only message format can be used, which is more broadly + compatible with OpenAI-compatible providers that don't accept multipart content. + + Args: + conversation (MutableSequence[Message]): The conversation to inspect. + + Returns: + bool: True if all messages are a single text/error piece, False otherwise. + """ + for turn in conversation: + if len(turn.message_pieces) != 1: + return False + if turn.message_pieces[0].converted_value_data_type not in _TEXT_DATA_TYPES: + return False + return True + + +def build_text_chat_messages(conversation: MutableSequence[Message]) -> list[dict[str, Any]]: + """ + Build chat messages using the simple ``{"role", "content": str}`` format. + + Many OpenAI-"compatible" providers don't support the multipart content format, so the + plain-text form is used whenever the conversation is text-only. + + Args: + conversation (MutableSequence[Message]): The conversation to convert. Each message + must have exactly one text or error piece. + + Returns: + list[dict[str, Any]]: The list of constructed chat messages. + + Raises: + ValueError: If any message does not have exactly one text/error piece. + """ + chat_messages: list[dict[str, Any]] = [] + for message in conversation: + if len(message.message_pieces) != 1: + raise ValueError("build_text_chat_messages only supports a single message piece per message.") + + message_piece = message.message_pieces[0] + if message_piece.converted_value_data_type not in _TEXT_DATA_TYPES: + raise ValueError( + f"build_text_chat_messages only supports text and error data types." + f" Received: {message_piece.converted_value_data_type}." + ) + + chat_message = ChatMessage(role=message_piece.api_role, content=message_piece.converted_value) + chat_messages.append(chat_message.model_dump(exclude_none=True)) + + return chat_messages + + +def build_text_content_entry(*, message_piece: MessagePiece) -> dict[str, Any]: + """ + Build a text content part for a multipart chat message. + + Args: + message_piece (MessagePiece): The text/error piece. + + Returns: + dict[str, Any]: A ``{"type": "text", "text": ...}`` content part. + """ + return {"type": "text", "text": message_piece.converted_value} + + +async def build_image_content_entry_async(*, message_piece: MessagePiece) -> dict[str, Any]: + """ + Build an image content part (as a base64 data URL) for a multipart chat message. + + Args: + message_piece (MessagePiece): The ``image_path`` piece. + + Returns: + dict[str, Any]: A ``{"type": "image_url", "image_url": {"url": ...}}`` content part. + """ + data_url = await convert_local_image_to_data_url(message_piece.converted_value) + return {"type": "image_url", "image_url": {"url": data_url}} + + +async def build_audio_content_entry_async(*, message_piece: MessagePiece) -> dict[str, Any]: + """ + Build an ``input_audio`` content part (base64-encoded) for a multipart chat message. + + Args: + message_piece (MessagePiece): The ``audio_path`` piece. + + Returns: + dict[str, Any]: A ``{"type": "input_audio", "input_audio": {"data", "format"}}`` part. + + Raises: + ValueError: If the audio file extension is not ``.wav`` or ``.mp3``. + """ + ext = DataTypeSerializer.get_extension(message_piece.converted_value) + if not ext or ext.lower() not in _INPUT_AUDIO_EXTENSIONS: + raise ValueError( + f"Unsupported audio format: {ext}. " + "OpenAI Chat Completions API input_audio only supports .wav and .mp3. " + "Note: This is different from the Whisper Speech-to-Text API which supports more formats." + ) + audio_serializer = data_serializer_factory( + category="prompt-memory-entries", + value=message_piece.converted_value, + data_type="audio_path", + extension=ext, + ) + base64_data = await audio_serializer.read_data_base64() + return {"type": "input_audio", "input_audio": {"data": base64_data, "format": ext.lower().lstrip(".")}} + + +def should_skip_audio_piece( + *, + message_piece: MessagePiece, + is_last_message: bool, + has_text_piece: bool, + prefer_transcript_for_history: bool, +) -> bool: + """ + Determine whether an ``audio_path`` piece should be omitted when building chat messages. + + Assistant audio is always skipped (Chat Completions only accepts audio in user messages; + the assistant transcript text piece carries the content). Historical user audio is skipped + when ``prefer_transcript_for_history`` is set and a transcript text piece is present. + + Args: + message_piece (MessagePiece): The piece to evaluate. + is_last_message (bool): Whether this is the last (current) message in the conversation. + has_text_piece (bool): Whether the message also contains a text (transcript) piece. + prefer_transcript_for_history (bool): Whether to drop historical user audio in favor of + its transcript. + + Returns: + bool: True if the audio should be skipped, False otherwise. + """ + if message_piece.converted_value_data_type != "audio_path": + return False + + if message_piece.api_role == "assistant": + return True + + return bool( + message_piece.api_role == "user" and not is_last_message and has_text_piece and prefer_transcript_for_history + ) + + +async def build_multimodal_chat_messages_async( + conversation: MutableSequence[Message], + *, + prefer_transcript_for_history: bool = False, +) -> list[dict[str, Any]]: + """ + Build chat messages using the multipart ``content`` format (text, image, and audio parts). + + Args: + conversation (MutableSequence[Message]): The conversation to convert. + prefer_transcript_for_history (bool): Whether to drop historical user audio in favor of + its transcript (see ``should_skip_audio_piece``). + + Returns: + list[dict[str, Any]]: The list of constructed chat messages. + + Raises: + ValueError: If a message has no determinable role, or contains an unsupported data type. + """ + chat_messages: list[dict[str, Any]] = [] + last_message_index = len(conversation) - 1 + + for message_index, message in enumerate(conversation): + message_pieces = message.message_pieces + is_last_message = message_index == last_message_index + has_text_piece = any(mp.converted_value_data_type == "text" for mp in message_pieces) + + content: list[dict[str, Any]] = [] + role = None + for message_piece in message_pieces: + role = message_piece.api_role + + if should_skip_audio_piece( + message_piece=message_piece, + is_last_message=is_last_message, + has_text_piece=has_text_piece, + prefer_transcript_for_history=prefer_transcript_for_history, + ): + continue + + data_type = message_piece.converted_value_data_type + if data_type in _TEXT_DATA_TYPES: + content.append(build_text_content_entry(message_piece=message_piece)) + elif data_type == "image_path": + content.append(await build_image_content_entry_async(message_piece=message_piece)) + elif data_type == "audio_path": + content.append(await build_audio_content_entry_async(message_piece=message_piece)) + else: + raise ValueError(f"Multimodal data type {data_type} is not yet supported.") + + if not role: + raise ValueError("No role could be determined from the message pieces.") + + chat_messages.append(ChatMessage(role=role, content=content).model_dump(exclude_none=True)) + + return chat_messages + + +def build_response_format(*, json_config: _JsonResponseConfig) -> dict[str, Any] | None: + """ + Build the ``response_format`` request parameter from a JSON response config. + + Args: + json_config (_JsonResponseConfig): The JSON response configuration derived from the + request metadata. + + Returns: + dict[str, Any] | None: A ``json_schema`` or ``json_object`` response-format dict, or + None when JSON output was not requested. + """ + if not json_config.enabled: + return None + + if json_config.schema: + return { + "type": "json_schema", + "json_schema": { + "name": json_config.schema_name, + "schema": json_config.schema, + "strict": json_config.strict, + }, + } + + return {"type": "json_object"} diff --git a/pyrit/prompt_target/common/chat_completions_response_parser.py b/pyrit/prompt_target/common/chat_completions_response_parser.py new file mode 100644 index 0000000000..c539b83ba5 --- /dev/null +++ b/pyrit/prompt_target/common/chat_completions_response_parser.py @@ -0,0 +1,472 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Helpers for parsing responses in the OpenAI *Chat Completions* wire format. + +Both the OpenAI SDK and LiteLLM return responses with the same shape +(``response.choices[0].message`` with ``content`` / ``tool_calls``, plus ``response.usage``), +so validation, piece construction, token-usage capture, and content-filter handling can be +shared across every target that speaks that format. +""" + +import base64 +import json +import logging +from collections.abc import Mapping +from typing import Any + +from pyrit.exceptions import ( + EmptyResponseException, + PyritException, + handle_bad_request_exception, +) +from pyrit.models import ( + Message, + MessagePiece, + TokenUsage, + construct_response_from_request, + data_serializer_factory, +) + +logger = logging.getLogger(__name__) + +# Finish reasons that represent a well-formed (non-error) completion. ``content_filter`` is +# included because it is handled separately (see ``is_content_filter_response``) before +# validation; a target should check for content filtering first. +DEFAULT_VALID_FINISH_REASONS: frozenset[str] = frozenset({"stop", "length", "content_filter", "tool_calls"}) + + +def detect_response_content(message: Any) -> tuple[bool, bool, bool]: + """ + Detect which content types are present in a Chat Completions ``message``. + + Args: + message (Any): The ``response.choices[0].message`` object. + + Returns: + tuple[bool, bool, bool]: ``(has_content, has_audio, has_tool_calls)``. + """ + has_content = bool(getattr(message, "content", None)) + has_audio = getattr(message, "audio", None) is not None + has_tool_calls = bool(getattr(message, "tool_calls", None)) + return has_content, has_audio, has_tool_calls + + +def validate_chat_completion_response( + *, + response: Any, + valid_finish_reasons: frozenset[str] = DEFAULT_VALID_FINISH_REASONS, +) -> None: + """ + Validate an OpenAI-compatible Chat Completions response. + + Checks for missing choices, an unexpected ``finish_reason``, and at least one usable + response type (text content, audio, or tool calls). Content-filter responses should be + handled by the caller *before* calling this (via ``is_content_filter_response``). + + Args: + response (Any): The Chat Completions response object. + valid_finish_reasons (frozenset[str]): Accepted ``finish_reason`` values. + + Raises: + PyritException: If there are no choices or the ``finish_reason`` is unexpected. + EmptyResponseException: If the response has no content, audio, or tool calls. + """ + if not getattr(response, "choices", None): + raise PyritException(message="No choices returned in the completion response.") + + choice = response.choices[0] + finish_reason = getattr(choice, "finish_reason", None) + if finish_reason not in valid_finish_reasons: + detail = response.model_dump_json() if hasattr(response, "model_dump_json") else str(response) + raise PyritException(message=f"Unknown finish_reason {finish_reason} from response: {detail}") + + has_content, has_audio, has_tool_calls = detect_response_content(getattr(choice, "message", None)) + if not (has_content or has_audio or has_tool_calls): + logger.error("The chat returned an empty response (no content, audio, or tool_calls).") + raise EmptyResponseException(message="The chat returned an empty response (no content, audio, or tool_calls).") + + +def _build_text_piece(*, content: str, request: MessagePiece) -> MessagePiece: + """ + Build a single text response piece. + + Args: + content (str): The text content. + request (MessagePiece): The originating request piece. + + Returns: + MessagePiece: The constructed text piece. + """ + return construct_response_from_request( + request=request, + response_text_pieces=[content], + response_type="text", + ).message_pieces[0] + + +def _build_tool_pieces(*, message: Any, request: MessagePiece) -> list[MessagePiece]: + """ + Build function_call response pieces from a message's ``tool_calls``. + + Args: + message (Any): The ``response.choices[0].message`` object. + request (MessagePiece): The originating request piece. + + Returns: + list[MessagePiece]: The constructed function_call pieces (may be empty). + """ + pieces: list[MessagePiece] = [] + tool_calls = getattr(message, "tool_calls", None) + if not tool_calls: + return pieces + + for tool_call in tool_calls: + tool_call_data = { + "type": "function", + "id": tool_call.id, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments, + }, + } + pieces.append( + construct_response_from_request( + request=request, + response_text_pieces=[json.dumps(tool_call_data)], + response_type="function_call", + ).message_pieces[0] + ) + return pieces + + +def build_text_and_tool_pieces(*, response: Any, request: MessagePiece) -> list[MessagePiece]: + """ + Build response pieces for text content and tool/function calls. + + Audio and other modality-specific parsing is intentionally left to the caller so that + provider-specific handling (e.g. saving audio files) stays out of this shared helper. + + Args: + response (Any): The Chat Completions response object. + request (MessagePiece): The originating request piece (for lineage propagation). + + Returns: + list[MessagePiece]: Text and function_call pieces (may be empty). + """ + message = response.choices[0].message + pieces: list[MessagePiece] = [] + + content = getattr(message, "content", None) + if content: + pieces.append(_build_text_piece(content=content, request=request)) + + pieces.extend(_build_tool_pieces(message=message, request=request)) + return pieces + + +async def save_audio_response_async(*, audio_data_base64: str, audio_format: str = "wav") -> str: + """ + Decode and persist base64 audio from a Chat Completions response to a file. + + Args: + audio_data_base64 (str): Base64-encoded audio data from ``message.audio.data``. + audio_format (str): The audio format (e.g. ``"wav"``, ``"mp3"``, ``"pcm16"``). Raw + ``pcm16`` is wrapped in a WAV container (24kHz mono). + + Returns: + str: The file path where the audio was saved. + """ + audio_bytes = base64.b64decode(audio_data_base64) + extension = f".{audio_format}" if audio_format != "pcm16" else ".wav" + + audio_serializer = data_serializer_factory( + category="prompt-memory-entries", + data_type="audio_path", + extension=extension, + ) + + if audio_format == "pcm16": + # Raw PCM needs WAV headers - OpenAI uses 24kHz mono PCM16 + await audio_serializer.save_formatted_audio( + data=audio_bytes, + num_channels=1, + sample_width=2, + sample_rate=24000, + ) + else: + await audio_serializer.save_data(audio_bytes) + + return audio_serializer.value + + +async def build_audio_pieces_async( + *, message: Any, request: MessagePiece, audio_format: str = "wav" +) -> list[MessagePiece]: + """ + Build response pieces for an audio message: a transcript text piece and a saved audio file. + + Args: + message (Any): The ``response.choices[0].message`` object. + request (MessagePiece): The originating request piece. + audio_format (str): The audio format used to persist the audio data. + + Returns: + list[MessagePiece]: The transcript and/or ``audio_path`` pieces (may be empty). + """ + pieces: list[MessagePiece] = [] + audio_response = getattr(message, "audio", None) + if audio_response is None: + return pieces + + transcript = getattr(audio_response, "transcript", None) + if transcript: + pieces.append( + construct_response_from_request( + request=request, + response_text_pieces=[transcript], + response_type="text", + prompt_metadata={"transcription": "audio"}, + ).message_pieces[0] + ) + + audio_data = getattr(audio_response, "data", None) + if audio_data: + audio_path = await save_audio_response_async(audio_data_base64=audio_data, audio_format=audio_format) + pieces.append( + construct_response_from_request( + request=request, + response_text_pieces=[audio_path], + response_type="audio_path", + ).message_pieces[0] + ) + + return pieces + + +async def build_response_pieces_async( + *, response: Any, request: MessagePiece, audio_format: str = "wav" +) -> list[MessagePiece]: + """ + Build all response pieces (text, audio, and tool calls) from a Chat Completions response. + + Pieces are ordered text, then audio (transcript + file), then tool calls. + + Args: + response (Any): The Chat Completions response object. + request (MessagePiece): The originating request piece. + audio_format (str): The audio format used to persist any audio data. + + Returns: + list[MessagePiece]: All constructed response pieces (may be empty). + """ + message = response.choices[0].message + pieces: list[MessagePiece] = [] + + content = getattr(message, "content", None) + if content: + pieces.append(_build_text_piece(content=content, request=request)) + + pieces.extend(await build_audio_pieces_async(message=message, request=request, audio_format=audio_format)) + pieces.extend(_build_tool_pieces(message=message, request=request)) + return pieces + + +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. + + 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: + return + + token_usage = token_usage_from_chat_completion(usage) + pieces[0].prompt_metadata.update(token_usage.to_metadata()) + + +def _read(source: Any, name: str) -> Any: + """ + Read ``name`` from ``source``, which may be a mapping or an attribute object. + + Args: + source (Any): The usage object (may be None). + name (str): The field name to read. + + Returns: + Any: The field value, or None when absent. + """ + if isinstance(source, Mapping): + return source.get(name) + return getattr(source, name, None) + + +def _usage_field(source: Any, *names: str) -> int | None: + """ + Return the first int-valued field among ``names`` on ``source``, else None. + + ``source`` may be either a mapping (for example, a ``model_dump``'d usage payload) or an + attribute object (the OpenAI/LiteLLM SDK ``Usage`` type), so both access styles are supported. + Booleans are rejected even though ``bool`` is a subclass of ``int``. + + Args: + source (Any): The usage object or nested details object (may be None). + names (str): Candidate field names, tried in order. + + Returns: + int | None: The first integer value found, or None. + """ + for name in names: + value = _read(source, name) + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + +def token_usage_from_chat_completion(usage: Any) -> TokenUsage: + """ + Build a ``TokenUsage`` from a Chat Completions ``usage`` payload (OpenAI or LiteLLM). + + Reads the top-level ``prompt_tokens`` / ``completion_tokens`` / ``total_tokens`` counts and the + nested ``prompt_tokens_details`` / ``completion_tokens_details`` breakdowns, deriving + ``total_tokens`` when the provider omits it. Non-OpenAI providers routed through LiteLLM + (for example, Anthropic) surface prompt-cache tokens at the top level rather than inside the + details object, so ``cache_read_input_tokens`` / ``cache_creation_input_tokens`` are picked up + as well. Unmodeled detail counts (audio, predicted-output) ride along in ``extra``. + + This parser is specific to the Chat Completions wire format. The Responses API reports usage + under different names (``input_tokens`` / ``output_tokens``); a target that speaks that format + should parse it in its own module rather than overloading this function. + + Args: + usage (Any): The Chat Completions usage object (attribute object or mapping). + + Returns: + TokenUsage: The parsed token usage. + """ + input_tokens = _usage_field(usage, "prompt_tokens") + output_tokens = _usage_field(usage, "completion_tokens") + total_tokens = _usage_field(usage, "total_tokens") + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + + prompt_details = _read(usage, "prompt_tokens_details") + completion_details = _read(usage, "completion_tokens_details") + + cached_tokens = _usage_field(prompt_details, "cached_tokens") + if cached_tokens is None: + cached_tokens = _usage_field(usage, "cache_read_input_tokens") + reasoning_tokens = _usage_field(completion_details, "reasoning_tokens") + + extra: dict[str, int] = {} + _add_extra(extra, "input_audio_tokens", _usage_field(prompt_details, "audio_tokens")) + _add_extra(extra, "cache_write_tokens", _usage_field(usage, "cache_creation_input_tokens")) + _add_extra(extra, "output_audio_tokens", _usage_field(completion_details, "audio_tokens")) + _add_extra(extra, "accepted_prediction_tokens", _usage_field(completion_details, "accepted_prediction_tokens")) + _add_extra(extra, "rejected_prediction_tokens", _usage_field(completion_details, "rejected_prediction_tokens")) + + return TokenUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + reasoning_tokens=reasoning_tokens, + cached_tokens=cached_tokens, + extra=extra, + ) + + +def _add_extra(target: dict[str, int], name: str, value: int | None) -> None: + """ + Insert ``name``->``value`` into ``target`` only when ``value`` is not None. + + Args: + target (dict[str, int]): The destination mapping. + name (str): The key to set. + value (int | None): The value to store, ignored when None. + """ + if value is not None: + target[name] = value + + +def is_content_filter_response(response: Any) -> bool: + """ + Return whether a Chat Completions response was blocked by a content filter. + + Args: + response (Any): The Chat Completions response object. + + Returns: + bool: True if ``finish_reason == "content_filter"``, False otherwise. + """ + try: + return bool(response.choices) and response.choices[0].finish_reason == "content_filter" + except (AttributeError, IndexError): + return False + + +def extract_partial_content(response: Any) -> str | None: + """ + Extract any partial text the model produced before a content filter triggered. + + Args: + response (Any): The Chat Completions response object. + + Returns: + str | None: The partial text, or None if none was generated. + """ + try: + choice = response.choices[0] + if choice.message and choice.message.content: + return choice.message.content + except (AttributeError, IndexError): + pass + return None + + +def build_content_filter_message( + *, + response: Any, + request: MessagePiece, + partial_content: str | None = None, +) -> Message: + """ + Build an ``error``-type Message for a content-filtered response. + + Rather than raising, blocked responses are surfaced as an error Message so attacks can + continue. When ``partial_content`` is available it is attached to each piece as + ``prompt_metadata["partial_content"]`` so scorers with ``score_blocked_content=True`` can + still evaluate what the model produced. + + Args: + response (Any): The Chat Completions response object (or an object exposing + ``model_dump_json``) describing the block. + request (MessagePiece): The originating request piece. + partial_content (str | None): Any partial model output recovered before the block. + + Returns: + Message: The constructed error Message with ``error="blocked"``. + """ + logger.warning("Output content filtered by content policy.") + + response_text = response.model_dump_json() if hasattr(response, "model_dump_json") else str(response) + error_message = handle_bad_request_exception( + response_text=response_text, + request=request, + error_code=200, + is_content_filter=True, + ) + + if partial_content: + for piece in error_message.message_pieces: + piece.prompt_metadata["partial_content"] = partial_content + + return error_message diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py new file mode 100644 index 0000000000..5294caba05 --- /dev/null +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -0,0 +1,583 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import itertools +import logging +import os +from collections.abc import Awaitable, Callable, MutableSequence +from typing import Any, NoReturn, cast + +from pyrit.auth import ensure_async_token_provider +from pyrit.exceptions import ( + EmptyResponseException, + PyritException, + RateLimitException, + get_retry_max_num_attempts, + handle_bad_request_exception, +) +from pyrit.identifiers import ComponentIdentifier +from pyrit.models import ( + Message, + MessagePiece, + PromptDataType, +) +from pyrit.prompt_target.common.chat_completions_message_builder import ( + build_multimodal_chat_messages_async, + build_response_format, + build_text_chat_messages, + is_text_only_conversation, +) +from pyrit.prompt_target.common.chat_completions_response_parser import ( + build_content_filter_message, + build_response_pieces_async, + capture_token_usage, + extract_partial_content, + is_content_filter_response, + validate_chat_completion_response, +) +from pyrit.prompt_target.common.prompt_target import PromptTarget +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, + validate_temperature, + validate_top_p, +) +from pyrit.prompt_target.openai.openai_chat_audio_config import OpenAIChatAudioConfig + +logger = logging.getLogger(__name__) + +# Conservative capability profile used when litellm metadata can't tell us more (unknown +# model, or litellm not importable at construction). Deliberately text-only with no JSON so +# we never advertise a capability we can't honor. +_TEXT_INPUT: frozenset[frozenset[PromptDataType]] = cast( + "frozenset[frozenset[PromptDataType]]", + frozenset({frozenset({"text"})}), +) + + +def _build_input_modalities(*, image: bool, audio: bool) -> frozenset[frozenset[PromptDataType]]: + """ + Build the set of supported input-modality combinations from capability flags. + + Always includes text. Enumerates every non-empty combination of the present modalities + (text, image, audio) so callers can advertise mixed-modality messages. + + Args: + image (bool): Whether image input is supported. + audio (bool): Whether audio input is supported. + + Returns: + frozenset[frozenset[PromptDataType]]: The supported input-modality combinations. + """ + present: list[PromptDataType] = ["text"] + if image: + present.append("image_path") + if audio: + present.append("audio_path") + + combos = [ + frozenset(combo) for size in range(1, len(present) + 1) for combo in itertools.combinations(present, size) + ] + return frozenset(combos) + + +def _build_output_modalities(*, audio: bool) -> frozenset[frozenset[PromptDataType]]: + """ + Build the set of supported output-modality combinations from capability flags. + + Args: + audio (bool): Whether audio output is supported. + + Returns: + frozenset[frozenset[PromptDataType]]: The supported output-modality combinations. + """ + output: list[frozenset[PromptDataType]] = [cast("frozenset[PromptDataType]", frozenset({"text"}))] + if audio: + output.append(cast("frozenset[PromptDataType]", frozenset({"audio_path"}))) + output.append(cast("frozenset[PromptDataType]", frozenset({"text", "audio_path"}))) + return frozenset(output) + + +class LiteLLMChatTarget(PromptTarget): + """ + Chat target that uses the LiteLLM SDK to access 100+ LLM providers. + + Unlike ``OpenAIChatTarget`` (which uses the OpenAI SDK directly), this target calls + ``litellm.acompletion()`` so it can route to any provider LiteLLM supports (Anthropic, + AWS Bedrock, Google Vertex, Cohere, etc.) without requiring a separate proxy server. + + LiteLLM speaks the OpenAI *Chat Completions* wire format, so this target shares its + request-building and response-parsing logic with ``OpenAIChatTarget`` via the helpers in + ``pyrit.prompt_target.common.chat_completions_message_builder`` and + ``pyrit.prompt_target.common.chat_completions_response_parser``. + + Install the optional dependency:: + + pip install pyrit[litellm] + + LiteLLM reads provider API keys from environment variables automatically + (e.g. ``ANTHROPIC_API_KEY``, ``AWS_ACCESS_KEY_ID``). You can also pass ``api_key`` + explicitly, or a callable/token-provider for Entra-style auth. + + Args: + model_name: LiteLLM model string (e.g. ``"anthropic/claude-sonnet-4-6"``, + ``"bedrock/anthropic.claude-v2"``, ``"vertex_ai/gemini-pro"``). Falls back to the + ``LITELLM_MODEL`` environment variable. + api_key: Optional API key, or a callable that returns an access token (sync or async). + When omitted, falls back to the ``LITELLM_API_KEY`` environment variable and then + to LiteLLM's own provider-specific environment variable lookup. + endpoint: Optional base URL override (e.g. for a self-hosted proxy or LiteLLM gateway). + Falls back to the ``LITELLM_ENDPOINT`` environment variable. + headers: Optional extra HTTP headers forwarded to the provider (``extra_headers``). + temperature: Sampling temperature (0-2). + top_p: Nucleus sampling probability (0-1). + max_tokens: Maximum number of tokens to generate. Deprecated by OpenAI in favor of + ``max_completion_tokens``; not compatible with reasoning (o1-series) models. + max_completion_tokens: Upper bound on tokens generated for a completion, including + reasoning tokens. Use this for reasoning (o1-series) models. Mutually exclusive with + ``max_tokens``. + frequency_penalty: Penalize frequently generated tokens. + presence_penalty: Penalize tokens already present in the conversation. + seed: Best-effort deterministic sampling seed. + n: Number of completions to generate. + stop: Stop sequence(s). + audio_response_config: Optional audio-output configuration (voice + format). When set, + audio modality is enabled on the request (``modalities``/``audio``) for models that + support it (e.g. ``gpt-4o-audio-preview``), and audio responses are saved as + ``audio_path`` pieces alongside their transcript. + extra_body_parameters: Additional provider parameters merged into the request body + (passthrough). Unsupported params are silently dropped per-provider, so anything a + given provider does not accept is ignored rather than raising. These may also override + the target's defaults (e.g. pass ``{"drop_params": False}`` to opt into strict param + validation, or ``{"timeout": 30}`` to set a per-request timeout). + underlying_model: The underlying model name (e.g. ``"gpt-4o"``) used for capability + lookup and identification when the provider/model string differs from a known model. + max_requests_per_minute: Client-side request cap. + custom_configuration: Override the derived target configuration. + """ + + # Fallback only. The real per-instance configuration is normally derived from LiteLLM's + # model metadata at construction time (see ``_derive_capabilities_from_litellm``). + _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_editable_history=True, + supports_system_prompt=True, + input_modalities=_TEXT_INPUT, + ) + ) + + def __init__( + self, + *, + model_name: str | None = None, + api_key: str | Callable[[], str | Awaitable[str]] | None = None, + endpoint: str | None = None, + headers: dict[str, str] | None = None, + temperature: float | None = None, + top_p: float | None = None, + max_tokens: int | None = None, + max_completion_tokens: int | None = None, + frequency_penalty: float | None = None, + presence_penalty: float | None = None, + seed: int | None = None, + n: int | None = None, + stop: str | list[str] | None = None, + audio_response_config: OpenAIChatAudioConfig | None = None, + extra_body_parameters: dict[str, Any] | None = None, + underlying_model: str | None = None, + max_requests_per_minute: int | None = None, + custom_configuration: TargetConfiguration | None = None, + ) -> None: + """ + Initialize a LiteLLMChatTarget. + + Raises: + ValueError: If model_name is not provided and LITELLM_MODEL env var is not set, or if + both max_tokens and max_completion_tokens are provided. + """ + resolved_model = model_name or os.environ.get("LITELLM_MODEL", "") + if not resolved_model: + raise ValueError("model_name is required. Pass it directly or set the LITELLM_MODEL environment variable.") + + if max_tokens is not None and max_completion_tokens is not None: + raise ValueError("Cannot provide both max_tokens and max_completion_tokens.") + + validate_temperature(temperature) + validate_top_p(top_p) + + super().__init__( + model_name=resolved_model, + underlying_model=underlying_model, + max_requests_per_minute=max_requests_per_minute, + custom_configuration=custom_configuration, + ) + + # Resolve api_key: explicit value/callable > LITELLM_API_KEY env var > None (LiteLLM + # then reads provider-specific env vars itself). ``ensure_async_token_provider`` wraps a + # sync token provider so we can uniformly await it at request time. + if api_key is None: + api_key = os.environ.get("LITELLM_API_KEY") + self._api_key = ensure_async_token_provider(api_key) + + self._endpoint = endpoint or os.environ.get("LITELLM_ENDPOINT") + self._headers = headers + self._temperature = temperature + self._top_p = top_p + self._max_tokens = max_tokens + self._max_completion_tokens = max_completion_tokens + self._frequency_penalty = frequency_penalty + self._presence_penalty = presence_penalty + self._seed = seed + self._n = n + self._stop = stop + self._audio_response_config = audio_response_config + + # Merge audio-output config into the passthrough body (modalities + audio params), so it + # rides the same passthrough OpenAIChatTarget uses. + if audio_response_config: + audio_params = audio_response_config.to_extra_body_parameters() + extra_body_parameters = {**audio_params, **extra_body_parameters} if extra_body_parameters else audio_params + + self._extra_body_parameters = extra_body_parameters + + # Delegate transient/rate-limit retry to LiteLLM (provider-aware: honors ``Retry-After`` + # and per-provider rate-limit semantics), rather than stacking PyRIT's + # ``pyrit_target_retry`` (which would double-retry). Derive the count from PyRIT's global + # ``RETRY_MAX_NUM_ATTEMPTS`` convention; ``num_retries`` is a retry count so it is + # attempts minus one. Per-request timeout is left to LiteLLM's own default; advanced + # callers can override it via ``extra_body_parameters={"timeout": ...}``. + self._num_retries = max(get_retry_max_num_attempts() - 1, 0) + + # Capability precedence: custom_configuration > known underlying_model profile > + # LiteLLM-derived default > conservative fallback. The base __init__ already applied the + # first two; only derive from LiteLLM metadata when neither was supplied. + if custom_configuration is None: + known = TargetCapabilities.get_known_capabilities(underlying_model) if underlying_model else None + if known is None: + derived = self._derive_capabilities_from_litellm(resolved_model) + if derived is not None: + self._configuration = TargetConfiguration(capabilities=derived) + + @staticmethod + def _import_litellm() -> Any: + try: + import litellm + except ImportError as e: + raise ImportError( + "The litellm package is required for LiteLLMChatTarget. Install it with: pip install pyrit[litellm]" + ) from e + return litellm + + def _derive_capabilities_from_litellm(self, model: str) -> TargetCapabilities | None: + """ + Best-effort capability derivation from LiteLLM's model metadata. + + Uses LiteLLM's own model-capability helpers (``supports_vision``, + ``supports_response_schema``, ``get_supported_openai_params``) rather than reinventing a + per-provider capability table. Returns None if LiteLLM is unavailable so the caller falls + back to ``_DEFAULT_CONFIGURATION``. + + Args: + model (str): The LiteLLM model string to inspect. + + Returns: + TargetCapabilities | None: The derived capabilities, or None if LiteLLM is + unavailable. + """ + try: + litellm = self._import_litellm() + except ImportError: + return None + + def _supports(attr: str) -> bool: + fn = getattr(litellm, attr, None) + if fn is None: + return False + try: + return bool(fn(model)) + except Exception: + return False + + supports_vision = _supports("supports_vision") + supports_json_schema = _supports("supports_response_schema") + supports_audio_input = _supports("supports_audio_input") + supports_audio_output = _supports("supports_audio_output") + + try: + supported_params = litellm.get_supported_openai_params(model=model) or [] + except Exception: + supported_params = [] + supports_json_output = supports_json_schema or ("response_format" in supported_params) + + return TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_editable_history=True, + supports_system_prompt=True, + supports_json_output=supports_json_output, + supports_json_schema=supports_json_schema, + input_modalities=_build_input_modalities(image=supports_vision, audio=supports_audio_input), + output_modalities=_build_output_modalities(audio=supports_audio_output), + ) + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the identifier with LiteLLM-specific behavioral parameters. + + The API key is intentionally excluded. + + Returns: + ComponentIdentifier: The identifier for this target instance. + """ + return self._create_identifier( + params={ + "endpoint": self._endpoint, + "temperature": self._temperature, + "top_p": self._top_p, + "max_tokens": self._max_tokens, + "max_completion_tokens": self._max_completion_tokens, + "frequency_penalty": self._frequency_penalty, + "presence_penalty": self._presence_penalty, + "seed": self._seed, + "n": self._n, + "stop": self._stop, + }, + ) + + def is_json_response_supported(self) -> bool: + """ + Whether this target honors a JSON ``response_format`` request. + + Returns: + bool: True if the target advertises JSON output support. + """ + return self.capabilities.supports_json_output + + # Not decorated with ``pyrit_target_retry``: LiteLLM owns transient/rate-limit retry via + # ``num_retries`` (see ``__init__``). Stacking both would multiply the retry count. + @limit_requests_per_minute + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + litellm = self._import_litellm() + + message = normalized_conversation[-1] + request_piece: MessagePiece = message.message_pieces[0] + + logger.info(f"Sending prompt to LiteLLM target ({self._model_name}): {message}") + + json_config = self._get_json_response_config(message_piece=request_piece) + messages = await self._build_chat_messages_async(normalized_conversation) + api_key = await self._resolve_api_key_async() + body = self._construct_request_body(messages=messages, json_config=json_config, api_key=api_key) + + try: + response = await litellm.acompletion(**body) + except Exception as exc: + return self._handle_litellm_exception(exc=exc, request=request_piece) + + # Content filtering is red-team critical: surface it as an error Message (not an + # exception) so attacks can continue and blocked-content scorers can still score. + if is_content_filter_response(response): + return [ + build_content_filter_message( + response=response, + request=request_piece, + partial_content=extract_partial_content(response), + ) + ] + + validate_chat_completion_response(response=response) + return [await self._construct_message_from_response(response=response, request=request_piece)] + + async def _resolve_api_key_async(self) -> str | None: + """ + Resolve the api_key to a concrete string, awaiting async token providers. + + Returns: + str | None: The resolved API key, or None when no key is configured. + """ + api_key = self._api_key + if api_key is None or isinstance(api_key, str): + return api_key + + result = api_key() + if isinstance(result, Awaitable): + result = await result + return result + + async def _build_chat_messages_async(self, conversation: MutableSequence[Message]) -> list[dict[str, Any]]: + # Text-only conversations use the simpler {"role", "content": str} form, which the widest + # set of OpenAI-"compatible" providers accept. + if is_text_only_conversation(conversation): + return build_text_chat_messages(conversation) + + prefer_transcript_for_history = bool( + self._audio_response_config and self._audio_response_config.prefer_transcript_for_history + ) + return await build_multimodal_chat_messages_async( + conversation, prefer_transcript_for_history=prefer_transcript_for_history + ) + + def _construct_request_body( + self, + *, + messages: list[dict[str, Any]], + json_config: Any, + api_key: str | None = None, + ) -> dict[str, Any]: + body: dict[str, Any] = { + "model": self._model_name, + "messages": messages, + # Always drop provider-unsupported params. As a cross-provider target we send the full + # OpenAI param set, but providers support different subsets; without this LiteLLM would + # raise on any unsupported param. Advanced callers can flip it via extra_body_parameters. + "drop_params": True, + "api_key": api_key, + "api_base": self._endpoint, + "extra_headers": self._headers, + "temperature": self._temperature, + "top_p": self._top_p, + "max_tokens": self._max_tokens, + "max_completion_tokens": self._max_completion_tokens, + "frequency_penalty": self._frequency_penalty, + "presence_penalty": self._presence_penalty, + "seed": self._seed, + "n": self._n, + "stop": self._stop, + "num_retries": self._num_retries, + "response_format": build_response_format(json_config=json_config), + } + + # Passthrough for arbitrary provider params (may override defaults above, e.g. drop_params). + if self._extra_body_parameters: + body.update(self._extra_body_parameters) + + return {k: v for k, v in body.items() if v is not None} + + async def _construct_message_from_response(self, *, response: Any, request: MessagePiece) -> Message: + audio_format = self._audio_response_config.audio_format if self._audio_response_config else "wav" + 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) + self._capture_response_cost(pieces=pieces, response=response) + return Message(message_pieces=pieces) + + def _capture_response_cost(self, *, pieces: list[MessagePiece], response: Any) -> None: + """ + Record LiteLLM's computed per-call dollar cost into the first piece's metadata. + + LiteLLM attaches the spend for a completion at ``response._hidden_params["response_cost"]`` + and, failing that, can recompute it via ``litellm.completion_cost``. Cost is provider- and + 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. + + Args: + pieces (list[MessagePiece]): The constructed response pieces. + response (Any): The LiteLLM completion response object. + """ + if not pieces: + return + cost = self._extract_response_cost(response=response) + if cost is None: + return + pieces[0].prompt_metadata["token_usage_cost"] = str(cost) + + @staticmethod + def _extract_response_cost(*, response: Any) -> float | None: + """ + Pull the per-call cost from a LiteLLM response, or None when it cannot be determined. + + Prefers LiteLLM's authoritative post-call ``_hidden_params["response_cost"]`` (which may be + ``0.0`` for free/local models) and falls back to recomputing via ``litellm.completion_cost``. + + Args: + response (Any): The LiteLLM completion response object. + + Returns: + float | None: The cost in dollars, or None if unavailable. + """ + hidden_params = getattr(response, "_hidden_params", None) + if isinstance(hidden_params, dict) and hidden_params.get("response_cost") is not None: + try: + return float(hidden_params["response_cost"]) + except (TypeError, ValueError): + return None + try: + litellm = LiteLLMChatTarget._import_litellm() + cost = litellm.completion_cost(completion_response=response) + return float(cost) if cost else None + except Exception: + return None + + def _handle_litellm_exception(self, *, exc: Exception, request: MessagePiece) -> list[Message]: + """ + Translate a LiteLLM exception into either a blocked-content error Message or a PyRIT + exception. LiteLLM re-exports the OpenAI SDK exception classes, so we match on those + types (``isinstance``) rather than fragile string/qualname comparisons. + + Args: + exc (Exception): The exception raised by ``litellm.acompletion``. + request (MessagePiece): The originating request piece. + + Returns: + list[Message]: A single error Message when the failure is a content-policy block. + + Raises: + RateLimitException: For rate-limit and transient provider errors. + PyritException: For authentication and all other errors. + """ + litellm = self._import_litellm() + exceptions = litellm.exceptions + + # Content policy violations are surfaced as an error Message (not raised) so attacks + # continue and blocked-content scorers can score the refusal. + if self._is_content_policy_error(exc=exc, exceptions=exceptions): + status_code = getattr(exc, "status_code", 400) or 400 + return [ + handle_bad_request_exception( + response_text=str(exc), + request=request, + error_code=status_code, + is_content_filter=True, + ) + ] + + return self._raise_translated_exception(exc=exc, exceptions=exceptions) + + @staticmethod + def _is_content_policy_error(*, exc: Exception, exceptions: Any) -> bool: + content_policy_error = getattr(exceptions, "ContentPolicyViolationError", None) + if content_policy_error is not None and isinstance(exc, content_policy_error): + return True + text = str(exc).lower() + return "content_filter" in text or "content policy" in text or "content_policy" in text + + @staticmethod + def _raise_translated_exception(*, exc: Exception, exceptions: Any) -> NoReturn: + rate_limit_error = getattr(exceptions, "RateLimitError", ()) + transient_errors = tuple( + err + for err in ( + getattr(exceptions, "APIConnectionError", None), + getattr(exceptions, "Timeout", None), + getattr(exceptions, "InternalServerError", None), + getattr(exceptions, "ServiceUnavailableError", None), + ) + if err is not None + ) + authentication_error = getattr(exceptions, "AuthenticationError", ()) + + if isinstance(exc, rate_limit_error): + raise RateLimitException(status_code=429, message=f"Rate limited by provider: {exc}") from exc + if transient_errors and isinstance(exc, transient_errors): + status_code = getattr(exc, "status_code", 503) or 503 + raise RateLimitException(status_code=status_code, message=f"Transient provider error: {exc}") from exc + if isinstance(exc, authentication_error): + raise PyritException(message=f"Authentication failed: {exc}") from exc + + raise PyritException(message=f"LiteLLM error: {exc}") from exc diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 35454d6e69..6ddef6bea8 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/pyrit/prompt_target/openai/openai_chat_target.py @@ -1,28 +1,36 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import base64 -import json import logging from collections.abc import MutableSequence from typing import Any, Optional -from pyrit.common.data_url_converter import convert_local_image_to_data_url from pyrit.exceptions import ( EmptyResponseException, - PyritException, pyrit_target_retry, ) from pyrit.identifiers import ComponentIdentifier from pyrit.models import ( - ChatMessage, - DataTypeSerializer, Message, MessagePiece, - construct_response_from_request, - data_serializer_factory, ) from pyrit.models.json_response_config import _JsonResponseConfig +from pyrit.prompt_target.common.chat_completions_message_builder import ( + build_multimodal_chat_messages_async, + build_response_format, + build_text_chat_messages, + is_text_only_conversation, + should_skip_audio_piece, +) +from pyrit.prompt_target.common.chat_completions_response_parser import ( + build_response_pieces_async, + capture_token_usage, + detect_response_content, + extract_partial_content, + is_content_filter_response, + save_audio_response_async, + validate_chat_completion_response, +) from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration @@ -253,12 +261,7 @@ def _check_content_filter(self, response: Any) -> bool: Returns: True if content was filtered, False otherwise. """ - try: - if response.choices and response.choices[0].finish_reason == "content_filter": - return True - except (AttributeError, IndexError): - pass - return False + return is_content_filter_response(response) def _extract_partial_content(self, response: Any) -> Optional[str]: """ @@ -273,12 +276,7 @@ def _extract_partial_content(self, response: Any) -> Optional[str]: Returns: The partial text content, or None if no content was generated. """ - try: - if response.choices and response.choices[0].message and response.choices[0].message.content: - return response.choices[0].message.content - except (AttributeError, IndexError): - pass - return None + return extract_partial_content(response) def _validate_response(self, response: Any, request: MessagePiece) -> Optional[Message]: """ @@ -300,30 +298,7 @@ def _validate_response(self, response: Any, request: MessagePiece) -> Optional[M PyritException: For unexpected response structures or finish reasons. EmptyResponseException: When the API returns an empty response. """ - # Check for missing choices - if not hasattr(response, "choices") or not response.choices: - raise PyritException(message="No choices returned in the completion response.") - - choice = response.choices[0] - finish_reason = choice.finish_reason - - # Check finish_reason (content_filter is handled by _check_content_filter) - # "tool_calls" is valid when the model invokes functions - valid_finish_reasons = ["stop", "length", "content_filter", "tool_calls"] - if finish_reason not in valid_finish_reasons: - raise PyritException( - message=f"Unknown finish_reason {finish_reason} from response: {response.model_dump_json()}" - ) - - # Check for at least one valid response type - has_content, has_audio, has_tool_calls = self._detect_response_content(choice.message) - - if not (has_content or has_audio or has_tool_calls): - logger.error("The chat returned an empty response (no content, audio, or tool_calls).") - raise EmptyResponseException( - message="The chat returned an empty response (no content, audio, or tool_calls)." - ) - + validate_chat_completion_response(response=response) return None def _detect_response_content(self, message: Any) -> tuple[bool, bool, bool]: @@ -336,10 +311,7 @@ def _detect_response_content(self, message: Any) -> tuple[bool, bool, bool]: Returns: Tuple of (has_content, has_audio, has_tool_calls) booleans. """ - has_content = bool(message.content) - has_audio = hasattr(message, "audio") and message.audio is not None - has_tool_calls = hasattr(message, "tool_calls") and message.tool_calls - return has_content, has_audio, has_tool_calls + return detect_response_content(message) def _should_skip_sending_audio( self, @@ -362,20 +334,14 @@ def _should_skip_sending_audio( if message_piece.converted_value_data_type != "audio_path": return False - api_role = message_piece.api_role - - # Skip audio for assistant messages - OpenAI only allows audio in user messages. - # For assistant responses, the transcript text piece should already be included. - if api_role == "assistant": - return True - - # Skip historical user audio if prefer_transcript_for_history is enabled and we have a transcript - return bool( - api_role == "user" - and not is_last_message - and has_text_piece - and self._audio_response_config - and self._audio_response_config.prefer_transcript_for_history + prefer_transcript_for_history = bool( + self._audio_response_config and self._audio_response_config.prefer_transcript_for_history + ) + return should_skip_audio_piece( + message_piece=message_piece, + is_last_message=is_last_message, + has_text_piece=has_text_piece, + prefer_transcript_for_history=prefer_transcript_for_history, ) async def _construct_message_from_response(self, response: Any, request: MessagePiece) -> Message: @@ -397,75 +363,14 @@ async def _construct_message_from_response(self, response: Any, request: Message Raises: EmptyResponseException: If the response contains no content, audio, or tool calls. """ - message = response.choices[0].message - has_content, has_audio, has_tool_calls = self._detect_response_content(message) - - pieces: list[MessagePiece] = [] - - # Handle text content - if has_content: - text_piece = construct_response_from_request( - request=request, - response_text_pieces=[message.content], - response_type="text", - ).message_pieces[0] - pieces.append(text_piece) - - # Handle audio response (transcript + saved audio file) - if has_audio: - audio_response = message.audio - - # Add transcript as text piece with metadata - audio_transcript: Optional[str] = getattr(audio_response, "transcript", None) - if audio_transcript: - transcript_piece = construct_response_from_request( - request=request, - response_text_pieces=[audio_transcript], - response_type="text", - prompt_metadata={"transcription": "audio"}, - ).message_pieces[0] - pieces.append(transcript_piece) - - # Save audio data and add as audio_path piece - audio_data: Optional[str] = getattr(audio_response, "data", None) - if audio_data: - audio_path = await self._save_audio_response_async(audio_data_base64=audio_data) - audio_piece = construct_response_from_request( - request=request, - response_text_pieces=[audio_path], - response_type="audio_path", - ).message_pieces[0] - pieces.append(audio_piece) - - # Handle tool calls; for completions it is always function at the time of writing - if has_tool_calls: - for tool_call in message.tool_calls: - tool_call_data = { - "type": "function", - "id": tool_call.id, - "function": { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments, - }, - } - tool_call_json = json.dumps(tool_call_data) - tool_piece = construct_response_from_request( - request=request, - response_text_pieces=[tool_call_json], - response_type="function_call", - ).message_pieces[0] - pieces.append(tool_piece) + audio_format = self._audio_response_config.audio_format if self._audio_response_config else "wav" + 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.") # Capture token usage from the API response and store in the first piece's metadata - if hasattr(response, "usage") and response.usage and pieces: - pieces[0].prompt_metadata["token_usage_model_name"] = getattr(response, "model", "unknown") - pieces[0].prompt_metadata["token_usage_prompt_tokens"] = getattr(response.usage, "prompt_tokens", 0) - pieces[0].prompt_metadata["token_usage_completion_tokens"] = getattr(response.usage, "completion_tokens", 0) - pieces[0].prompt_metadata["token_usage_total_tokens"] = getattr(response.usage, "total_tokens", 0) - pieces[0].prompt_metadata["token_usage_cached_tokens"] = getattr(response.usage, "cached_tokens", 0) + capture_token_usage(pieces=pieces, response=response) return Message(message_pieces=pieces) @@ -479,31 +384,8 @@ async def _save_audio_response_async(self, *, audio_data_base64: str) -> str: Returns: str: The file path where the audio was saved. """ - audio_bytes = base64.b64decode(audio_data_base64) - - # Determine the format from config, default to wav audio_format = self._audio_response_config.audio_format if self._audio_response_config else "wav" - extension = f".{audio_format}" if audio_format != "pcm16" else ".wav" - - audio_serializer = data_serializer_factory( - category="prompt-memory-entries", - data_type="audio_path", - extension=extension, - ) - - if audio_format == "pcm16": - # Raw PCM needs WAV headers - OpenAI uses 24kHz mono PCM16 - await audio_serializer.save_formatted_audio( - data=audio_bytes, - num_channels=1, - sample_width=2, - sample_rate=24000, - ) - else: - # wav, mp3, flac, opus are already properly formatted - await audio_serializer.save_data(audio_bytes) - - return audio_serializer.value + return await save_audio_response_async(audio_data_base64=audio_data_base64, audio_format=audio_format) async def _build_chat_messages_async(self, conversation: MutableSequence[Message]) -> list[dict[str, Any]]: """ @@ -529,12 +411,7 @@ def _is_text_message_format(self, conversation: MutableSequence[Message]) -> boo Returns: bool: True if the message piece is in text message format, False otherwise. """ - for turn in conversation: - if len(turn.message_pieces) != 1: - return False - if turn.message_pieces[0].converted_value_data_type not in ("text", "error"): - return False - return True + return is_text_only_conversation(conversation) def _build_chat_messages_for_text(self, conversation: MutableSequence[Message]) -> list[dict[str, Any]]: """ @@ -551,25 +428,7 @@ def _build_chat_messages_for_text(self, conversation: MutableSequence[Message]) ValueError: If any message does not have exactly one text piece. ValueError: If any message piece is not of type text. """ - chat_messages: list[dict[str, Any]] = [] - for message in conversation: - # validated to only have one text entry - - if len(message.message_pieces) != 1: - raise ValueError("_build_chat_messages_for_text only supports a single message piece.") - - message_piece = message.message_pieces[0] - - if message_piece.converted_value_data_type not in ("text", "error"): - raise ValueError( - f"_build_chat_messages_for_text only supports text and error data types." - f" Received: {message_piece.converted_value_data_type}." - ) - - chat_message = ChatMessage(role=message_piece.api_role, content=message_piece.converted_value) - chat_messages.append(chat_message.model_dump(exclude_none=True)) - - return chat_messages + return build_text_chat_messages(conversation) async def _build_chat_messages_for_multi_modal_async( self, conversation: MutableSequence[Message] @@ -587,68 +446,12 @@ async def _build_chat_messages_for_multi_modal_async( ValueError: If any message does not have a role. ValueError: If any message piece has an unsupported data type. """ - chat_messages: list[dict[str, Any]] = [] - last_message_index = len(conversation) - 1 - - for message_index, message in enumerate(conversation): - message_pieces = message.message_pieces - is_last_message = message_index == last_message_index - - # Check if this message has a text piece (transcript) alongside audio - has_text_piece = any(mp.converted_value_data_type == "text" for mp in message_pieces) - - content = [] - role = None - for message_piece in message_pieces: - role = message_piece.api_role - - if self._should_skip_sending_audio( - message_piece=message_piece, - is_last_message=is_last_message, - has_text_piece=has_text_piece, - ): - continue - - if message_piece.converted_value_data_type in ("text", "error"): - entry = {"type": "text", "text": message_piece.converted_value} - content.append(entry) - elif message_piece.converted_value_data_type == "image_path": - data_base64_encoded_url = await convert_local_image_to_data_url(message_piece.converted_value) - image_url_entry = {"url": data_base64_encoded_url} - entry = {"type": "image_url", "image_url": image_url_entry} - content.append(entry) - elif message_piece.converted_value_data_type == "audio_path": - ext = DataTypeSerializer.get_extension(message_piece.converted_value) - # OpenAI SDK: openai/types/chat/chat_completion_content_part_input_audio_param.py - # defines format: Required[Literal["wav", "mp3"]] - if not ext or ext.lower() not in [".wav", ".mp3"]: - raise ValueError( - f"Unsupported audio format: {ext}. " - "OpenAI Chat Completions API input_audio only supports .wav and .mp3. " - "Note: This is different from the Whisper Speech-to-Text API which supports more formats." - ) - audio_serializer = data_serializer_factory( - category="prompt-memory-entries", - value=message_piece.converted_value, - data_type="audio_path", - extension=ext, - ) - base64_data = await audio_serializer.read_data_base64() - audio_format = ext.lower().lstrip(".") - input_audio_entry = {"data": base64_data, "format": audio_format} - entry = {"type": "input_audio", "input_audio": input_audio_entry} - content.append(entry) - else: - raise ValueError( - f"Multimodal data type {message_piece.converted_value_data_type} is not yet supported." - ) - - if not role: - raise ValueError("No role could be determined from the message pieces.") - - chat_message = ChatMessage(role=role, content=content) - chat_messages.append(chat_message.model_dump(exclude_none=True)) - return chat_messages + prefer_transcript_for_history = bool( + self._audio_response_config and self._audio_response_config.prefer_transcript_for_history + ) + return await build_multimodal_chat_messages_async( + conversation, prefer_transcript_for_history=prefer_transcript_for_history + ) async def _construct_request_body( self, *, conversation: MutableSequence[Message], json_config: _JsonResponseConfig @@ -678,17 +481,4 @@ async def _construct_request_body( return {k: v for k, v in body_parameters.items() if v is not None} def _build_response_format(self, json_config: _JsonResponseConfig) -> Optional[dict[str, Any]]: - if not json_config.enabled: - return None - - if json_config.schema: - return { - "type": "json_schema", - "json_schema": { - "name": json_config.schema_name, - "schema": json_config.schema, - "strict": json_config.strict, - }, - } - - return {"type": "json_object"} + return build_response_format(json_config=json_config) diff --git a/tests/integration/targets/test_litellm_chat_target_integration.py b/tests/integration/targets/test_litellm_chat_target_integration.py new file mode 100644 index 0000000000..1a63cb0ab4 --- /dev/null +++ b/tests/integration/targets/test_litellm_chat_target_integration.py @@ -0,0 +1,239 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Integration tests for LiteLLMChatTarget. + +These reuse the existing ``PLATFORM_OPENAI_CHAT_*`` environment configuration to exercise +LiteLLM against a real OpenAI-compatible endpoint (LiteLLM's ``openai/`` provider prefix). + +They verify: +- Basic text completion +- Tool calling via the ``extra_body_parameters`` passthrough +- Multimodal image input (vision) +- Multimodal audio input/output +- Token-usage metadata capture (parsed back through ``TokenUsage``) + +The chat/vision model defaults to ``gpt-5.4`` and the audio model to ``gpt-audio``; both can be +overridden with the ``PLATFORM_OPENAI_CHAT_MODEL`` / ``PLATFORM_OPENAI_AUDIO_MODEL`` env vars to +match whatever deployment names the target endpoint exposes. +""" + +import json +import os +import uuid + +import pytest + +from pyrit.common.path import HOME_PATH +from pyrit.models import Message, MessagePiece, TokenUsage +from pyrit.prompt_target import ( + LiteLLMChatTarget, + OpenAIChatAudioConfig, + TargetCapabilities, + TargetConfiguration, +) + +# Assets reused for multimodal parity checks. +SAMPLE_IMAGE_FILE = HOME_PATH / "assets" / "pyrit_architecture.png" +SAMPLE_AUDIO_FILE = HOME_PATH / "assets" / "converted_audio.wav" + + +@pytest.fixture() +def platform_litellm_args(): + """ + Requires: + - PLATFORM_OPENAI_CHAT_ENDPOINT: An OpenAI-compatible endpoint (e.g. https://api.openai.com/v1) + - PLATFORM_OPENAI_CHAT_KEY: The API key + """ + endpoint = os.environ.get("PLATFORM_OPENAI_CHAT_ENDPOINT") + api_key = os.environ.get("PLATFORM_OPENAI_CHAT_KEY") + + if not endpoint or not api_key: + pytest.skip("PLATFORM_OPENAI_CHAT_ENDPOINT and PLATFORM_OPENAI_CHAT_KEY must be set") + + model = os.environ.get("PLATFORM_OPENAI_CHAT_MODEL", "gpt-5.4") + return { + "model_name": f"openai/{model}", + "endpoint": endpoint, + "api_key": api_key, + } + + +@pytest.fixture() +def platform_litellm_audio_args(): + """ + Requires: + - PLATFORM_OPENAI_CHAT_ENDPOINT: An OpenAI-compatible endpoint + - PLATFORM_OPENAI_CHAT_KEY: The API key + """ + endpoint = os.environ.get("PLATFORM_OPENAI_CHAT_ENDPOINT") + api_key = os.environ.get("PLATFORM_OPENAI_CHAT_KEY") + + if not endpoint or not api_key: + pytest.skip("PLATFORM_OPENAI_CHAT_ENDPOINT and PLATFORM_OPENAI_CHAT_KEY must be set") + + model = os.environ.get("PLATFORM_OPENAI_AUDIO_MODEL", "gpt-audio") + return { + "model_name": f"openai/{model}", + "endpoint": endpoint, + "api_key": api_key, + } + + +@pytest.mark.run_only_if_all_tests +async def test_litellm_chat_target_text_completion(sqlite_instance, platform_litellm_args): + target = LiteLLMChatTarget(**platform_litellm_args) + + user_piece = MessagePiece( + role="user", + original_value="Reply with exactly the word: pong", + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), + ) + + result = await target.send_prompt_async(message=user_piece.to_message()) + + assert result is not None + assert len(result) >= 1 + text_pieces = [p for p in result[0].message_pieces if p.converted_value_data_type == "text"] + assert len(text_pieces) >= 1 + assert "pong" in text_pieces[0].converted_value.lower() + + +@pytest.mark.run_only_if_all_tests +async def test_litellm_chat_target_tool_calling(sqlite_instance, platform_litellm_args): + tools = [ + { + "type": "function", + "function": { + "name": "get_stock_price", + "description": "Get the current stock price for a given ticker symbol", + "parameters": { + "type": "object", + "properties": { + "ticker": {"type": "string", "description": "The stock ticker symbol, e.g. AAPL"}, + }, + "required": ["ticker"], + }, + }, + }, + ] + + target = LiteLLMChatTarget( + **platform_litellm_args, + extra_body_parameters={"tools": tools, "tool_choice": "auto"}, + ) + + user_piece = MessagePiece( + role="user", + original_value="What's the current stock price for Microsoft (MSFT)?", + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), + ) + + result = await target.send_prompt_async(message=user_piece.to_message()) + + tool_call_pieces = [p for p in result[0].message_pieces if p.converted_value_data_type == "function_call"] + assert len(tool_call_pieces) >= 1, "Response should contain at least one tool call" + tool_call_data = json.loads(tool_call_pieces[0].converted_value) + assert tool_call_data["function"]["name"] == "get_stock_price" + assert "msft" in tool_call_data["function"]["arguments"].lower() + + +@pytest.mark.run_only_if_all_tests +async def test_litellm_chat_target_token_usage_in_metadata(sqlite_instance, platform_litellm_args): + target = LiteLLMChatTarget(**platform_litellm_args) + + user_piece = MessagePiece( + role="user", + original_value="Say hello in one word.", + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), + ) + + result = await target.send_prompt_async(message=user_piece.to_message()) + + metadata = result[0].message_pieces[0].prompt_metadata + usage = TokenUsage.from_metadata(metadata) + assert usage is not None + assert usage.input_tokens is not None and usage.input_tokens > 0 + assert usage.output_tokens is not None and usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.run_only_if_all_tests +async def test_litellm_chat_target_image_input(sqlite_instance, platform_litellm_args): + """Send a text + image message and verify the vision model returns a text response.""" + target = LiteLLMChatTarget( + **platform_litellm_args, + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset( + {frozenset({"text", "image_path"}), frozenset({"text"}), frozenset({"image_path"})} + ), + ) + ), + ) + + conv_id = str(uuid.uuid4()) + text_piece = MessagePiece( + role="user", + original_value="Describe what this image shows in one short sentence.", + original_value_data_type="text", + conversation_id=conv_id, + ) + image_piece = MessagePiece( + role="user", + original_value=str(SAMPLE_IMAGE_FILE), + original_value_data_type="image_path", + conversation_id=conv_id, + ) + + result = await target.send_prompt_async(message=Message(message_pieces=[text_piece, image_piece])) + + assert result is not None + text_pieces = [p for p in result[0].message_pieces if p.converted_value_data_type == "text"] + assert len(text_pieces) >= 1 + assert text_pieces[0].converted_value.strip() + + +@pytest.mark.run_only_if_all_tests +async def test_litellm_chat_target_audio_input_output(sqlite_instance, platform_litellm_audio_args): + """Verify audio output generation and audio input handling against an audio-capable model.""" + audio_config = OpenAIChatAudioConfig(voice="alloy", audio_format="wav") + + target = LiteLLMChatTarget( + **platform_litellm_audio_args, + audio_response_config=audio_config, + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + input_modalities=frozenset( + {frozenset({"text", "audio_path"}), frozenset({"text"}), frozenset({"audio_path"})} + ), + ) + ), + ) + + conv_id = str(uuid.uuid4()) + + text_piece = MessagePiece( + role="user", + original_value="Hello! What's your name?", + original_value_data_type="text", + conversation_id=conv_id, + ) + result1 = await target.send_prompt_async(message=text_piece.to_message()) + audio_pieces1 = [p for p in result1[0].message_pieces if p.converted_value_data_type == "audio_path"] + assert len(audio_pieces1) >= 1, "First response should contain audio" + + audio_piece = MessagePiece( + role="user", + original_value=str(SAMPLE_AUDIO_FILE), + original_value_data_type="audio_path", + conversation_id=conv_id, + ) + result2 = await target.send_prompt_async(message=audio_piece.to_message()) + audio_pieces2 = [p for p in result2[0].message_pieces if p.converted_value_data_type == "audio_path"] + assert len(audio_pieces2) >= 1, "Second response should contain audio" diff --git a/tests/integration/targets/test_openai_chat_target_integration.py b/tests/integration/targets/test_openai_chat_target_integration.py index 5ae43605d9..563ef62df5 100644 --- a/tests/integration/targets/test_openai_chat_target_integration.py +++ b/tests/integration/targets/test_openai_chat_target_integration.py @@ -16,7 +16,7 @@ import pytest from pyrit.common.path import HOME_PATH -from pyrit.models import MessagePiece +from pyrit.models import MessagePiece, TokenUsage from pyrit.prompt_target import OpenAIChatAudioConfig, OpenAIChatTarget, TargetCapabilities, TargetConfiguration # Path to sample audio file for testing @@ -41,7 +41,7 @@ def platform_openai_audio_args(): return { "endpoint": endpoint, "api_key": api_key, - "model_name": "gpt-audio", + "model_name": os.environ.get("PLATFORM_OPENAI_AUDIO_MODEL", "gpt-audio"), } @@ -63,7 +63,7 @@ def platform_openai_chat_args(): return { "endpoint": endpoint, "api_key": api_key, - "model_name": "gpt-4o", + "model_name": os.environ.get("PLATFORM_OPENAI_CHAT_MODEL", "gpt-5.4"), } @@ -216,9 +216,9 @@ async def test_openai_chat_target_token_usage_in_metadata(sqlite_instance, platf Test that token usage metadata is captured from a real API response. This test verifies that: - 1. Token usage keys are present in prompt_metadata of the response - 2. Token counts are non-negative integers - 3. Model name is a non-empty string + 1. Token usage is recoverable via ``TokenUsage.from_metadata`` + 2. Token counts are positive integers + 3. The total equals input + output """ target = OpenAIChatTarget(**platform_openai_chat_args) @@ -236,20 +236,10 @@ async def test_openai_chat_target_token_usage_in_metadata(sqlite_instance, platf assert len(result) >= 1 first_piece = result[0].message_pieces[0] - metadata = first_piece.prompt_metadata - - # Verify token usage keys are present - assert "token_usage_model_name" in metadata, "Response should contain token_usage_model_name in metadata" - assert "token_usage_prompt_tokens" in metadata, "Response should contain token_usage_prompt_tokens in metadata" - assert "token_usage_completion_tokens" in metadata - assert "token_usage_total_tokens" in metadata - - # Verify values are reasonable - assert isinstance(metadata["token_usage_model_name"], str) - assert len(metadata["token_usage_model_name"]) > 0 - assert metadata["token_usage_prompt_tokens"] > 0 - assert metadata["token_usage_completion_tokens"] > 0 - assert metadata["token_usage_total_tokens"] > 0 - assert metadata["token_usage_total_tokens"] == ( - metadata["token_usage_prompt_tokens"] + metadata["token_usage_completion_tokens"] - ) + usage = TokenUsage.from_metadata(first_piece.prompt_metadata) + + assert usage is not None, "Response should contain token-usage metadata" + assert usage.input_tokens is not None and usage.input_tokens > 0 + 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 diff --git a/tests/unit/models/test_token_usage.py b/tests/unit/models/test_token_usage.py new file mode 100644 index 0000000000..439c92efd2 --- /dev/null +++ b/tests/unit/models/test_token_usage.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.models import TokenUsage + + +def test_to_metadata_uses_input_output_key_names_and_omits_none(): + usage = TokenUsage(input_tokens=10, output_tokens=20, total_tokens=30, cached_tokens=5) + metadata = usage.to_metadata() + assert metadata["token_usage_input_tokens"] == 10 + assert metadata["token_usage_output_tokens"] == 20 + assert metadata["token_usage_total_tokens"] == 30 + assert metadata["token_usage_cached_tokens"] == 5 + assert "token_usage_reasoning_tokens" not in metadata + + +def test_to_metadata_includes_extra(): + usage = TokenUsage(input_tokens=1, output_tokens=2, extra={"output_audio_tokens": 9}) + metadata = usage.to_metadata() + assert metadata["token_usage_input_tokens"] == 1 + assert metadata["token_usage_output_audio_tokens"] == 9 + + +def test_round_trip_through_metadata(): + original = TokenUsage( + input_tokens=10, + output_tokens=20, + total_tokens=30, + reasoning_tokens=4, + cached_tokens=5, + extra={"output_audio_tokens": 3}, + ) + restored = TokenUsage.from_metadata(original.to_metadata()) + assert restored == original + + +def test_from_metadata_reads_input_output_suffixes(): + metadata = {"token_usage_input_tokens": 8, "token_usage_output_tokens": 12} + restored = TokenUsage.from_metadata(metadata) + assert restored is not None + assert restored.input_tokens == 8 + assert restored.output_tokens == 12 + + +def test_from_metadata_routes_unknown_int_keys_to_extra(): + metadata = {"token_usage_input_tokens": 10, "token_usage_output_audio_tokens": 4} + restored = TokenUsage.from_metadata(metadata) + assert restored is not None + assert restored.extra == {"output_audio_tokens": 4} + + +def test_from_metadata_ignores_cost_and_unrelated_keys(): + metadata = { + "token_usage_input_tokens": 10, + "token_usage_cost": "0.0021", + "unrelated_key": 99, + } + restored = TokenUsage.from_metadata(metadata) + assert restored is not None + assert restored.input_tokens == 10 + assert "cost" not in restored.extra + assert restored.extra == {} + + +def test_from_metadata_returns_none_without_token_usage_keys(): + assert TokenUsage.from_metadata({"partial_content": "x"}) is 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 new file mode 100644 index 0000000000..d4de5c1ae5 --- /dev/null +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -0,0 +1,462 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit tests for the shared OpenAI Chat Completions wire-format helpers.""" + +import base64 +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.exceptions import EmptyResponseException, PyritException +from pyrit.models import Message, MessagePiece +from pyrit.models.json_response_config import _JsonResponseConfig +from pyrit.prompt_target.common.chat_completions_message_builder import ( + build_multimodal_chat_messages_async, + build_response_format, + build_text_chat_messages, + build_text_content_entry, + is_text_only_conversation, + should_skip_audio_piece, +) +from pyrit.prompt_target.common.chat_completions_response_parser import ( + build_audio_pieces_async, + build_content_filter_message, + build_response_pieces_async, + build_text_and_tool_pieces, + capture_token_usage, + extract_partial_content, + is_content_filter_response, + save_audio_response_async, + token_usage_from_chat_completion, + validate_chat_completion_response, +) + +pytestmark = pytest.mark.usefixtures("patch_central_database") + + +def _text_message(text="hi", role="user"): + return MessagePiece( + role=role, conversation_id="c", original_value=text, original_value_data_type="text" + ).to_message() + + +def _request_piece(text="ask"): + return MessagePiece(role="user", conversation_id="c", original_value=text, original_value_data_type="text") + + +def _mock_response(content="hello", finish_reason="stop", tool_calls=None): + resp = MagicMock() + resp.choices = [MagicMock()] + resp.choices[0].finish_reason = finish_reason + resp.choices[0].message.content = content + resp.choices[0].message.tool_calls = tool_calls + resp.choices[0].message.audio = None + resp.model = "some-model" + resp.model_dump_json = MagicMock(return_value=json.dumps({"finish_reason": finish_reason})) + return resp + + +# --------------------------------------------------------------------------- +# message builder +# --------------------------------------------------------------------------- + + +def test_is_text_only_conversation_true(): + assert is_text_only_conversation([_text_message("a"), _text_message("b", role="assistant")]) is True + + +def test_is_text_only_conversation_false_for_multi_piece(): + text_piece = MessagePiece(role="user", conversation_id="c", original_value="a", original_value_data_type="text") + image_piece = MessagePiece( + role="user", conversation_id="c", original_value="x.png", original_value_data_type="image_path" + ) + message = Message(message_pieces=[text_piece, image_piece]) + assert is_text_only_conversation([message]) is False + + +def test_build_text_chat_messages_preserves_roles(): + messages = [_text_message("hello", "user"), _text_message("hi", "assistant")] + result = build_text_chat_messages(messages) + assert result[0] == {"role": "user", "content": "hello"} + assert result[1] == {"role": "assistant", "content": "hi"} + + +def test_build_text_content_entry(): + piece = _request_piece("describe this") + assert build_text_content_entry(message_piece=piece) == {"type": "text", "text": "describe this"} + + +def test_build_response_format_disabled_returns_none(): + config = _JsonResponseConfig.from_metadata(metadata={}) + assert build_response_format(json_config=config) is None + + +def test_build_response_format_json_object(): + config = _JsonResponseConfig.from_metadata(metadata={"response_format": "json"}) + assert build_response_format(json_config=config) == {"type": "json_object"} + + +def test_build_response_format_json_schema(): + schema = {"type": "object", "properties": {"a": {"type": "string"}}} + config = _JsonResponseConfig.from_metadata(metadata={"response_format": "json", "json_schema": json.dumps(schema)}) + result = build_response_format(json_config=config) + assert result is not None + assert result["type"] == "json_schema" + assert result["json_schema"]["schema"] == schema + + +# --------------------------------------------------------------------------- +# response parser +# --------------------------------------------------------------------------- + + +def test_validate_response_no_choices_raises(): + resp = MagicMock() + resp.choices = [] + with pytest.raises(PyritException, match="No choices"): + validate_chat_completion_response(response=resp) + + +def test_validate_response_unknown_finish_reason_raises(): + with pytest.raises(PyritException, match="Unknown finish_reason"): + validate_chat_completion_response(response=_mock_response(finish_reason="banana")) + + +def test_validate_response_empty_raises(): + resp = _mock_response(content=None) + with pytest.raises(EmptyResponseException): + validate_chat_completion_response(response=resp) + + +def test_validate_response_accepts_valid(): + for reason in ("stop", "length", "tool_calls", "content_filter"): + validate_chat_completion_response(response=_mock_response(finish_reason=reason)) + + +def test_build_text_and_tool_pieces_text(): + pieces = build_text_and_tool_pieces(response=_mock_response("answer"), request=_request_piece()) + assert len(pieces) == 1 + assert pieces[0].converted_value == "answer" + assert pieces[0].converted_value_data_type == "text" + + +def test_build_text_and_tool_pieces_tool_call(): + tool_call = MagicMock() + tool_call.id = "call_1" + tool_call.function.name = "get_weather" + tool_call.function.arguments = '{"loc": "SF"}' + resp = _mock_response(content=None, tool_calls=[tool_call]) + pieces = build_text_and_tool_pieces(response=resp, request=_request_piece()) + assert len(pieces) == 1 + assert pieces[0].converted_value_data_type == "function_call" + parsed = json.loads(pieces[0].converted_value) + assert parsed["function"]["name"] == "get_weather" + + +def test_capture_token_usage_populates_metadata(): + resp = _mock_response("ok") + 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 = build_text_and_tool_pieces(response=resp, request=_request_piece()) + capture_token_usage(pieces=pieces, response=resp) + metadata = pieces[0].prompt_metadata + assert metadata["token_usage_input_tokens"] == 3 + assert metadata["token_usage_output_tokens"] == 4 + assert metadata["token_usage_total_tokens"] == 7 + assert metadata["token_usage_cached_tokens"] == 1 + assert metadata["token_usage_reasoning_tokens"] == 2 + assert "token_usage_model_name" not in metadata + + +def test_capture_token_usage_noop_without_usage(): + resp = _mock_response("ok") + resp.usage = None + pieces = build_text_and_tool_pieces(response=resp, request=_request_piece()) + capture_token_usage(pieces=pieces, response=resp) + assert "token_usage_total_tokens" not in pieces[0].prompt_metadata + + +# --------------------------------------------------------------------------- +# token_usage_from_chat_completion (Chat Completions usage parsing) +# --------------------------------------------------------------------------- + + +def _usage(**kwargs): + """Build an attribute-style stand-in for a provider usage object.""" + return SimpleNamespace(**kwargs) + + +def test_token_usage_maps_prompt_completion_and_total(): + result = token_usage_from_chat_completion(_usage(prompt_tokens=10, completion_tokens=20, total_tokens=30)) + assert result.input_tokens == 10 + assert result.output_tokens == 20 + assert result.total_tokens == 30 + assert result.cached_tokens is None + assert result.reasoning_tokens is None + assert result.extra == {} + + +def test_token_usage_derives_total_when_missing(): + result = token_usage_from_chat_completion(_usage(prompt_tokens=4, completion_tokens=6)) + assert result.total_tokens == 10 + + +def test_token_usage_reads_nested_details(): + usage = _usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=_usage(cached_tokens=40, audio_tokens=8), + completion_tokens_details=_usage( + reasoning_tokens=12, audio_tokens=3, accepted_prediction_tokens=2, rejected_prediction_tokens=1 + ), + ) + result = token_usage_from_chat_completion(usage) + assert result.cached_tokens == 40 + assert result.reasoning_tokens == 12 + assert result.extra == { + "input_audio_tokens": 8, + "output_audio_tokens": 3, + "accepted_prediction_tokens": 2, + "rejected_prediction_tokens": 1, + } + + +def test_token_usage_accepts_mapping_payload(): + usage = { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + "prompt_tokens_details": {"cached_tokens": 2}, + "completion_tokens_details": {"reasoning_tokens": 3}, + } + result = token_usage_from_chat_completion(usage) + assert result.input_tokens == 5 + assert result.output_tokens == 7 + assert result.cached_tokens == 2 + assert result.reasoning_tokens == 3 + + +def test_token_usage_reads_litellm_top_level_cache_fields(): + usage = _usage( + prompt_tokens=100, + completion_tokens=20, + total_tokens=120, + cache_read_input_tokens=30, + cache_creation_input_tokens=15, + ) + result = token_usage_from_chat_completion(usage) + assert result.cached_tokens == 30 + assert result.extra == {"cache_write_tokens": 15} + + +def test_token_usage_prefers_nested_cached_over_top_level(): + usage = _usage( + prompt_tokens=100, + completion_tokens=20, + prompt_tokens_details=_usage(cached_tokens=40), + cache_read_input_tokens=30, + ) + result = token_usage_from_chat_completion(usage) + assert result.cached_tokens == 40 + + +def test_token_usage_preserves_zero_cached_tokens(): + usage = _usage(prompt_tokens=100, completion_tokens=20, prompt_tokens_details=_usage(cached_tokens=0)) + result = token_usage_from_chat_completion(usage) + assert result.cached_tokens == 0 + + +def test_token_usage_ignores_non_int_and_bool(): + result = token_usage_from_chat_completion(_usage(prompt_tokens=True, completion_tokens="5", total_tokens=None)) + assert result.input_tokens is None + assert result.output_tokens is None + assert result.total_tokens is None + + +def test_token_usage_handles_missing_details(): + result = token_usage_from_chat_completion(_usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)) + assert result.cached_tokens is None + assert result.reasoning_tokens is None + assert result.extra == {} + + +def test_token_usage_ignores_responses_api_names(): + # The Responses API shape (input_tokens/output_tokens) is intentionally not parsed here. + result = token_usage_from_chat_completion(_usage(input_tokens=7, output_tokens=3, total_tokens=10)) + assert result.input_tokens is None + assert result.output_tokens is None + assert result.total_tokens == 10 + + +def test_is_content_filter_response_true(): + assert is_content_filter_response(_mock_response(finish_reason="content_filter")) is True + + +def test_is_content_filter_response_false(): + assert is_content_filter_response(_mock_response(finish_reason="stop")) is False + + +def test_extract_partial_content_returns_text(): + assert extract_partial_content(_mock_response(content="partial")) == "partial" + + +def test_extract_partial_content_none_when_absent(): + assert extract_partial_content(_mock_response(content=None)) is None + + +def test_build_content_filter_message_creates_error_with_partial(): + resp = _mock_response(content="partial answer", finish_reason="content_filter") + message = build_content_filter_message(response=resp, request=_request_piece(), partial_content="partial answer") + piece = message.message_pieces[0] + assert piece.converted_value_data_type == "error" + assert piece.prompt_metadata["partial_content"] == "partial answer" + + +# --------------------------------------------------------------------------- +# audio helpers +# --------------------------------------------------------------------------- + + +def _audio_piece(role="user"): + return MessagePiece( + role=role, conversation_id="c", original_value="clip.wav", original_value_data_type="audio_path" + ) + + +def test_should_skip_audio_piece_non_audio_type_false(): + assert ( + should_skip_audio_piece( + message_piece=_request_piece(), + is_last_message=False, + has_text_piece=True, + prefer_transcript_for_history=True, + ) + is False + ) + + +def test_should_skip_audio_piece_assistant_always_skipped(): + assert ( + should_skip_audio_piece( + message_piece=_audio_piece(role="assistant"), + is_last_message=True, + has_text_piece=False, + prefer_transcript_for_history=False, + ) + is True + ) + + +def test_should_skip_audio_piece_user_history_with_transcript_skipped(): + assert ( + should_skip_audio_piece( + message_piece=_audio_piece(), + is_last_message=False, + has_text_piece=True, + prefer_transcript_for_history=True, + ) + is True + ) + + +def test_should_skip_audio_piece_current_user_message_kept(): + assert ( + should_skip_audio_piece( + message_piece=_audio_piece(), + is_last_message=True, + has_text_piece=True, + prefer_transcript_for_history=True, + ) + is False + ) + + +async def test_build_multimodal_chat_messages_includes_audio(): + text_piece = MessagePiece(role="user", conversation_id="c", original_value="hi", original_value_data_type="text") + message = Message(message_pieces=[text_piece, _audio_piece()]) + with patch( + "pyrit.prompt_target.common.chat_completions_message_builder.build_audio_content_entry_async", + new=AsyncMock(return_value={"type": "input_audio", "input_audio": {"data": "x", "format": "wav"}}), + ): + result = await build_multimodal_chat_messages_async([message], prefer_transcript_for_history=False) + content_types = [part["type"] for part in result[0]["content"]] + assert content_types == ["text", "input_audio"] + + +async def test_save_audio_response_async_wav(): + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: + serializer = MagicMock() + serializer.value = "/path/audio.wav" + serializer.save_data = AsyncMock() + mock_factory.return_value = serializer + + result = await save_audio_response_async( + audio_data_base64=base64.b64encode(b"abc").decode("utf-8"), audio_format="wav" + ) + + mock_factory.assert_called_once_with(category="prompt-memory-entries", data_type="audio_path", extension=".wav") + serializer.save_data.assert_awaited_once_with(b"abc") + assert result == "/path/audio.wav" + + +async def test_save_audio_response_async_pcm16_wraps_wav(): + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: + serializer = MagicMock() + serializer.value = "/path/audio.wav" + serializer.save_formatted_audio = AsyncMock() + mock_factory.return_value = serializer + + result = await save_audio_response_async( + audio_data_base64=base64.b64encode(b"pcmdata").decode("utf-8"), audio_format="pcm16" + ) + + mock_factory.assert_called_once_with(category="prompt-memory-entries", data_type="audio_path", extension=".wav") + serializer.save_formatted_audio.assert_awaited_once_with( + data=b"pcmdata", num_channels=1, sample_width=2, sample_rate=24000 + ) + assert result == "/path/audio.wav" + + +async def test_build_audio_pieces_async_transcript_and_file(): + message = MagicMock() + message.audio.transcript = "the transcript" + message.audio.data = base64.b64encode(b"audio").decode("utf-8") + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: + serializer = MagicMock() + serializer.value = "/path/audio.wav" + serializer.save_data = AsyncMock() + mock_factory.return_value = serializer + + pieces = await build_audio_pieces_async(message=message, request=_request_piece(), audio_format="wav") + + assert [p.converted_value_data_type for p in pieces] == ["text", "audio_path"] + assert pieces[0].converted_value == "the transcript" + assert pieces[0].prompt_metadata.get("transcription") == "audio" + assert pieces[1].converted_value == "/path/audio.wav" + + +async def test_build_response_pieces_async_orders_text_audio_tool(): + tool_call = MagicMock() + tool_call.id = "call_1" + tool_call.function.name = "fn" + tool_call.function.arguments = "{}" + resp = _mock_response(content="text answer", tool_calls=[tool_call]) + resp.choices[0].message.audio = MagicMock() + resp.choices[0].message.audio.transcript = "spoken" + resp.choices[0].message.audio.data = base64.b64encode(b"audio").decode("utf-8") + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: + serializer = MagicMock() + serializer.value = "/path/audio.wav" + serializer.save_data = AsyncMock() + mock_factory.return_value = serializer + + 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"] diff --git a/tests/unit/prompt_target/target/test_litellm_chat_target.py b/tests/unit/prompt_target/target/test_litellm_chat_target.py new file mode 100644 index 0000000000..4ae4efa7bb --- /dev/null +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -0,0 +1,677 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import base64 +import json +import sys +import types +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.exceptions import ( + EmptyResponseException, + PyritException, + RateLimitException, + get_retry_max_num_attempts, +) +from pyrit.models import Message, MessagePiece +from pyrit.models.json_response_config import _JsonResponseConfig +from pyrit.prompt_target import ( + OpenAIChatAudioConfig, + TargetCapabilities, + TargetConfiguration, +) +from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget + +# --------------------------------------------------------------------------- +# LiteLLM stub +# --------------------------------------------------------------------------- + + +class _StubLiteLLMError(Exception): + def __init__(self, message: str = "", **kwargs: object) -> None: + super().__init__(message) + self.status_code = kwargs.get("status_code") + + +_EXCEPTION_NAMES = [ + "RateLimitError", + "APIConnectionError", + "Timeout", + "AuthenticationError", + "InternalServerError", + "ServiceUnavailableError", + "BadRequestError", + "ContentPolicyViolationError", +] + + +def _make_litellm_stub( + *, + supports_vision: bool = True, + supports_response_schema: bool = False, + supports_audio_input: bool = False, + supports_audio_output: bool = False, +): + mod = types.ModuleType("litellm") + mod.acompletion = AsyncMock(name="litellm.acompletion") + mod.supports_vision = MagicMock(return_value=supports_vision) + mod.supports_response_schema = MagicMock(return_value=supports_response_schema) + mod.supports_audio_input = MagicMock(return_value=supports_audio_input) + mod.supports_audio_output = MagicMock(return_value=supports_audio_output) + mod.get_supported_openai_params = MagicMock( + return_value=["temperature", "top_p", "max_tokens", "response_format", "seed", "n", "stop"] + ) + + exc_mod = types.ModuleType("litellm.exceptions") + for name in _EXCEPTION_NAMES: + setattr(exc_mod, name, type(name, (_StubLiteLLMError,), {"__module__": "litellm.exceptions"})) + mod.exceptions = exc_mod + return mod, exc_mod + + +@pytest.fixture +def litellm_stub(): + mod, exc_mod = _make_litellm_stub() + with patch.dict(sys.modules, {"litellm": mod, "litellm.exceptions": exc_mod}): + yield mod + + +@pytest.fixture +def target(patch_central_database, litellm_stub) -> LiteLLMChatTarget: + return LiteLLMChatTarget(model_name="anthropic/claude-sonnet-4-6") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _mock_response(content="hello", finish_reason="stop", model="anthropic/claude-sonnet-4-6"): + resp = MagicMock() + resp.choices = [MagicMock()] + resp.choices[0].finish_reason = finish_reason + resp.choices[0].message.content = content + resp.choices[0].message.tool_calls = None + resp.choices[0].message.audio = None + resp.model = model + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.total_tokens = 15 + resp.usage.cached_tokens = 0 + resp.model_dump_json = MagicMock(return_value=json.dumps({"finish_reason": finish_reason, "content": content})) + return resp + + +def _mock_audio_response(transcript="hello there", model="openai/gpt-4o-audio-preview"): + resp = _mock_response(content=None, model=model) + resp.choices[0].message.content = None + audio = MagicMock() + audio.transcript = transcript + audio.data = base64.b64encode(b"fake audio bytes").decode("utf-8") + resp.choices[0].message.audio = audio + return resp + + +def _mock_tool_call_response(): + resp = _mock_response(content=None) + resp.choices[0].message.content = None + tool_call = MagicMock() + tool_call.id = "call_123" + tool_call.function.name = "get_weather" + tool_call.function.arguments = '{"location": "SF"}' + resp.choices[0].message.tool_calls = [tool_call] + return resp + + +def _user_message(text="test prompt", conversation_id="convo"): + piece = MessagePiece( + role="user", + conversation_id=conversation_id, + original_value=text, + original_value_data_type="text", + ) + return piece.to_message() + + +def _disabled_json_config() -> _JsonResponseConfig: + return _JsonResponseConfig.from_metadata(metadata={}) + + +# --------------------------------------------------------------------------- +# Constructor +# --------------------------------------------------------------------------- + + +def test_init_requires_model_name(patch_central_database, litellm_stub): + with pytest.raises(ValueError, match="model_name is required"): + LiteLLMChatTarget() + + +def test_init_reads_model_env_var(patch_central_database, litellm_stub): + with patch.dict("os.environ", {"LITELLM_MODEL": "openai/gpt-4o"}): + t = LiteLLMChatTarget() + assert t._model_name == "openai/gpt-4o" + + +def test_init_explicit_model_overrides_env(patch_central_database, litellm_stub): + with patch.dict("os.environ", {"LITELLM_MODEL": "openai/gpt-4o"}): + t = LiteLLMChatTarget(model_name="anthropic/claude-haiku-4-5") + assert t._model_name == "anthropic/claude-haiku-4-5" + + +def test_init_api_key_from_env(patch_central_database, litellm_stub): + with patch.dict("os.environ", {"LITELLM_API_KEY": "sk-env"}): + t = LiteLLMChatTarget(model_name="openai/gpt-4o") + assert t._api_key == "sk-env" + + +def test_init_endpoint_from_env(patch_central_database, litellm_stub): + with patch.dict("os.environ", {"LITELLM_ENDPOINT": "http://localhost:4000"}): + t = LiteLLMChatTarget(model_name="openai/gpt-4o") + assert t._endpoint == "http://localhost:4000" + + +def test_drop_params_always_set_in_body(target): + body = target._construct_request_body( + messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config() + ) + assert body["drop_params"] is True + + +def test_drop_params_can_be_overridden_via_extra_body(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", extra_body_parameters={"drop_params": False}) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["drop_params"] is False + + +def test_num_retries_default_from_pyrit_convention(target): + assert target._num_retries == max(get_retry_max_num_attempts() - 1, 0) + + +def test_init_rejects_out_of_range_temperature(patch_central_database, litellm_stub): + with pytest.raises((ValueError, PyritException)): + LiteLLMChatTarget(model_name="openai/gpt-4o", temperature=5) + + +# --------------------------------------------------------------------------- +# Capability derivation +# --------------------------------------------------------------------------- + + +def test_capabilities_vision_model_includes_image(target): + supported = {t for combo in target.capabilities.input_modalities for t in combo} + assert "image_path" in supported + + +def test_capabilities_text_only_model_excludes_image(patch_central_database): + mod, exc_mod = _make_litellm_stub(supports_vision=False) + with patch.dict(sys.modules, {"litellm": mod, "litellm.exceptions": exc_mod}): + t = LiteLLMChatTarget(model_name="text/only-model") + supported = {ty for combo in t.capabilities.input_modalities for ty in combo} + assert supported == {"text"} + + +def test_capabilities_json_output_derived_from_supported_params(target): + # stub get_supported_openai_params includes "response_format" + assert target.capabilities.supports_json_output is True + + +def test_capabilities_audio_model_includes_audio_modalities(patch_central_database): + mod, exc_mod = _make_litellm_stub(supports_audio_input=True, supports_audio_output=True) + with patch.dict(sys.modules, {"litellm": mod, "litellm.exceptions": exc_mod}): + t = LiteLLMChatTarget(model_name="openai/gpt-4o-audio-preview") + input_types = {ty for combo in t.capabilities.input_modalities for ty in combo} + output_types = {ty for combo in t.capabilities.output_modalities for ty in combo} + assert "audio_path" in input_types + assert "audio_path" in output_types + + +def test_capabilities_text_only_model_excludes_audio(target): + output_types = {ty for combo in target.capabilities.output_modalities for ty in combo} + input_types = {ty for combo in target.capabilities.input_modalities for ty in combo} + assert output_types == {"text"} + assert "audio_path" not in input_types + + +def test_custom_configuration_overrides_derivation(patch_central_database, litellm_stub): + custom = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + t = LiteLLMChatTarget(model_name="openai/gpt-4o", custom_configuration=custom) + supported = {ty for combo in t.capabilities.input_modalities for ty in combo} + assert supported == {"text"} + assert t.capabilities.supports_json_output is False + + +# --------------------------------------------------------------------------- +# Identifier +# --------------------------------------------------------------------------- + + +def test_identifier_includes_behavioral_params_and_excludes_key(patch_central_database, litellm_stub): + t = LiteLLMChatTarget( + model_name="openai/gpt-4o", + api_key="sk-secret", + endpoint="http://localhost:4000", + temperature=0.5, + ) + params = t.get_identifier().params + assert params["temperature"] == 0.5 + assert params["endpoint"] == "http://localhost:4000" + assert not any("key" in key.lower() for key in params) + assert "sk-secret" not in json.dumps(params) + + +# --------------------------------------------------------------------------- +# Request body +# --------------------------------------------------------------------------- + + +def test_construct_request_body_basics(target): + messages = [{"role": "user", "content": "hi"}] + body = target._construct_request_body(messages=messages, json_config=_disabled_json_config()) + assert body["model"] == "anthropic/claude-sonnet-4-6" + assert body["messages"] == messages + assert body["drop_params"] is True + assert body["num_retries"] == target._num_retries + + +def test_construct_request_body_forwards_api_key(target): + body = target._construct_request_body( + messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config(), api_key="sk-x" + ) + assert body["api_key"] == "sk-x" + + +def test_construct_request_body_omits_none_values(target): + body = target._construct_request_body( + messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config() + ) + assert "api_key" not in body + assert "temperature" not in body + assert "response_format" not in body + + +def test_construct_request_body_forwards_optional_params(patch_central_database, litellm_stub): + t = LiteLLMChatTarget( + model_name="openai/gpt-4o", + temperature=0.5, + top_p=0.9, + max_tokens=100, + frequency_penalty=0.1, + presence_penalty=0.2, + seed=7, + n=2, + stop=["END"], + endpoint="http://localhost:4000", + ) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["temperature"] == 0.5 + assert body["top_p"] == 0.9 + assert body["max_tokens"] == 100 + assert body["frequency_penalty"] == 0.1 + assert body["presence_penalty"] == 0.2 + assert body["seed"] == 7 + assert body["n"] == 2 + assert body["stop"] == ["END"] + assert body["api_base"] == "http://localhost:4000" + + +def test_construct_request_body_forwards_max_completion_tokens(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/o1", max_completion_tokens=512) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["max_completion_tokens"] == 512 + assert "max_tokens" not in body + + +def test_max_tokens_and_max_completion_tokens_mutually_exclusive(patch_central_database, litellm_stub): + with pytest.raises(ValueError, match="Cannot provide both max_tokens and max_completion_tokens"): + LiteLLMChatTarget(model_name="openai/gpt-4o", max_tokens=10, max_completion_tokens=10) + + +def test_construct_request_body_passthrough_extra_params(patch_central_database, litellm_stub): + t = LiteLLMChatTarget( + model_name="openai/gpt-4o", + extra_body_parameters={"tools": [{"type": "function"}], "tool_choice": "auto"}, + ) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["tools"] == [{"type": "function"}] + assert body["tool_choice"] == "auto" + + +def test_construct_request_body_forwards_headers(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", headers={"X-Trace": "abc"}) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["extra_headers"] == {"X-Trace": "abc"} + + +# --------------------------------------------------------------------------- +# Send prompt (through the public API) +# --------------------------------------------------------------------------- + + +async def test_send_prompt_returns_text_response(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_response("The answer is 4.")) + + result = await target.send_prompt_async(message=_user_message("What is 2+2?")) + + assert len(result) == 1 + assert result[0].message_pieces[0].converted_value == "The answer is 4." + litellm_stub.acompletion.assert_awaited_once() + call_kwargs = litellm_stub.acompletion.call_args.kwargs + assert call_kwargs["model"] == "anthropic/claude-sonnet-4-6" + assert call_kwargs["num_retries"] == target._num_retries + + +async def test_send_prompt_handles_tool_calls(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_tool_call_response()) + + result = await target.send_prompt_async(message=_user_message("What's the weather?")) + + piece = result[0].message_pieces[0] + assert piece.converted_value_data_type == "function_call" + parsed = json.loads(piece.converted_value) + assert parsed["function"]["name"] == "get_weather" + + +async def test_send_prompt_captures_token_usage(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_response("ok")) + + result = await target.send_prompt_async(message=_user_message("hi")) + + metadata = result[0].message_pieces[0].prompt_metadata + assert metadata["token_usage_input_tokens"] == 10 + assert metadata["token_usage_output_tokens"] == 5 + assert metadata["token_usage_total_tokens"] == 15 + + +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} + litellm_stub.acompletion = AsyncMock(return_value=response) + + result = await target.send_prompt_async(message=_user_message("hi")) + + metadata = result[0].message_pieces[0].prompt_metadata + assert float(metadata["token_usage_cost"]) == pytest.approx(0.00042) + + +async def test_send_prompt_captures_response_cost_falls_back_to_completion_cost(target, litellm_stub): + response = _mock_response("ok") + # No usable _hidden_params dict -> should recompute via litellm.completion_cost. + response._hidden_params = None + litellm_stub.completion_cost = MagicMock(return_value=0.0009) + litellm_stub.acompletion = AsyncMock(return_value=response) + + result = await target.send_prompt_async(message=_user_message("hi")) + + metadata = result[0].message_pieces[0].prompt_metadata + assert float(metadata["token_usage_cost"]) == pytest.approx(0.0009) + litellm_stub.completion_cost.assert_called_once() + + +async def test_send_prompt_omits_cost_when_unavailable(target, litellm_stub): + response = _mock_response("ok") + response._hidden_params = None + litellm_stub.completion_cost = MagicMock(side_effect=Exception("no pricing map")) + litellm_stub.acompletion = AsyncMock(return_value=response) + + result = await target.send_prompt_async(message=_user_message("hi")) + + metadata = result[0].message_pieces[0].prompt_metadata + assert "token_usage_cost" not in metadata + + +async def test_send_prompt_resolves_callable_api_key(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key=lambda: "token-abc") + litellm_stub.acompletion = AsyncMock(return_value=_mock_response("ok")) + + await t.send_prompt_async(message=_user_message("hi")) + + assert litellm_stub.acompletion.call_args.kwargs["api_key"] == "token-abc" + + +# --------------------------------------------------------------------------- +# Multimodal message building +# --------------------------------------------------------------------------- + + +async def test_build_chat_messages_text_only_uses_string_content(target): + messages = [_user_message("hello")] + result = await target._build_chat_messages_async(messages) + assert result[0]["role"] == "user" + assert result[0]["content"] == "hello" + + +async def test_build_chat_messages_multimodal_text_and_image(target): + text_piece = MessagePiece( + role="user", conversation_id="c", original_value="describe", original_value_data_type="text" + ) + image_piece = MessagePiece( + role="user", conversation_id="c", original_value="x.png", original_value_data_type="image_path" + ) + message = Message(message_pieces=[text_piece, image_piece]) + + with patch( + "pyrit.prompt_target.common.chat_completions_message_builder.build_image_content_entry_async", + new=AsyncMock(return_value={"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}}), + ): + result = await target._build_chat_messages_async([message]) + + content_types = [part["type"] for part in result[0]["content"]] + assert "text" in content_types + assert "image_url" in content_types + + +async def test_build_chat_messages_rejects_unsupported_type(target): + piece = MessagePiece( + role="user", conversation_id="c", original_value="v.mp4", original_value_data_type="video_path" + ) + text_piece = MessagePiece(role="user", conversation_id="c", original_value="t", original_value_data_type="text") + message = Message(message_pieces=[text_piece, piece]) + with pytest.raises(ValueError, match="Multimodal data type video_path is not yet supported"): + await target._build_chat_messages_async([message]) + + +async def test_build_chat_messages_multimodal_text_and_audio(target): + text_piece = MessagePiece( + role="user", conversation_id="c", original_value="transcribe", original_value_data_type="text" + ) + audio_piece = MessagePiece( + role="user", conversation_id="c", original_value="clip.wav", original_value_data_type="audio_path" + ) + message = Message(message_pieces=[text_piece, audio_piece]) + + with patch( + "pyrit.prompt_target.common.chat_completions_message_builder.build_audio_content_entry_async", + new=AsyncMock(return_value={"type": "input_audio", "input_audio": {"data": "xxx", "format": "wav"}}), + ): + result = await target._build_chat_messages_async([message]) + + content_types = [part["type"] for part in result[0]["content"]] + assert "text" in content_types + assert "input_audio" in content_types + + +def test_audio_response_config_adds_modalities_to_request_body(patch_central_database, litellm_stub): + t = LiteLLMChatTarget( + model_name="openai/gpt-4o-audio-preview", + audio_response_config=OpenAIChatAudioConfig(voice="alloy", audio_format="wav"), + ) + body = t._construct_request_body(messages=[{"role": "user", "content": "hi"}], json_config=_disabled_json_config()) + assert body["modalities"] == ["text", "audio"] + assert body["audio"] == {"voice": "alloy", "format": "wav"} + + +async def test_send_prompt_parses_audio_response(patch_central_database, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_audio_response()) + target = LiteLLMChatTarget( + model_name="openai/gpt-4o-audio-preview", + audio_response_config=OpenAIChatAudioConfig(voice="alloy", audio_format="wav"), + ) + piece = MessagePiece(role="user", conversation_id="c", original_value="say hi", original_value_data_type="text") + + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: + mock_serializer = MagicMock() + mock_serializer.value = "/path/to/audio.wav" + mock_serializer.save_data = AsyncMock() + mock_factory.return_value = mock_serializer + + responses = await target.send_prompt_async(message=piece.to_message()) + + pieces = responses[0].message_pieces + data_types = [p.converted_value_data_type for p in pieces] + assert "text" in data_types + assert "audio_path" in data_types + transcript_piece = next(p for p in pieces if p.converted_value_data_type == "text") + assert transcript_piece.prompt_metadata.get("transcription") == "audio" + audio_piece = next(p for p in pieces if p.converted_value_data_type == "audio_path") + assert audio_piece.converted_value == "/path/to/audio.wav" + + +# --------------------------------------------------------------------------- +# JSON output +# --------------------------------------------------------------------------- + + +async def test_send_prompt_emits_response_format_for_json(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_response("{}")) + piece = MessagePiece( + role="user", + conversation_id="c", + original_value="give json", + original_value_data_type="text", + prompt_metadata={"response_format": "json"}, + ) + await target.send_prompt_async(message=piece.to_message()) + assert litellm_stub.acompletion.call_args.kwargs["response_format"] == {"type": "json_object"} + + +# --------------------------------------------------------------------------- +# Content filtering (surfaced as error Message, not raised) +# --------------------------------------------------------------------------- + + +async def test_content_filter_finish_reason_returns_error_message(target, litellm_stub): + litellm_stub.acompletion = AsyncMock( + return_value=_mock_response(content="partial answer", finish_reason="content_filter") + ) + + result = await target.send_prompt_async(message=_user_message("bad prompt")) + + piece = result[0].message_pieces[0] + assert piece.converted_value_data_type == "error" + assert piece.prompt_metadata.get("partial_content") == "partial answer" + + +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) + + result = await target.send_prompt_async(message=_user_message("bad prompt")) + + assert result[0].message_pieces[0].converted_value_data_type == "error" + + +# --------------------------------------------------------------------------- +# Empty / malformed responses +# --------------------------------------------------------------------------- + + +async def test_empty_response_raises(target, litellm_stub): + empty = _mock_response(content=None) + empty.choices[0].message.content = None + empty.choices[0].message.tool_calls = None + empty.choices[0].message.audio = None + litellm_stub.acompletion = AsyncMock(return_value=empty) + + with pytest.raises(EmptyResponseException): + await target.send_prompt_async(message=_user_message()) + + +async def test_no_choices_raises_pyrit_exception(target, litellm_stub): + bad = MagicMock() + bad.choices = [] + litellm_stub.acompletion = AsyncMock(return_value=bad) + + with pytest.raises(PyritException, match="No choices"): + await target.send_prompt_async(message=_user_message()) + + +async def test_unknown_finish_reason_raises(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_response(finish_reason="banana")) + + with pytest.raises(PyritException, match="Unknown finish_reason"): + await target.send_prompt_async(message=_user_message()) + + +# --------------------------------------------------------------------------- +# Exception translation +# --------------------------------------------------------------------------- + + +async def test_rate_limit_error_translated(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(side_effect=litellm_stub.exceptions.RateLimitError("rate limited")) + with pytest.raises(RateLimitException, match="Rate limited"): + await target.send_prompt_async(message=_user_message()) + + +async def test_auth_error_translated(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(side_effect=litellm_stub.exceptions.AuthenticationError("bad key")) + with pytest.raises(PyritException, match="Authentication failed"): + await target.send_prompt_async(message=_user_message()) + + +async def test_connection_error_translated_to_transient(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(side_effect=litellm_stub.exceptions.APIConnectionError("reset")) + with pytest.raises(RateLimitException, match="Transient"): + await target.send_prompt_async(message=_user_message()) + + +async def test_timeout_translated_to_transient(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(side_effect=litellm_stub.exceptions.Timeout("timed out")) + with pytest.raises(RateLimitException, match="Transient"): + await target.send_prompt_async(message=_user_message()) + + +async def test_unknown_error_wrapped_in_pyrit_exception(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(side_effect=RuntimeError("something broke")) + with pytest.raises(PyritException, match="LiteLLM error"): + await target.send_prompt_async(message=_user_message()) + + +# --------------------------------------------------------------------------- +# API key resolution +# --------------------------------------------------------------------------- + + +async def test_resolve_api_key_string(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key="sk-x") + assert await t._resolve_api_key_async() == "sk-x" + + +async def test_resolve_api_key_none(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o") + assert await t._resolve_api_key_async() is None + + +async def test_resolve_api_key_sync_callable(patch_central_database, litellm_stub): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key=lambda: "tok") + assert await t._resolve_api_key_async() == "tok" + + +async def test_resolve_api_key_async_callable(patch_central_database, litellm_stub): + async def provider() -> str: + return "atok" + + t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key=provider) + assert await t._resolve_api_key_async() == "atok" + + +# --------------------------------------------------------------------------- +# Misc +# --------------------------------------------------------------------------- + + +def test_is_json_response_supported_reflects_capability(target): + assert target.is_json_response_supported() is True 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 0540bb68ca..2ef46121cf 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -1500,7 +1500,7 @@ async def test_save_audio_response_async_wav_format(patch_central_database): audio_bytes = b"fake wav audio data" audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.wav" mock_serializer.save_data = AsyncMock() @@ -1535,7 +1535,7 @@ async def test_save_audio_response_async_mp3_format(patch_central_database): audio_bytes = b"fake mp3 audio data" audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.mp3" mock_serializer.save_data = AsyncMock() @@ -1570,7 +1570,7 @@ async def test_save_audio_response_async_pcm16_format(patch_central_database): audio_bytes = b"\x00\x01\x02\x03" * 100 # Raw PCM16 data audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.wav" mock_serializer.save_formatted_audio = AsyncMock() @@ -1667,7 +1667,7 @@ async def test_save_audio_response_async_flac_format(patch_central_database): audio_bytes = b"fake flac audio data" audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.flac" mock_serializer.save_data = AsyncMock() @@ -1698,7 +1698,7 @@ async def test_save_audio_response_async_opus_format(patch_central_database): audio_bytes = b"fake opus audio data" audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.opus" mock_serializer.save_data = AsyncMock() @@ -1728,7 +1728,7 @@ async def test_save_audio_response_async_no_config_defaults_to_wav(patch_central audio_bytes = b"fake audio data" audio_data_base64 = base64.b64encode(audio_bytes).decode("utf-8") - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/saved/audio.wav" mock_serializer.save_data = AsyncMock() @@ -1769,7 +1769,7 @@ async def test_construct_message_from_response_audio_transcript_has_metadata( mock_response.choices[0].message.audio.data = base64.b64encode(b"fake audio").decode("utf-8") mock_response.choices[0].message.tool_calls = None - with patch("pyrit.prompt_target.openai.openai_chat_target.data_serializer_factory") as mock_factory: + with patch("pyrit.prompt_target.common.chat_completions_response_parser.data_serializer_factory") as mock_factory: mock_serializer = MagicMock() mock_serializer.value = "/path/to/audio.wav" mock_serializer.save_data = AsyncMock() @@ -2003,21 +2003,22 @@ async def test_construct_message_from_response_captures_token_usage( ): """Test that token usage from the API response is stored in prompt_metadata.""" mock_response = create_mock_completion(content="Hello") - mock_response.model = "gpt-4o-2024-05-13" mock_response.usage = MagicMock() mock_response.usage.prompt_tokens = 10 mock_response.usage.completion_tokens = 20 mock_response.usage.total_tokens = 30 - mock_response.usage.cached_tokens = 5 + # cached_tokens is nested under prompt_tokens_details in the real API, not top-level. + mock_response.usage.prompt_tokens_details.cached_tokens = 5 + mock_response.usage.completion_tokens_details.reasoning_tokens = 7 result = await target._construct_message_from_response(mock_response, dummy_text_message_piece) piece = result.message_pieces[0] - assert piece.prompt_metadata["token_usage_model_name"] == "gpt-4o-2024-05-13" - assert piece.prompt_metadata["token_usage_prompt_tokens"] == 10 - assert piece.prompt_metadata["token_usage_completion_tokens"] == 20 + assert piece.prompt_metadata["token_usage_input_tokens"] == 10 + assert piece.prompt_metadata["token_usage_output_tokens"] == 20 assert piece.prompt_metadata["token_usage_total_tokens"] == 30 assert piece.prompt_metadata["token_usage_cached_tokens"] == 5 + assert piece.prompt_metadata["token_usage_reasoning_tokens"] == 7 async def test_construct_message_from_response_no_usage_no_metadata( @@ -2030,29 +2031,26 @@ async def test_construct_message_from_response_no_usage_no_metadata( result = await target._construct_message_from_response(mock_response, dummy_text_message_piece) piece = result.message_pieces[0] - assert "token_usage_model_name" not in piece.prompt_metadata - assert "token_usage_prompt_tokens" not in piece.prompt_metadata + assert "token_usage_input_tokens" not in piece.prompt_metadata -async def test_construct_message_from_response_token_usage_defaults_on_missing_attrs( +async def test_construct_message_from_response_token_usage_omits_missing_attrs( target: OpenAIChatTarget, dummy_text_message_piece: MessagePiece ): - """Test that missing usage attributes default to 0 and missing model defaults to 'unknown'.""" + """Test that fields the provider does not report are omitted rather than defaulted.""" mock_response = create_mock_completion(content="Hello") - # Create a usage object without cached_tokens + # Create a usage object without cached/reasoning detail breakdowns. mock_usage = MagicMock(spec=[]) mock_usage.prompt_tokens = 5 mock_usage.completion_tokens = 10 mock_usage.total_tokens = 15 mock_response.usage = mock_usage - # Remove model attribute to test default - del mock_response.model result = await target._construct_message_from_response(mock_response, dummy_text_message_piece) piece = result.message_pieces[0] - assert piece.prompt_metadata["token_usage_model_name"] == "unknown" - assert piece.prompt_metadata["token_usage_prompt_tokens"] == 5 - assert piece.prompt_metadata["token_usage_completion_tokens"] == 10 + assert piece.prompt_metadata["token_usage_input_tokens"] == 5 + assert piece.prompt_metadata["token_usage_output_tokens"] == 10 assert piece.prompt_metadata["token_usage_total_tokens"] == 15 - assert piece.prompt_metadata["token_usage_cached_tokens"] == 0 + assert "token_usage_cached_tokens" not in piece.prompt_metadata + assert "token_usage_reasoning_tokens" not in piece.prompt_metadata