diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index 6a5d8bea73..6e0da9ea09 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -392,6 +392,7 @@ def _build_identifier(self) -> ComponentIdentifier: "seed": self._seed, "n": self._n, "stop": self._stop, + "extra_body_parameters": self._extra_body_parameters, }, ) diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index 23ea5dbe32..fc1ea4070c 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/pyrit/prompt_target/openai/openai_chat_target.py @@ -199,6 +199,7 @@ def _build_identifier(self) -> ComponentIdentifier: "presence_penalty": self._presence_penalty, "seed": self._seed, "n": self._n, + "extra_body_parameters": self._extra_body_parameters, }, ) 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 6c5bbf338b..c935fbd1da 100644 --- a/tests/unit/prompt_target/target/test_litellm_chat_target.py +++ b/tests/unit/prompt_target/target/test_litellm_chat_target.py @@ -326,6 +326,14 @@ def test_identifier_includes_behavioral_params_and_excludes_key(patch_central_da assert "sk-secret" not in json.dumps(params) +@pytest.mark.usefixtures("patch_central_database", "litellm_stub") +def test_identifier_distinguishes_extra_body_parameters() -> None: + low = LiteLLMChatTarget(model_name="openai/gpt-4o", extra_body_parameters={"reasoning_effort": "low"}) + high = LiteLLMChatTarget(model_name="openai/gpt-4o", extra_body_parameters={"reasoning_effort": "high"}) + assert low.get_identifier().params["extra_body_parameters"] == {"reasoning_effort": "low"} + assert low.get_identifier().hash != high.get_identifier().hash + + # --------------------------------------------------------------------------- # Request body # --------------------------------------------------------------------------- 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 589b645268..381bd05ce0 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -8,6 +8,7 @@ from collections.abc import MutableSequence from tempfile import NamedTemporaryFile from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -1279,6 +1280,22 @@ async def test_construct_message_from_response(target: OpenAIChatTarget, dummy_t # Tests for underlying_model parameter and get_identifier +@pytest.mark.usefixtures("patch_central_database") +def test_get_identifier_distinguishes_extra_body_parameters() -> None: + def build(**kwargs: Any) -> OpenAIChatTarget: + return OpenAIChatTarget(model_name="o4-mini", endpoint="https://mock.azure.com/", api_key="k", **kwargs) + + low = build(extra_body_parameters={"reasoning_effort": "low"}) + high = build(extra_body_parameters={"reasoning_effort": "high"}) + alloy = build(audio_response_config=OpenAIChatAudioConfig(voice="alloy")) + echo = build(audio_response_config=OpenAIChatAudioConfig(voice="echo")) + + assert low.get_identifier().params["extra_body_parameters"] == {"reasoning_effort": "low"} + assert low.get_identifier().hash != high.get_identifier().hash + assert alloy.get_identifier().hash != echo.get_identifier().hash + assert build().get_identifier().params.get("extra_body_parameters") is None + + def test_get_identifier_uses_model_name_when_no_underlying_model(patch_central_database): """Test that get_identifier uses model_name when underlying_model is not provided.""" # Clear the environment variable to ensure it doesn't interfere with the test