From 42d212b3d5483f2df9cc96ba59d9a3b072451421 Mon Sep 17 00:00:00 2001 From: RheagalFire Date: Thu, 9 Jul 2026 22:57:57 +0530 Subject: [PATCH 1/4] feat: add LiteLLM as AI gateway prompt target --- pyproject.toml | 3 + pyrit/prompt_target/__init__.py | 2 + pyrit/prompt_target/litellm_chat_target.py | 270 +++++++++++++++++ .../target/test_litellm_chat_target.py | 286 ++++++++++++++++++ 4 files changed, 561 insertions(+) create mode 100644 pyrit/prompt_target/litellm_chat_target.py create mode 100644 tests/unit/prompt_target/target/test_litellm_chat_target.py diff --git a/pyproject.toml b/pyproject.toml index bdc563d000..499db5db06 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 = [ diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 82f897c156..222ff5325a 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -28,6 +28,7 @@ from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute +from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget from pyrit.prompt_target.http_target.http_target import HTTPTarget from pyrit.prompt_target.http_target.http_target_callback_functions import ( @@ -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/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py new file mode 100644 index 0000000000..a34d55ae4e --- /dev/null +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -0,0 +1,270 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +import logging +from collections.abc import MutableSequence +from typing import Any, Optional + +from pyrit.exceptions import ( + EmptyResponseException, + PyritException, + RateLimitException, + pyrit_target_retry, +) +from pyrit.models import ( + ChatMessage, + Message, + MessagePiece, + construct_response_from_request, +) +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 + +logger = logging.getLogger(__name__) + + +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. + + Install the optional dependency:: + + pip install pyrit[litellm] + + Example:: + + target = LiteLLMChatTarget( + model_name="anthropic/claude-sonnet-4-6", + ) + + 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. + + 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. When omitted, LiteLLM reads provider-specific + environment variables. + api_base: Optional base URL override (e.g. for a self-hosted proxy). + temperature: Sampling temperature (0-2). + top_p: Nucleus sampling probability (0-1). + max_tokens: Maximum number of tokens to generate. + drop_params: When ``True`` (the default), LiteLLM silently drops + parameters that the target provider does not support, preventing + cross-provider compatibility errors. + """ + + _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_output=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + ) + ) + + def __init__( + self, + *, + model_name: Optional[str] = None, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + max_tokens: Optional[int] = None, + drop_params: bool = True, + max_requests_per_minute: Optional[int] = None, + custom_configuration: Optional[TargetConfiguration] = None, + ) -> None: + """ + Initialize a LiteLLMChatTarget. + + Raises: + ValueError: If model_name is not provided and LITELLM_MODEL env var is not set. + """ + import os + + 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." + ) + + validate_temperature(temperature) + validate_top_p(top_p) + + super().__init__( + model_name=resolved_model, + max_requests_per_minute=max_requests_per_minute, + custom_configuration=custom_configuration, + ) + + self._api_key = api_key + self._api_base = api_base + self._temperature = temperature + self._top_p = top_p + self._max_tokens = max_tokens + self._drop_params = drop_params + + @limit_requests_per_minute + @pyrit_target_retry + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + 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 + + message = normalized_conversation[-1] + request_piece: MessagePiece = message.message_pieces[0] + + logger.info(f"Sending prompt to LiteLLM target ({self._model_name}): {message}") + + messages = self._build_chat_messages(normalized_conversation) + body = self._construct_request_body(messages) + + try: + response = await litellm.acompletion(**body) + except Exception as exc: + self._translate_litellm_exception(exc) + + self._validate_response(response) + + return [self._construct_message_from_response(response, request_piece)] + + def _build_chat_messages(self, conversation: MutableSequence[Message]) -> list[dict[str, Any]]: + chat_messages: list[dict[str, Any]] = [] + for msg in conversation: + for piece in msg.message_pieces: + chat_message = ChatMessage(role=piece.api_role, content=piece.converted_value) + chat_messages.append(chat_message.model_dump(exclude_none=True)) + return chat_messages + + def _construct_request_body(self, messages: list[dict[str, Any]]) -> dict[str, Any]: + body: dict[str, Any] = { + "model": self._model_name, + "messages": messages, + "drop_params": self._drop_params, + } + + if self._api_key: + body["api_key"] = self._api_key + if self._api_base: + body["api_base"] = self._api_base + if self._temperature is not None: + body["temperature"] = self._temperature + if self._top_p is not None: + body["top_p"] = self._top_p + if self._max_tokens is not None: + body["max_tokens"] = self._max_tokens + + return body + + def _translate_litellm_exception(self, exc: Exception) -> None: + qualname = f"{type(exc).__module__}.{type(exc).__qualname__}" + + rate_limit_types = { + "litellm.exceptions.RateLimitError", + } + transient_types = { + "litellm.exceptions.APIConnectionError", + "litellm.exceptions.Timeout", + "litellm.exceptions.InternalServerError", + "litellm.exceptions.ServiceUnavailableError", + } + auth_types = { + "litellm.exceptions.AuthenticationError", + } + + if qualname in rate_limit_types: + raise RateLimitException( + message=f"Rate limited by provider: {exc}", + status_code=429, + ) from exc + + if qualname in auth_types: + raise PyritException(message=f"Authentication failed: {exc}") from exc + + if qualname in transient_types: + raise RateLimitException( + message=f"Transient provider error (will retry): {exc}", + status_code=getattr(exc, "status_code", 503), + ) from exc + + raise PyritException(message=f"LiteLLM error: {exc}") from exc + + def _validate_response(self, response: Any) -> None: + if not hasattr(response, "choices") or not response.choices: + raise EmptyResponseException(message="No choices returned from LiteLLM.") + + choice = response.choices[0] + finish_reason = getattr(choice, "finish_reason", None) + valid_reasons = {"stop", "length", "tool_calls", "content_filter"} + if finish_reason and finish_reason not in valid_reasons: + raise PyritException( + message=f"Unexpected finish_reason '{finish_reason}' from LiteLLM response." + ) + + content = getattr(choice.message, "content", None) if hasattr(choice, "message") else None + tool_calls = getattr(choice.message, "tool_calls", None) if hasattr(choice, "message") else None + + if not content and not tool_calls: + raise EmptyResponseException(message="LiteLLM returned an empty response.") + + def _construct_message_from_response(self, response: Any, request_piece: MessagePiece) -> Message: + choice = response.choices[0] + content = getattr(choice.message, "content", None) or "" + + pieces: list[MessagePiece] = [] + + if content: + text_msg = construct_response_from_request( + request=request_piece, + response_text_pieces=[content], + response_type="text", + ) + pieces.append(text_msg.message_pieces[0]) + + tool_calls = getattr(choice.message, "tool_calls", None) + if tool_calls: + 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, + }, + } + tool_msg = construct_response_from_request( + request=request_piece, + response_text_pieces=[json.dumps(tool_call_data)], + response_type="function_call", + ) + pieces.append(tool_msg.message_pieces[0]) + + if not pieces: + raise EmptyResponseException(message="Failed to extract any response content from LiteLLM.") + + 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) + + return Message(message_pieces=pieces) 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..24f2c0eae2 --- /dev/null +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -0,0 +1,286 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +import sys +import types +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from pyrit.exceptions.exception_classes import ( + EmptyResponseException, + PyritException, + RateLimitException, +) +from pyrit.models import Message, MessagePiece + + +def _make_litellm_stub(): + """Install a fake litellm module so the target can be imported without the real package.""" + mod = types.ModuleType("litellm") + mod.acompletion = AsyncMock(name="litellm.acompletion") + + exc_mod = types.ModuleType("litellm.exceptions") + exc_mod.RateLimitError = type("RateLimitError", (Exception,), {"__module__": "litellm.exceptions"}) + exc_mod.APIConnectionError = type("APIConnectionError", (Exception,), {"__module__": "litellm.exceptions"}) + exc_mod.Timeout = type("Timeout", (Exception,), {"__module__": "litellm.exceptions"}) + exc_mod.AuthenticationError = type("AuthenticationError", (Exception,), {"__module__": "litellm.exceptions"}) + exc_mod.InternalServerError = type("InternalServerError", (Exception,), {"__module__": "litellm.exceptions"}) + exc_mod.ServiceUnavailableError = type( + "ServiceUnavailableError", (Exception,), {"__module__": "litellm.exceptions"} + ) + + mod.exceptions = exc_mod + sys.modules["litellm"] = mod + sys.modules["litellm.exceptions"] = exc_mod + return mod + + +LITELLM_STUB = _make_litellm_stub() + +from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget # noqa: E402 + + +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.model = model + resp.usage = MagicMock() + resp.usage.prompt_tokens = 10 + resp.usage.completion_tokens = 5 + resp.usage.total_tokens = 15 + 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 _make_message(text="test prompt", role="user"): + piece = MessagePiece( + role=role, + conversation_id="test_convo", + original_value=text, + converted_value=text, + original_value_data_type="text", + converted_value_data_type="text", + ) + return Message(message_pieces=[piece]) + + +@pytest.fixture +def target(patch_central_database) -> LiteLLMChatTarget: + return LiteLLMChatTarget(model_name="anthropic/claude-sonnet-4-6") + + +# ── constructor tests ───────────────────────────────────────────────── + +def test_init_requires_model_name(patch_central_database): + with pytest.raises(ValueError, match="model_name is required"): + LiteLLMChatTarget() + + +def test_init_reads_env_var(patch_central_database, monkeypatch): + monkeypatch.setenv("LITELLM_MODEL", "openai/gpt-4o") + t = LiteLLMChatTarget() + assert t._model_name == "openai/gpt-4o" + + +def test_init_explicit_overrides_env(patch_central_database, monkeypatch): + monkeypatch.setenv("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_drop_params_defaults_true(target): + assert target._drop_params is True + + +def test_drop_params_can_be_disabled(patch_central_database): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", drop_params=False) + assert t._drop_params is False + + +# ── request body tests ──────────────────────────────────────────────── + +def test_construct_request_body_includes_drop_params(target): + messages = [{"role": "user", "content": "hi"}] + body = target._construct_request_body(messages) + assert body["drop_params"] is True + assert body["model"] == "anthropic/claude-sonnet-4-6" + assert body["messages"] == messages + + +def test_construct_request_body_forwards_api_key(patch_central_database): + t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key="sk-test") + body = t._construct_request_body([{"role": "user", "content": "hi"}]) + assert body["api_key"] == "sk-test" + + +def test_construct_request_body_omits_api_key_when_blank(target): + body = target._construct_request_body([{"role": "user", "content": "hi"}]) + assert "api_key" not in body + + +def test_construct_request_body_forwards_optional_params(patch_central_database): + t = LiteLLMChatTarget( + model_name="openai/gpt-4o", + temperature=0.5, + top_p=0.9, + max_tokens=100, + api_base="http://localhost:4000", + ) + body = t._construct_request_body([{"role": "user", "content": "hi"}]) + assert body["temperature"] == 0.5 + assert body["top_p"] == 0.9 + assert body["max_tokens"] == 100 + assert body["api_base"] == "http://localhost:4000" + + +# ── send prompt tests ──────────────────────────────────────────────── + +@pytest.mark.asyncio +async def test_send_prompt_returns_text_response(target): + LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_response("The answer is 4.")) + conversation = [_make_message("What is 2+2?")] + + result = await target._send_prompt_to_target_async(normalized_conversation=conversation) + + 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["drop_params"] is True + + +@pytest.mark.asyncio +async def test_send_prompt_handles_tool_calls(target): + LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_tool_call_response()) + conversation = [_make_message("What's the weather?")] + + result = await target._send_prompt_to_target_async(normalized_conversation=conversation) + + assert len(result) == 1 + 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" + + +@pytest.mark.asyncio +async def test_send_prompt_captures_token_usage(target): + LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_response("ok")) + conversation = [_make_message("hi")] + + result = await target._send_prompt_to_target_async(normalized_conversation=conversation) + + metadata = result[0].message_pieces[0].prompt_metadata + assert metadata["token_usage_prompt_tokens"] == 10 + assert metadata["token_usage_completion_tokens"] == 5 + assert metadata["token_usage_total_tokens"] == 15 + + +# ── empty/malformed response tests ─────────────────────────────────── + +@pytest.mark.asyncio +async def test_empty_response_raises(target): + empty_resp = _mock_response(content=None) + empty_resp.choices[0].message.content = None + empty_resp.choices[0].message.tool_calls = None + LITELLM_STUB.acompletion = AsyncMock(return_value=empty_resp) + + with pytest.raises(EmptyResponseException): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +@pytest.mark.asyncio +async def test_no_choices_raises(target): + bad_resp = MagicMock() + bad_resp.choices = [] + LITELLM_STUB.acompletion = AsyncMock(return_value=bad_resp) + + with pytest.raises(EmptyResponseException): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +# ── exception translation tests ────────────────────────────────────── + +@pytest.mark.asyncio +async def test_rate_limit_translated(target): + exc = LITELLM_STUB.exceptions.RateLimitError("rate limited") + LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + + with pytest.raises(RateLimitException, match="Rate limited"): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +@pytest.mark.asyncio +async def test_auth_error_translated(target): + exc = LITELLM_STUB.exceptions.AuthenticationError("bad key") + LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + + with pytest.raises(PyritException, match="Authentication failed"): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +@pytest.mark.asyncio +async def test_connection_error_translated_to_retryable(target): + exc = LITELLM_STUB.exceptions.APIConnectionError("connection reset") + LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + + with pytest.raises(RateLimitException, match="Transient"): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +@pytest.mark.asyncio +async def test_timeout_translated_to_retryable(target): + exc = LITELLM_STUB.exceptions.Timeout("timed out") + LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + + with pytest.raises(RateLimitException, match="Transient"): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +@pytest.mark.asyncio +async def test_unknown_error_wrapped_in_pyrit_exception(target): + LITELLM_STUB.acompletion = AsyncMock(side_effect=RuntimeError("something broke")) + + with pytest.raises(PyritException, match="LiteLLM error"): + await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) + + +# ── validation edge cases ──────────────────────────────────────────── + +def test_unexpected_finish_reason_raises(target): + resp = _mock_response(finish_reason="unexpected_value") + with pytest.raises(PyritException, match="Unexpected finish_reason"): + target._validate_response(resp) + + +def test_valid_finish_reasons_accepted(target): + for reason in ("stop", "length", "tool_calls", "content_filter"): + resp = _mock_response(finish_reason=reason) + target._validate_response(resp) + + +# ── message building tests ─────────────────────────────────────────── + +def test_build_chat_messages_preserves_roles(target): + msgs = [_make_message("hello", "user"), _make_message("hi there", "assistant")] + result = target._build_chat_messages(msgs) + assert result[0]["role"] == "user" + assert result[0]["content"] == "hello" + assert result[1]["role"] == "assistant" + assert result[1]["content"] == "hi there" From 4d413a1c7b79dbf5bb22ce67f957746428142dc3 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 10 Jul 2026 14:00:14 -0700 Subject: [PATCH 2/4] feat: harden LiteLLMChatTarget (multimodal, params, auth, capabilities, token usage) Extends the LiteLLM target for parity with OpenAIChatTarget and shares logic instead of reinventing it: - Extract shared Chat Completions helpers (chat_completions_message_builder, chat_completions_response_parser) used by both OpenAIChatTarget and LiteLLMChatTarget for request building and response parsing (text, image, audio, tool calls, content-filter handling). - Add multimodal support (image + audio input, audio output via audio_response_config) with capabilities derived from LiteLLM's model registry and a conservative text-only fallback. - Support the full OpenAI parameter set plus an extra_body_parameters passthrough; auth via sync/async token providers; identifiers that exclude the api_key; underlying_model capability lookup; and LiteLLM-owned retry (num_retries from PyRIT's global convention) to avoid double-retrying. - Add a provider-neutral TokenUsage value object (input/output/total/reasoning/cached + extra) and capture it for both targets; capture LiteLLM per-call cost. - Add litellm as an optional extra and include it in the 'all' extra. - Modernize type syntax (X | None), tidy docstrings per the style guide, and add unit + integration tests (image/audio on a gpt-5 deployment). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .env_example | 4 + pyproject.toml | 1 + pyrit/models/__init__.py | 2 + pyrit/models/token_usage.py | 199 ++++++ pyrit/prompt_target/__init__.py | 2 +- .../chat_completions_message_builder.py | 267 +++++++ .../chat_completions_response_parser.py | 369 ++++++++++ pyrit/prompt_target/litellm_chat_target.py | 639 ++++++++++++----- .../openai/openai_chat_target.py | 292 ++------ .../test_litellm_chat_target_integration.py | 239 +++++++ .../test_openai_chat_target_integration.py | 36 +- tests/unit/models/test_token_usage.py | 147 ++++ .../target/test_chat_completions_helpers.py | 347 +++++++++ .../target/test_litellm_chat_target.py | 667 ++++++++++++++---- .../target/test_openai_chat_target.py | 44 +- 15 files changed, 2656 insertions(+), 599 deletions(-) create mode 100644 pyrit/models/token_usage.py create mode 100644 pyrit/prompt_target/common/chat_completions_message_builder.py create mode 100644 pyrit/prompt_target/common/chat_completions_response_parser.py create mode 100644 tests/integration/targets/test_litellm_chat_target_integration.py create mode 100644 tests/unit/models/test_token_usage.py create mode 100644 tests/unit/prompt_target/target/test_chat_completions_helpers.py diff --git a/.env_example b/.env_example index b925fb097c..25cb650214 100644 --- a/.env_example +++ b/.env_example @@ -31,6 +31,10 @@ PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_GPT4O_MODEL="gpt-4o" +# Model overrides for the chat integration tests (OpenAIChatTarget / LiteLLMChatTarget). +# Optional: if unset, the integration tests default to "gpt-5" (chat/vision) and "gpt-audio" (audio). +PLATFORM_OPENAI_CHAT_MODEL="gpt-5" +PLATFORM_OPENAI_AUDIO_MODEL="gpt-audio" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately # Example: https://xxxx.openai.azure.com/openai/v1 diff --git a/pyproject.toml b/pyproject.toml index 499db5db06..eecbddbb9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,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..a903f803be --- /dev/null +++ b/pyrit/models/token_usage.py @@ -0,0 +1,199 @@ +# 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_" + +#: Attribute names a provider ``usage`` object may use for the input/output counts. Chat +#: Completions (OpenAI/LiteLLM) reports ``prompt_tokens``/``completion_tokens``; the Responses API +#: (and Anthropic/Gemini) reports ``input_tokens``/``output_tokens``. ``from_provider_usage`` +#: accepts either, but PyRIT normalizes everything to the ``input``/``output`` vocabulary. +_PROVIDER_INPUT_NAMES = ("input_tokens", "prompt_tokens") +_PROVIDER_OUTPUT_NAMES = ("output_tokens", "completion_tokens") + +#: 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"}) + + +def _first_int(source: Any, *names: str) -> int | None: + """ + Return the first attribute among ``names`` on ``source`` that is an int, else None. + + Booleans are rejected even though ``bool`` is a subclass of ``int``. + + Args: + source (Any): The object to read attributes from. + names (str): Candidate attribute names, tried in order. + + Returns: + int | None: The first integer value found, or None. + """ + for name in names: + value = getattr(source, name, None) + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + +@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; + :meth:`from_provider_usage` accepts either provider shape. The object is persisted onto a + :class:`~pyrit.models.MessagePiece`'s ``prompt_metadata`` via :meth:`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``. + + 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 :meth:`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_provider_usage(cls, usage: Any) -> TokenUsage: + """ + Build a :class:`TokenUsage` from a provider ``usage`` object (OpenAI/LiteLLM/Anthropic shape). + + Reads top-level counts (mapping ``prompt``->``input`` and ``completion``->``output``, and + accepting Responses-API ``input``/``output`` names directly) and the nested + ``*_tokens_details`` breakdowns. ``total_tokens`` is derived from input + output when the + provider omits it. Unmodeled detail counts (audio, predicted-output, cache-write) are + collected into ``extra`` with disambiguated names. + + Args: + usage (Any): The provider usage object. + + Returns: + TokenUsage: The parsed token usage. + """ + input_tokens = _first_int(usage, *_PROVIDER_INPUT_NAMES) + output_tokens = _first_int(usage, *_PROVIDER_OUTPUT_NAMES) + total_tokens = _first_int(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 + + input_details = getattr(usage, "prompt_tokens_details", None) or getattr(usage, "input_tokens_details", None) + output_details = getattr(usage, "completion_tokens_details", None) or getattr( + usage, "output_tokens_details", None + ) + + cached_tokens = _first_int(input_details, "cached_tokens") if input_details is not None else None + reasoning_tokens = _first_int(output_details, "reasoning_tokens") if output_details is not None else None + + extra: dict[str, int] = {} + if input_details is not None: + _put(extra, "input_audio_tokens", _first_int(input_details, "audio_tokens")) + _put(extra, "cache_write_tokens", _first_int(input_details, "cache_write_tokens", "cache_creation_tokens")) + if output_details is not None: + _put(extra, "output_audio_tokens", _first_int(output_details, "audio_tokens")) + _put(extra, "accepted_prediction_tokens", _first_int(output_details, "accepted_prediction_tokens")) + _put(extra, "rejected_prediction_tokens", _first_int(output_details, "rejected_prediction_tokens")) + + return cls( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + reasoning_tokens=reasoning_tokens, + cached_tokens=cached_tokens, + extra=extra, + ) + + @classmethod + def from_metadata(cls, metadata: dict[str, Any]) -> TokenUsage | None: + """ + Reconstruct a :class:`TokenUsage` from a :class:`~pyrit.models.MessagePiece`'s ``prompt_metadata``. + + Reads the ``token_usage_input_tokens`` / ``token_usage_output_tokens`` keys written by + :meth:`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 + + +def _put(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 diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 222ff5325a..55bbb5aaad 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -28,7 +28,6 @@ from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute -from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget from pyrit.prompt_target.http_target.http_target import HTTPTarget from pyrit.prompt_target.http_target.http_target_callback_functions import ( @@ -37,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 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..c42f85913a --- /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 :func:`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..e218b33dd6 --- /dev/null +++ b/pyrit/prompt_target/common/chat_completions_response_parser.py @@ -0,0 +1,369 @@ +# 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 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 provider usage through :class:`~pyrit.models.TokenUsage`, which reads both the + top-level counts and the nested ``prompt_tokens_details`` / ``completion_tokens_details`` + breakdowns (so cached and reasoning tokens are captured, not just prompt/completion/total). + 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 = TokenUsage.from_provider_usage(usage) + pieces[0].prompt_metadata.update(token_usage.to_metadata()) + + +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 index a34d55ae4e..0c97f3b539 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -1,270 +1,583 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import json +import itertools import logging -from collections.abc import MutableSequence -from typing import Any, Optional +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, - pyrit_target_retry, + get_retry_max_num_attempts, + handle_bad_request_exception, ) +from pyrit.identifiers import ComponentIdentifier from pyrit.models import ( - ChatMessage, Message, MessagePiece, - construct_response_from_request, + 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.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. + 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 + :mod:`pyrit.prompt_target.common.chat_completions_message_builder` and + :mod:`pyrit.prompt_target.common.chat_completions_response_parser`. Install the optional dependency:: pip install pyrit[litellm] - Example:: - - target = LiteLLMChatTarget( - model_name="anthropic/claude-sonnet-4-6", - ) - 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. + (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. When omitted, LiteLLM reads provider-specific - environment variables. - api_base: Optional base URL override (e.g. for a self-hosted proxy). + ``"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. - drop_params: When ``True`` (the default), LiteLLM silently drops - parameters that the target provider does not support, preventing - cross-provider compatibility errors. + 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_editable_history=True, - supports_json_output=True, supports_multi_message_pieces=True, + supports_editable_history=True, supports_system_prompt=True, + input_modalities=_TEXT_INPUT, ) ) def __init__( self, *, - model_name: Optional[str] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - max_tokens: Optional[int] = None, - drop_params: bool = True, - max_requests_per_minute: Optional[int] = None, - custom_configuration: Optional[TargetConfiguration] = None, + 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. + 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. """ - import os - 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." - ) + 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, ) - self._api_key = api_key - self._api_base = api_base + # 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._drop_params = drop_params - - @limit_requests_per_minute - @pyrit_target_retry - async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + 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]" + "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}") - messages = self._build_chat_messages(normalized_conversation) - body = self._construct_request_body(messages) + 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: - self._translate_litellm_exception(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), + ) + ] - self._validate_response(response) + validate_chat_completion_response(response=response) + return [await self._construct_message_from_response(response=response, request=request_piece)] - return [self._construct_message_from_response(response, request_piece)] + async def _resolve_api_key_async(self) -> str | None: + """ + Resolve the api_key to a concrete string, awaiting async token providers. - def _build_chat_messages(self, conversation: MutableSequence[Message]) -> list[dict[str, Any]]: - chat_messages: list[dict[str, Any]] = [] - for msg in conversation: - for piece in msg.message_pieces: - chat_message = ChatMessage(role=piece.api_role, content=piece.converted_value) - chat_messages.append(chat_message.model_dump(exclude_none=True)) - return chat_messages + 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]]) -> dict[str, Any]: + 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, - "drop_params": self._drop_params, + # 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), } - if self._api_key: - body["api_key"] = self._api_key - if self._api_base: - body["api_base"] = self._api_base - if self._temperature is not None: - body["temperature"] = self._temperature - if self._top_p is not None: - body["top_p"] = self._top_p - if self._max_tokens is not None: - body["max_tokens"] = self._max_tokens - - return body + # Passthrough for arbitrary provider params (may override defaults above, e.g. drop_params). + if self._extra_body_parameters: + body.update(self._extra_body_parameters) - def _translate_litellm_exception(self, exc: Exception) -> None: - qualname = f"{type(exc).__module__}.{type(exc).__qualname__}" + return {k: v for k, v in body.items() if v is not None} - rate_limit_types = { - "litellm.exceptions.RateLimitError", - } - transient_types = { - "litellm.exceptions.APIConnectionError", - "litellm.exceptions.Timeout", - "litellm.exceptions.InternalServerError", - "litellm.exceptions.ServiceUnavailableError", - } - auth_types = { - "litellm.exceptions.AuthenticationError", - } - - if qualname in rate_limit_types: - raise RateLimitException( - message=f"Rate limited by provider: {exc}", - status_code=429, - ) from exc - - if qualname in auth_types: - raise PyritException(message=f"Authentication failed: {exc}") from exc - - if qualname in transient_types: - raise RateLimitException( - message=f"Transient provider error (will retry): {exc}", - status_code=getattr(exc, "status_code", 503), - ) from exc + 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) - raise PyritException(message=f"LiteLLM error: {exc}") from exc + 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. - def _validate_response(self, response: Any) -> None: - if not hasattr(response, "choices") or not response.choices: - raise EmptyResponseException(message="No choices returned from LiteLLM.") + 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``. - choice = response.choices[0] - finish_reason = getattr(choice, "finish_reason", None) - valid_reasons = {"stop", "length", "tool_calls", "content_filter"} - if finish_reason and finish_reason not in valid_reasons: - raise PyritException( - message=f"Unexpected finish_reason '{finish_reason}' from LiteLLM response." - ) + Args: + response (Any): The LiteLLM completion response object. - content = getattr(choice.message, "content", None) if hasattr(choice, "message") else None - tool_calls = getattr(choice.message, "tool_calls", None) if hasattr(choice, "message") else None + 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 - if not content and not tool_calls: - raise EmptyResponseException(message="LiteLLM returned an empty response.") + 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. - def _construct_message_from_response(self, response: Any, request_piece: MessagePiece) -> Message: - choice = response.choices[0] - content = getattr(choice.message, "content", None) or "" + Args: + exc (Exception): The exception raised by ``litellm.acompletion``. + request (MessagePiece): The originating request piece. - pieces: list[MessagePiece] = [] + Returns: + list[Message]: A single error Message when the failure is a content-policy block. - if content: - text_msg = construct_response_from_request( - request=request_piece, - response_text_pieces=[content], - response_type="text", - ) - pieces.append(text_msg.message_pieces[0]) - - tool_calls = getattr(choice.message, "tool_calls", None) - if tool_calls: - 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, - }, - } - tool_msg = construct_response_from_request( - request=request_piece, - response_text_pieces=[json.dumps(tool_call_data)], - response_type="function_call", + 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, ) - pieces.append(tool_msg.message_pieces[0]) - - if not pieces: - raise EmptyResponseException(message="Failed to extract any response content from LiteLLM.") - - 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 + ] + + 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), ) - pieces[0].prompt_metadata["token_usage_total_tokens"] = getattr(response.usage, "total_tokens", 0) + 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 - return Message(message_pieces=pieces) + 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..3992dfe851 --- /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 :class:`~pyrit.models.TokenUsage`) + +The chat/vision model defaults to ``gpt-5`` 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") + 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..11f2ee1cc1 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"), } @@ -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 :meth:`~pyrit.models.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..6e4eb3f4b0 --- /dev/null +++ b/tests/unit/models/test_token_usage.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from types import SimpleNamespace + +from pyrit.models import TokenUsage + + +def _usage(**kwargs): + """Build a lightweight stand-in for a provider usage object from keyword attributes.""" + return SimpleNamespace(**kwargs) + + +def test_from_provider_usage_maps_prompt_completion(): + usage = _usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + result = TokenUsage.from_provider_usage(usage) + 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_from_provider_usage_accepts_responses_api_names(): + usage = _usage(input_tokens=7, output_tokens=3, total_tokens=10) + result = TokenUsage.from_provider_usage(usage) + assert result.input_tokens == 7 + assert result.output_tokens == 3 + + +def test_from_provider_usage_derives_total_when_missing(): + usage = _usage(prompt_tokens=4, completion_tokens=6) + result = TokenUsage.from_provider_usage(usage) + assert result.total_tokens == 10 + + +def test_from_provider_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, cache_write_tokens=5), + completion_tokens_details=_usage( + reasoning_tokens=12, audio_tokens=3, accepted_prediction_tokens=2, rejected_prediction_tokens=1 + ), + ) + result = TokenUsage.from_provider_usage(usage) + assert result.cached_tokens == 40 + assert result.reasoning_tokens == 12 + assert result.extra == { + "input_audio_tokens": 8, + "cache_write_tokens": 5, + "output_audio_tokens": 3, + "accepted_prediction_tokens": 2, + "rejected_prediction_tokens": 1, + } + + +def test_from_provider_usage_reads_responses_api_detail_names(): + usage = _usage( + input_tokens=100, + output_tokens=50, + input_tokens_details=_usage(cached_tokens=9), + output_tokens_details=_usage(reasoning_tokens=11), + ) + result = TokenUsage.from_provider_usage(usage) + assert result.cached_tokens == 9 + assert result.reasoning_tokens == 11 + + +def test_from_provider_usage_ignores_non_int_and_bool(): + usage = _usage(prompt_tokens=True, completion_tokens="5", total_tokens=None) + result = TokenUsage.from_provider_usage(usage) + assert result.input_tokens is None + assert result.output_tokens is None + assert result.total_tokens is None + + +def test_from_provider_usage_handles_missing_details(): + usage = _usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) + result = TokenUsage.from_provider_usage(usage) + assert result.cached_tokens is None + assert result.reasoning_tokens is None + assert result.extra == {} + + +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..a4ffcfafc6 --- /dev/null +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -0,0 +1,347 @@ +# 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 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, + 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 + + +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 index 24f2c0eae2..4ae4efa7bb 100644 --- a/tests/unit/prompt_target/target/test_litellm_chat_target.py +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -1,45 +1,91 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import base64 import json import sys import types -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from pyrit.exceptions.exception_classes import ( +from pyrit.exceptions import ( EmptyResponseException, PyritException, RateLimitException, + get_retry_max_num_attempts, ) from pyrit.models import Message, MessagePiece - - -def _make_litellm_stub(): - """Install a fake litellm module so the target can be imported without the real package.""" +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") - - exc_mod = types.ModuleType("litellm.exceptions") - exc_mod.RateLimitError = type("RateLimitError", (Exception,), {"__module__": "litellm.exceptions"}) - exc_mod.APIConnectionError = type("APIConnectionError", (Exception,), {"__module__": "litellm.exceptions"}) - exc_mod.Timeout = type("Timeout", (Exception,), {"__module__": "litellm.exceptions"}) - exc_mod.AuthenticationError = type("AuthenticationError", (Exception,), {"__module__": "litellm.exceptions"}) - exc_mod.InternalServerError = type("InternalServerError", (Exception,), {"__module__": "litellm.exceptions"}) - exc_mod.ServiceUnavailableError = type( - "ServiceUnavailableError", (Exception,), {"__module__": "litellm.exceptions"} + 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 - sys.modules["litellm"] = mod - sys.modules["litellm.exceptions"] = exc_mod - return 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") -LITELLM_STUB = _make_litellm_stub() -from pyrit.prompt_target.litellm_chat_target import LiteLLMChatTarget # noqa: E402 +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- def _mock_response(content="hello", finish_reason="stop", model="anthropic/claude-sonnet-4-6"): @@ -48,11 +94,23 @@ def _mock_response(content="hello", finish_reason="stop", model="anthropic/claud 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 = MagicMock() 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 @@ -67,220 +125,553 @@ def _mock_tool_call_response(): return resp -def _make_message(text="test prompt", role="user"): +def _user_message(text="test prompt", conversation_id="convo"): piece = MessagePiece( - role=role, - conversation_id="test_convo", + role="user", + conversation_id=conversation_id, original_value=text, - converted_value=text, original_value_data_type="text", - converted_value_data_type="text", ) - return Message(message_pieces=[piece]) + return piece.to_message() -@pytest.fixture -def target(patch_central_database) -> LiteLLMChatTarget: - return LiteLLMChatTarget(model_name="anthropic/claude-sonnet-4-6") +def _disabled_json_config() -> _JsonResponseConfig: + return _JsonResponseConfig.from_metadata(metadata={}) -# ── constructor tests ───────────────────────────────────────────────── +# --------------------------------------------------------------------------- +# Constructor +# --------------------------------------------------------------------------- -def test_init_requires_model_name(patch_central_database): + +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_env_var(patch_central_database, monkeypatch): - monkeypatch.setenv("LITELLM_MODEL", "openai/gpt-4o") - t = 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_overrides_env(patch_central_database, monkeypatch): - monkeypatch.setenv("LITELLM_MODEL", "openai/gpt-4o") - t = LiteLLMChatTarget(model_name="anthropic/claude-haiku-4-5") +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_drop_params_defaults_true(target): - assert target._drop_params is True +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_drop_params_can_be_disabled(patch_central_database): - t = LiteLLMChatTarget(model_name="openai/gpt-4o", drop_params=False) - assert t._drop_params is False +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 -# ── request body tests ──────────────────────────────────────────────── +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_construct_request_body_includes_drop_params(target): + +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) - assert body["drop_params"] is True + 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(patch_central_database): - t = LiteLLMChatTarget(model_name="openai/gpt-4o", api_key="sk-test") - body = t._construct_request_body([{"role": "user", "content": "hi"}]) - assert body["api_key"] == "sk-test" +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_api_key_when_blank(target): - body = target._construct_request_body([{"role": "user", "content": "hi"}]) +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): +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, - api_base="http://localhost:4000", + frequency_penalty=0.1, + presence_penalty=0.2, + seed=7, + n=2, + stop=["END"], + endpoint="http://localhost:4000", ) - body = t._construct_request_body([{"role": "user", "content": "hi"}]) + 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" -# ── send prompt tests ──────────────────────────────────────────────── +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 + -@pytest.mark.asyncio -async def test_send_prompt_returns_text_response(target): - LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_response("The answer is 4.")) - conversation = [_make_message("What is 2+2?")] +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) - result = await target._send_prompt_to_target_async(normalized_conversation=conversation) + +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 + 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["drop_params"] is True + assert call_kwargs["num_retries"] == target._num_retries -@pytest.mark.asyncio -async def test_send_prompt_handles_tool_calls(target): - LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_tool_call_response()) - conversation = [_make_message("What's the weather?")] +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_to_target_async(normalized_conversation=conversation) + result = await target.send_prompt_async(message=_user_message("What's the weather?")) - assert len(result) == 1 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" -@pytest.mark.asyncio -async def test_send_prompt_captures_token_usage(target): - LITELLM_STUB.acompletion = AsyncMock(return_value=_mock_response("ok")) - conversation = [_make_message("hi")] +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_to_target_async(normalized_conversation=conversation) + result = await target.send_prompt_async(message=_user_message("hi")) metadata = result[0].message_pieces[0].prompt_metadata - assert metadata["token_usage_prompt_tokens"] == 10 - assert metadata["token_usage_completion_tokens"] == 5 + assert metadata["token_usage_input_tokens"] == 10 + assert metadata["token_usage_output_tokens"] == 5 assert metadata["token_usage_total_tokens"] == 15 -# ── empty/malformed response tests ─────────────────────────────────── +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) -@pytest.mark.asyncio -async def test_empty_response_raises(target): - empty_resp = _mock_response(content=None) - empty_resp.choices[0].message.content = None - empty_resp.choices[0].message.tool_calls = None - LITELLM_STUB.acompletion = AsyncMock(return_value=empty_resp) + 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) - with pytest.raises(EmptyResponseException): - await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) +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) -@pytest.mark.asyncio -async def test_no_choices_raises(target): - bad_resp = MagicMock() - bad_resp.choices = [] - LITELLM_STUB.acompletion = AsyncMock(return_value=bad_resp) + 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_to_target_async(normalized_conversation=[_make_message()]) + await target.send_prompt_async(message=_user_message()) -# ── exception translation tests ────────────────────────────────────── +async def test_no_choices_raises_pyrit_exception(target, litellm_stub): + bad = MagicMock() + bad.choices = [] + litellm_stub.acompletion = AsyncMock(return_value=bad) -@pytest.mark.asyncio -async def test_rate_limit_translated(target): - exc = LITELLM_STUB.exceptions.RateLimitError("rate limited") - LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + with pytest.raises(PyritException, match="No choices"): + await target.send_prompt_async(message=_user_message()) - with pytest.raises(RateLimitException, match="Rate limited"): - await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) +async def test_unknown_finish_reason_raises(target, litellm_stub): + litellm_stub.acompletion = AsyncMock(return_value=_mock_response(finish_reason="banana")) -@pytest.mark.asyncio -async def test_auth_error_translated(target): - exc = LITELLM_STUB.exceptions.AuthenticationError("bad key") - LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) + with pytest.raises(PyritException, match="Unknown finish_reason"): + await target.send_prompt_async(message=_user_message()) - with pytest.raises(PyritException, match="Authentication failed"): - await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) +# --------------------------------------------------------------------------- +# Exception translation +# --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_connection_error_translated_to_retryable(target): - exc = LITELLM_STUB.exceptions.APIConnectionError("connection reset") - LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) - with pytest.raises(RateLimitException, match="Transient"): - await target._send_prompt_to_target_async(normalized_conversation=[_make_message()]) +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()) -@pytest.mark.asyncio -async def test_timeout_translated_to_retryable(target): - exc = LITELLM_STUB.exceptions.Timeout("timed out") - LITELLM_STUB.acompletion = AsyncMock(side_effect=exc) +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_to_target_async(normalized_conversation=[_make_message()]) + await target.send_prompt_async(message=_user_message()) -@pytest.mark.asyncio -async def test_unknown_error_wrapped_in_pyrit_exception(target): - LITELLM_STUB.acompletion = AsyncMock(side_effect=RuntimeError("something broke")) +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_to_target_async(normalized_conversation=[_make_message()]) + await target.send_prompt_async(message=_user_message()) -# ── validation edge cases ──────────────────────────────────────────── +# --------------------------------------------------------------------------- +# API key resolution +# --------------------------------------------------------------------------- -def test_unexpected_finish_reason_raises(target): - resp = _mock_response(finish_reason="unexpected_value") - with pytest.raises(PyritException, match="Unexpected finish_reason"): - target._validate_response(resp) +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" -def test_valid_finish_reasons_accepted(target): - for reason in ("stop", "length", "tool_calls", "content_filter"): - resp = _mock_response(finish_reason=reason) - target._validate_response(resp) +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 -# ── message building tests ─────────────────────────────────────────── -def test_build_chat_messages_preserves_roles(target): - msgs = [_make_message("hello", "user"), _make_message("hi there", "assistant")] - result = target._build_chat_messages(msgs) - assert result[0]["role"] == "user" - assert result[0]["content"] == "hello" - assert result[1]["role"] == "assistant" - assert result[1]["content"] == "hi there" +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 From f3c68ba5557851918cf53a443ffcffe5e71fe16f Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 10 Jul 2026 14:19:31 -0700 Subject: [PATCH 3/4] refactor: separate provider token-usage parsing from TokenUsage value object Make TokenUsage a pure value object (fields + to_metadata/from_metadata) and move Chat Completions usage parsing into chat_completions_response_parser as token_usage_from_chat_completion, explicit to the one wire shape both chat targets actually send. - Drop the speculative Responses-API sniffing (no caller sends that shape); a Responses target should parse its own usage in its own module. - Tolerate dict-or-attribute usage payloads so a mapping no longer silently yields all-None counts. - Capture LiteLLM/Anthropic top-level cache fields (cache_read_input_tokens -> cached_tokens, cache_creation_input_tokens -> extra), preserving a zero cached count. - Move/expand parsing tests into test_chat_completions_helpers; test_token_usage now covers only metadata round-tripping. - Convert remaining Sphinx roles to plain double-backtick references. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/models/token_usage.py | 107 ++-------------- .../chat_completions_message_builder.py | 2 +- .../chat_completions_response_parser.py | 115 +++++++++++++++++- pyrit/prompt_target/litellm_chat_target.py | 4 +- tests/unit/models/test_token_usage.py | 81 ------------ .../target/test_chat_completions_helpers.py | 115 ++++++++++++++++++ 6 files changed, 238 insertions(+), 186 deletions(-) diff --git a/pyrit/models/token_usage.py b/pyrit/models/token_usage.py index a903f803be..32bc39ce1e 100644 --- a/pyrit/models/token_usage.py +++ b/pyrit/models/token_usage.py @@ -9,59 +9,36 @@ #: Prefix for all token-usage keys stored in a MessagePiece's ``prompt_metadata``. _METADATA_PREFIX = "token_usage_" -#: Attribute names a provider ``usage`` object may use for the input/output counts. Chat -#: Completions (OpenAI/LiteLLM) reports ``prompt_tokens``/``completion_tokens``; the Responses API -#: (and Anthropic/Gemini) reports ``input_tokens``/``output_tokens``. ``from_provider_usage`` -#: accepts either, but PyRIT normalizes everything to the ``input``/``output`` vocabulary. -_PROVIDER_INPUT_NAMES = ("input_tokens", "prompt_tokens") -_PROVIDER_OUTPUT_NAMES = ("output_tokens", "completion_tokens") - #: 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"}) -def _first_int(source: Any, *names: str) -> int | None: - """ - Return the first attribute among ``names`` on ``source`` that is an int, else None. - - Booleans are rejected even though ``bool`` is a subclass of ``int``. - - Args: - source (Any): The object to read attributes from. - names (str): Candidate attribute names, tried in order. - - Returns: - int | None: The first integer value found, or None. - """ - for name in names: - value = getattr(source, name, None) - if isinstance(value, int) and not isinstance(value, bool): - return value - return None - - @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; - :meth:`from_provider_usage` accepts either provider shape. The object is persisted onto a - :class:`~pyrit.models.MessagePiece`'s ``prompt_metadata`` via :meth:`to_metadata` using matching - ``token_usage_input_tokens`` / ``token_usage_output_tokens`` key names (one consistent + 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 :meth:`to_metadata`) rather than being coerced to a misleading zero. + omitted from ``to_metadata``) rather than being coerced to a misleading zero. """ input_tokens: int | None = None @@ -71,62 +48,13 @@ class TokenUsage: cached_tokens: int | None = None extra: dict[str, int] = field(default_factory=dict) - @classmethod - def from_provider_usage(cls, usage: Any) -> TokenUsage: - """ - Build a :class:`TokenUsage` from a provider ``usage`` object (OpenAI/LiteLLM/Anthropic shape). - - Reads top-level counts (mapping ``prompt``->``input`` and ``completion``->``output``, and - accepting Responses-API ``input``/``output`` names directly) and the nested - ``*_tokens_details`` breakdowns. ``total_tokens`` is derived from input + output when the - provider omits it. Unmodeled detail counts (audio, predicted-output, cache-write) are - collected into ``extra`` with disambiguated names. - - Args: - usage (Any): The provider usage object. - - Returns: - TokenUsage: The parsed token usage. - """ - input_tokens = _first_int(usage, *_PROVIDER_INPUT_NAMES) - output_tokens = _first_int(usage, *_PROVIDER_OUTPUT_NAMES) - total_tokens = _first_int(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 - - input_details = getattr(usage, "prompt_tokens_details", None) or getattr(usage, "input_tokens_details", None) - output_details = getattr(usage, "completion_tokens_details", None) or getattr( - usage, "output_tokens_details", None - ) - - cached_tokens = _first_int(input_details, "cached_tokens") if input_details is not None else None - reasoning_tokens = _first_int(output_details, "reasoning_tokens") if output_details is not None else None - - extra: dict[str, int] = {} - if input_details is not None: - _put(extra, "input_audio_tokens", _first_int(input_details, "audio_tokens")) - _put(extra, "cache_write_tokens", _first_int(input_details, "cache_write_tokens", "cache_creation_tokens")) - if output_details is not None: - _put(extra, "output_audio_tokens", _first_int(output_details, "audio_tokens")) - _put(extra, "accepted_prediction_tokens", _first_int(output_details, "accepted_prediction_tokens")) - _put(extra, "rejected_prediction_tokens", _first_int(output_details, "rejected_prediction_tokens")) - - return cls( - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=total_tokens, - reasoning_tokens=reasoning_tokens, - cached_tokens=cached_tokens, - extra=extra, - ) - @classmethod def from_metadata(cls, metadata: dict[str, Any]) -> TokenUsage | None: """ - Reconstruct a :class:`TokenUsage` from a :class:`~pyrit.models.MessagePiece`'s ``prompt_metadata``. + Reconstruct a ``TokenUsage`` from a ``MessagePiece``'s ``prompt_metadata``. Reads the ``token_usage_input_tokens`` / ``token_usage_output_tokens`` keys written by - :meth:`to_metadata`. Non-core integer ``token_usage_*`` keys are collected into ``extra``; + ``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: @@ -184,16 +112,3 @@ def to_metadata(self) -> dict[str, int]: for name, value in self.extra.items(): out[_METADATA_PREFIX + name] = value return out - - -def _put(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 diff --git a/pyrit/prompt_target/common/chat_completions_message_builder.py b/pyrit/prompt_target/common/chat_completions_message_builder.py index c42f85913a..149cc1023f 100644 --- a/pyrit/prompt_target/common/chat_completions_message_builder.py +++ b/pyrit/prompt_target/common/chat_completions_message_builder.py @@ -192,7 +192,7 @@ async def build_multimodal_chat_messages_async( 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 :func:`should_skip_audio_piece`). + its transcript (see ``should_skip_audio_piece``). Returns: list[dict[str, Any]]: The list of constructed chat messages. diff --git a/pyrit/prompt_target/common/chat_completions_response_parser.py b/pyrit/prompt_target/common/chat_completions_response_parser.py index e218b33dd6..c539b83ba5 100644 --- a/pyrit/prompt_target/common/chat_completions_response_parser.py +++ b/pyrit/prompt_target/common/chat_completions_response_parser.py @@ -13,6 +13,7 @@ import base64 import json import logging +from collections.abc import Mapping from typing import Any from pyrit.exceptions import ( @@ -276,11 +277,10 @@ 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 provider usage through :class:`~pyrit.models.TokenUsage`, which reads both the - top-level counts and the nested ``prompt_tokens_details`` / ``completion_tokens_details`` - breakdowns (so cached and reasoning tokens are captured, not just prompt/completion/total). - 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. + 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. @@ -290,10 +290,113 @@ def capture_token_usage(*, pieces: list[MessagePiece], response: Any) -> None: if not usage or not pieces: return - token_usage = TokenUsage.from_provider_usage(usage) + 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. diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index 0c97f3b539..5294caba05 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -109,8 +109,8 @@ class LiteLLMChatTarget(PromptTarget): 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 - :mod:`pyrit.prompt_target.common.chat_completions_message_builder` and - :mod:`pyrit.prompt_target.common.chat_completions_response_parser`. + ``pyrit.prompt_target.common.chat_completions_message_builder`` and + ``pyrit.prompt_target.common.chat_completions_response_parser``. Install the optional dependency:: diff --git a/tests/unit/models/test_token_usage.py b/tests/unit/models/test_token_usage.py index 6e4eb3f4b0..439c92efd2 100644 --- a/tests/unit/models/test_token_usage.py +++ b/tests/unit/models/test_token_usage.py @@ -1,90 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from types import SimpleNamespace - from pyrit.models import TokenUsage -def _usage(**kwargs): - """Build a lightweight stand-in for a provider usage object from keyword attributes.""" - return SimpleNamespace(**kwargs) - - -def test_from_provider_usage_maps_prompt_completion(): - usage = _usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) - result = TokenUsage.from_provider_usage(usage) - 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_from_provider_usage_accepts_responses_api_names(): - usage = _usage(input_tokens=7, output_tokens=3, total_tokens=10) - result = TokenUsage.from_provider_usage(usage) - assert result.input_tokens == 7 - assert result.output_tokens == 3 - - -def test_from_provider_usage_derives_total_when_missing(): - usage = _usage(prompt_tokens=4, completion_tokens=6) - result = TokenUsage.from_provider_usage(usage) - assert result.total_tokens == 10 - - -def test_from_provider_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, cache_write_tokens=5), - completion_tokens_details=_usage( - reasoning_tokens=12, audio_tokens=3, accepted_prediction_tokens=2, rejected_prediction_tokens=1 - ), - ) - result = TokenUsage.from_provider_usage(usage) - assert result.cached_tokens == 40 - assert result.reasoning_tokens == 12 - assert result.extra == { - "input_audio_tokens": 8, - "cache_write_tokens": 5, - "output_audio_tokens": 3, - "accepted_prediction_tokens": 2, - "rejected_prediction_tokens": 1, - } - - -def test_from_provider_usage_reads_responses_api_detail_names(): - usage = _usage( - input_tokens=100, - output_tokens=50, - input_tokens_details=_usage(cached_tokens=9), - output_tokens_details=_usage(reasoning_tokens=11), - ) - result = TokenUsage.from_provider_usage(usage) - assert result.cached_tokens == 9 - assert result.reasoning_tokens == 11 - - -def test_from_provider_usage_ignores_non_int_and_bool(): - usage = _usage(prompt_tokens=True, completion_tokens="5", total_tokens=None) - result = TokenUsage.from_provider_usage(usage) - assert result.input_tokens is None - assert result.output_tokens is None - assert result.total_tokens is None - - -def test_from_provider_usage_handles_missing_details(): - usage = _usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) - result = TokenUsage.from_provider_usage(usage) - assert result.cached_tokens is None - assert result.reasoning_tokens is None - assert result.extra == {} - - 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() diff --git a/tests/unit/prompt_target/target/test_chat_completions_helpers.py b/tests/unit/prompt_target/target/test_chat_completions_helpers.py index a4ffcfafc6..d4de5c1ae5 100644 --- a/tests/unit/prompt_target/target/test_chat_completions_helpers.py +++ b/tests/unit/prompt_target/target/test_chat_completions_helpers.py @@ -5,6 +5,7 @@ import base64 import json +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -29,6 +30,7 @@ extract_partial_content, is_content_filter_response, save_audio_response_async, + token_usage_from_chat_completion, validate_chat_completion_response, ) @@ -180,6 +182,119 @@ def test_capture_token_usage_noop_without_usage(): 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 From e27b6141b181d48cf5ae8906219946954f2c5466 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 10 Jul 2026 14:32:37 -0700 Subject: [PATCH 4/4] test: drop .env_example test-model overrides; default integration tests to gpt-5.4 Revert the PLATFORM_OPENAI_CHAT_MODEL / PLATFORM_OPENAI_AUDIO_MODEL additions to .env_example and default the chat/vision integration fixtures to the deployed gpt-5.4 model instead of the generic "gpt-5". Also convert two stray Sphinx roles in the integration test docstrings to plain double-backtick references. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .env_example | 4 ---- .../targets/test_litellm_chat_target_integration.py | 6 +++--- .../targets/test_openai_chat_target_integration.py | 4 ++-- 3 files changed, 5 insertions(+), 9 deletions(-) diff --git a/.env_example b/.env_example index 25cb650214..b925fb097c 100644 --- a/.env_example +++ b/.env_example @@ -31,10 +31,6 @@ PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_GPT4O_MODEL="gpt-4o" -# Model overrides for the chat integration tests (OpenAIChatTarget / LiteLLMChatTarget). -# Optional: if unset, the integration tests default to "gpt-5" (chat/vision) and "gpt-audio" (audio). -PLATFORM_OPENAI_CHAT_MODEL="gpt-5" -PLATFORM_OPENAI_AUDIO_MODEL="gpt-audio" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately # Example: https://xxxx.openai.azure.com/openai/v1 diff --git a/tests/integration/targets/test_litellm_chat_target_integration.py b/tests/integration/targets/test_litellm_chat_target_integration.py index 3992dfe851..1a63cb0ab4 100644 --- a/tests/integration/targets/test_litellm_chat_target_integration.py +++ b/tests/integration/targets/test_litellm_chat_target_integration.py @@ -12,9 +12,9 @@ - Tool calling via the ``extra_body_parameters`` passthrough - Multimodal image input (vision) - Multimodal audio input/output -- Token-usage metadata capture (parsed back through :class:`~pyrit.models.TokenUsage`) +- Token-usage metadata capture (parsed back through ``TokenUsage``) -The chat/vision model defaults to ``gpt-5`` and the audio model to ``gpt-audio``; both can be +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. """ @@ -52,7 +52,7 @@ def platform_litellm_args(): 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") + model = os.environ.get("PLATFORM_OPENAI_CHAT_MODEL", "gpt-5.4") return { "model_name": f"openai/{model}", "endpoint": endpoint, diff --git a/tests/integration/targets/test_openai_chat_target_integration.py b/tests/integration/targets/test_openai_chat_target_integration.py index 11f2ee1cc1..563ef62df5 100644 --- a/tests/integration/targets/test_openai_chat_target_integration.py +++ b/tests/integration/targets/test_openai_chat_target_integration.py @@ -63,7 +63,7 @@ def platform_openai_chat_args(): return { "endpoint": endpoint, "api_key": api_key, - "model_name": os.environ.get("PLATFORM_OPENAI_CHAT_MODEL", "gpt-5"), + "model_name": os.environ.get("PLATFORM_OPENAI_CHAT_MODEL", "gpt-5.4"), } @@ -216,7 +216,7 @@ 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 is recoverable via :meth:`~pyrit.models.TokenUsage.from_metadata` + 1. Token usage is recoverable via ``TokenUsage.from_metadata`` 2. Token counts are positive integers 3. The total equals input + output """