diff --git a/pyrit/auth/__init__.py b/pyrit/auth/__init__.py index ac37e89abf..28746eecd7 100644 --- a/pyrit/auth/__init__.py +++ b/pyrit/auth/__init__.py @@ -21,6 +21,7 @@ from pyrit.auth.azure_storage_auth import AzureStorageAuth from pyrit.auth.copilot_authenticator import CopilotAuthenticator from pyrit.auth.manual_copilot_authenticator import ManualCopilotAuthenticator +from pyrit.auth.openai_auth import resolve_openai_auth __all__ = [ "AsyncTokenProviderCredential", @@ -29,6 +30,7 @@ "AzureStorageAuth", "CopilotAuthenticator", "ManualCopilotAuthenticator", + "resolve_openai_auth", "TokenProviderCredential", "ensure_async_token_provider", "get_azure_token_provider", diff --git a/pyrit/auth/openai_auth.py b/pyrit/auth/openai_auth.py new file mode 100644 index 0000000000..aa281cc199 --- /dev/null +++ b/pyrit/auth/openai_auth.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from collections.abc import Awaitable, Callable +from typing import cast + +from pyrit.auth.azure_auth import ensure_async_token_provider, get_azure_openai_auth, is_azure_openai_endpoint +from pyrit.common import default_values + + +def resolve_openai_auth( + *, + endpoint: str, + api_key: str | Callable[[], str | Awaitable[str]] | None, + api_key_environment_variable: str, +) -> str | Callable[[], Awaitable[str]]: + """ + Resolve OpenAI authentication from a key, environment variable, or Azure Entra fallback. + + Args: + endpoint (str): The OpenAI-compatible endpoint URL. + api_key (str | Callable[[], str | Awaitable[str]] | None): The explicit API key or token provider. + api_key_environment_variable (str): Environment variable to use when ``api_key`` is not provided. + + Returns: + str | Callable[[], Awaitable[str]]: API key string or async-compatible token provider. + + Raises: + ValueError: If no key is provided and the endpoint is not a recognized Azure OpenAI endpoint. + """ + if api_key is not None and callable(api_key): + return cast("str | Callable[[], Awaitable[str]]", ensure_async_token_provider(api_key)) + + api_key_value = default_values.get_non_required_value( + env_var_name=api_key_environment_variable, passed_value=api_key + ) + if api_key_value: + return api_key_value + + if is_azure_openai_endpoint(endpoint): + return get_azure_openai_auth(endpoint) + + raise ValueError( + f"Environment variable {api_key_environment_variable} is required for non-Azure endpoints. " + "For recognized Azure OpenAI / AI Foundry endpoints, Entra ID authentication is used automatically." + ) diff --git a/pyrit/embedding/openai_text_embedding.py b/pyrit/embedding/openai_text_embedding.py index 86a774c404..5a84a880d6 100644 --- a/pyrit/embedding/openai_text_embedding.py +++ b/pyrit/embedding/openai_text_embedding.py @@ -8,7 +8,7 @@ import tenacity from openai import AsyncOpenAI -from pyrit.auth import ensure_async_token_provider +from pyrit.auth import resolve_openai_auth from pyrit.common import default_values from pyrit.models import ( EmbeddingData, @@ -40,7 +40,10 @@ def __init__( Args: api_key: The API key (string) or token provider (callable) for authentication. - For Azure with Entra auth, pass get_azure_openai_auth(endpoint) from pyrit.auth. + For recognized Azure OpenAI / AI Foundry endpoints, if no API key is provided + (via parameter or environment variable), Entra ID authentication is used automatically. + You can also explicitly pass a token provider from pyrit.auth + (e.g., get_azure_openai_auth(endpoint) for async). Defaults to OPENAI_EMBEDDING_KEY environment variable. endpoint: The API endpoint URL. For Azure: https://{resource}.openai.azure.com/openai/v1 @@ -48,6 +51,10 @@ def __init__( Defaults to OPENAI_EMBEDDING_ENDPOINT environment variable. model_name: The model/deployment name (e.g., "text-embedding-3-small"). Defaults to OPENAI_EMBEDDING_MODEL environment variable. + + Raises: + ValueError: If no API key is provided (via parameter or environment variable) and the + endpoint is not a recognized Azure OpenAI / AI Foundry endpoint. """ endpoint = default_values.get_required_value( env_var_name=self.ENDPOINT_URI_ENVIRONMENT_VARIABLE, passed_value=endpoint @@ -56,15 +63,13 @@ def __init__( env_var_name=self.MODEL_ENVIRONMENT_VARIABLE, passed_value=model_name ) - if api_key is None: - api_key = default_values.get_required_value( - env_var_name=self.API_KEY_ENVIRONMENT_VARIABLE, passed_value=api_key - ) - - # Wrap sync token providers for async compatibility; AsyncOpenAI accepts str or async callable - resolved_api_key = ensure_async_token_provider(api_key) + async_api_key = resolve_openai_auth( + endpoint=endpoint, + api_key=api_key, + api_key_environment_variable=self.API_KEY_ENVIRONMENT_VARIABLE, + ) self._async_client = AsyncOpenAI( - api_key=resolved_api_key, + api_key=async_api_key, base_url=endpoint, ) diff --git a/pyrit/prompt_target/openai/openai_target.py b/pyrit/prompt_target/openai/openai_target.py index 9b7356847e..b943c44a54 100644 --- a/pyrit/prompt_target/openai/openai_target.py +++ b/pyrit/prompt_target/openai/openai_target.py @@ -22,7 +22,7 @@ AuthenticationError, ) -from pyrit.auth import ensure_async_token_provider, get_azure_openai_auth, is_azure_openai_endpoint +from pyrit.auth import resolve_openai_auth from pyrit.common import default_values from pyrit.exceptions.exception_classes import ( RateLimitException, @@ -152,27 +152,11 @@ def __init__( custom_configuration=custom_configuration, ) - # API key: use passed value, env var, or fall back to Entra ID for Azure endpoints - resolved_api_key: str | Callable[[], str | Awaitable[str]] - if api_key is not None and callable(api_key): - resolved_api_key = api_key - else: - api_key_value = default_values.get_non_required_value( - env_var_name=self.api_key_environment_variable, passed_value=api_key - ) - if api_key_value: - resolved_api_key = api_key_value - elif is_azure_openai_endpoint(endpoint_value): - resolved_api_key = get_azure_openai_auth(endpoint_value) - else: - raise ValueError( - f"Environment variable {self.api_key_environment_variable} is required for non-Azure endpoints. " - "For recognized Azure OpenAI / AI Foundry endpoints, Entra ID authentication is used " - "automatically." - ) - - # Ensure api_key is async-compatible (wrap sync token providers if needed) - self._api_key = ensure_async_token_provider(resolved_api_key) + self._api_key = resolve_openai_auth( + endpoint=endpoint_value, + api_key=api_key, + api_key_environment_variable=self.api_key_environment_variable, + ) self._initialize_openai_client() diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 9984b33e9e..d6fd6df784 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -455,7 +455,7 @@ async def test_create_openai_target_with_entra_omits_key_and_target_mints_token( with patch.dict(os.environ, {}, clear=False): os.environ.pop("OPENAI_CHAT_KEY", None) with patch( - "pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", + "pyrit.auth.openai_auth.get_azure_openai_auth", return_value=_test_token_provider, ) as mock_get_auth: service = TargetService() @@ -483,7 +483,7 @@ async def test_create_openai_target_with_identity_drops_user_api_key(self, sqlit with patch.dict(os.environ, {}, clear=False): os.environ.pop("OPENAI_CHAT_KEY", None) with patch( - "pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", + "pyrit.auth.openai_auth.get_azure_openai_auth", return_value=_test_token_provider, ): service = TargetService() @@ -512,7 +512,7 @@ async def test_create_openai_target_with_identity_does_not_mutate_request_params with patch.dict(os.environ, {}, clear=False): os.environ.pop("OPENAI_CHAT_KEY", None) with patch( - "pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", + "pyrit.auth.openai_auth.get_azure_openai_auth", return_value=_test_token_provider, ): service = TargetService() diff --git a/tests/unit/embedding/test_azure_text_embedding.py b/tests/unit/embedding/test_azure_text_embedding.py index e572fa774d..d2f1ba4b63 100644 --- a/tests/unit/embedding/test_azure_text_embedding.py +++ b/tests/unit/embedding/test_azure_text_embedding.py @@ -2,10 +2,10 @@ # Licensed under the MIT license. import os -from unittest.mock import MagicMock, patch +from collections.abc import Awaitable, Callable +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from openai import OpenAIError from pyrit.embedding import OpenAITextEmbedding @@ -27,12 +27,12 @@ def test_valid_init_env(): def test_invalid_key_raises(): - """Test that an empty API key is rejected by the underlying OpenAI client.""" + """An empty API key on a non-Azure endpoint raises ValueError (no Entra fallback).""" os.environ[OpenAITextEmbedding.API_KEY_ENVIRONMENT_VARIABLE] = "" - with pytest.raises(OpenAIError, match="Missing credentials"): + with pytest.raises(ValueError, match="required for non-Azure endpoints"): OpenAITextEmbedding( api_key="", - endpoint="https://mock.azure.com/", + endpoint="https://api.openai.com/v1", model_name="gpt-4", ) @@ -100,3 +100,60 @@ def mock_token_provider(): assert async_call_args.kwargs["base_url"] == "https://mock.azure.com/" assert embedding._async_client == mock_async_client + + +_AZURE_ENDPOINT = "https://foo.openai.azure.com/openai/v1" +_NON_AZURE_ENDPOINT = "https://api.openai.com/v1" + + +def _build_embedding( + *, + endpoint: str = _AZURE_ENDPOINT, + api_key: str | Callable[[], str | Awaitable[str]] | None = "test-key", + model_name: str = "text-embedding-3-small", +) -> OpenAITextEmbedding: + """Build an OpenAITextEmbedding with a cleared environment so env vars don't leak in.""" + with patch.dict(os.environ, {}, clear=True): + return OpenAITextEmbedding(api_key=api_key, endpoint=endpoint, model_name=model_name) + + +@patch("pyrit.embedding.openai_text_embedding.AsyncOpenAI") +def test_explicit_string_api_key_used_directly(mock_async_openai): + """An explicit string api_key is passed to the async client as-is.""" + mock_async_openai.return_value = MagicMock() + + _build_embedding(api_key="my-secret-key") + + assert mock_async_openai.call_args.kwargs["api_key"] == "my-secret-key" + + +@patch("pyrit.embedding.openai_text_embedding.AsyncOpenAI") +def test_callable_token_provider_used_as_is(mock_async_openai): + """An async token provider is used directly and not overwritten by env resolution.""" + mock_async_openai.return_value = MagicMock() + + async def async_provider() -> str: + return "async-token" + + _build_embedding(api_key=async_provider) + + assert mock_async_openai.call_args.kwargs["api_key"] is async_provider + + +@patch("pyrit.embedding.openai_text_embedding.AsyncOpenAI") +def test_no_key_azure_endpoint_falls_back_to_entra(mock_async_openai): + """A recognized Azure endpoint with no key mints an Entra token provider.""" + mock_async_openai.return_value = MagicMock() + mock_auth = AsyncMock(return_value="entra-token") + + with patch("pyrit.auth.openai_auth.get_azure_openai_auth", return_value=mock_auth) as mock_get_auth: + _build_embedding(api_key=None, endpoint=_AZURE_ENDPOINT) + + mock_get_auth.assert_called_once_with(_AZURE_ENDPOINT) + assert mock_async_openai.call_args.kwargs["api_key"] is mock_auth + + +def test_no_key_non_azure_endpoint_raises(): + """A non-Azure endpoint with no key raises ValueError (no Entra fallback).""" + with pytest.raises(ValueError, match="required for non-Azure endpoints"): + _build_embedding(api_key=None, endpoint=_NON_AZURE_ENDPOINT) 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 1a37a9f27b..fa71f7e776 100644 --- a/tests/unit/prompt_target/target/test_openai_chat_target.py +++ b/tests/unit/prompt_target/target/test_openai_chat_target.py @@ -741,7 +741,7 @@ async def _provider() -> str: with ( patch.dict(os.environ, {}, clear=True), patch( - "pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", + "pyrit.auth.openai_auth.get_azure_openai_auth", return_value=_provider, ) as mock_get_auth, ): diff --git a/tests/unit/prompt_target/target/test_openai_target_auth.py b/tests/unit/prompt_target/target/test_openai_target_auth.py index a95e59cd5f..b2b98e0d73 100644 --- a/tests/unit/prompt_target/target/test_openai_target_auth.py +++ b/tests/unit/prompt_target/target/test_openai_target_auth.py @@ -80,7 +80,7 @@ def test_non_azure_endpoint_without_key_raises(self): def test_azure_endpoint_falls_back_to_entra(self): """Azure endpoints without a key fall back to get_azure_openai_auth.""" mock_auth = AsyncMock(return_value="entra-token") - with patch("pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", return_value=mock_auth): + with patch("pyrit.auth.openai_auth.get_azure_openai_auth", return_value=mock_auth): target = _build_target( endpoint="https://myresource.openai.azure.com/openai/v1", api_key=None,