FIX: Entra ID auth fallback for Azure endpoints in OpenAITextEmbedding

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",

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."
+    )

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,14 +40,21 @@ 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
                 For platform OpenAI: https://api.openai.com/v1
                 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,
         )
 

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()
 

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()

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)

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,
     ):

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,