From ba3fa38e970eb8ada7e600d1bb5c1e867d4bd06c Mon Sep 17 00:00:00 2001 From: Kunal Karmakar Date: Sun, 31 May 2026 15:08:01 +0000 Subject: [PATCH] Add support for Azure AAD based Auth --- docs/configuration.md | 66 +++++++++++ nanobot/providers/azure_openai_provider.py | 72 +++++++++++- nanobot/providers/factory.py | 6 +- pyproject.toml | 3 + tests/providers/test_azure_openai_provider.py | 108 ++++++++++++++++-- 5 files changed, 241 insertions(+), 14 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index bc1ad8c0b..fa6c02e1f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -205,6 +205,72 @@ Valid `apiType` values are exactly `auto`, `chat_completions`, and `responses`. +
+Azure OpenAI + +The `azure_openai` provider talks to your Azure OpenAI resource via the OpenAI **Responses API** (`/openai/v1/responses`). Model names map to **deployment names**, not OpenAI model IDs. Two authentication modes are supported. + +**Mode 1: Static API key** (simplest) + +```json +{ + "providers": { + "azure_openai": { + "apiKey": "${AZURE_OPENAI_API_KEY}", + "apiBase": "https://my-resource.openai.azure.com" + } + }, + "agents": { + "defaults": { + "provider": "azure_openai", + "model": "my-gpt-5-deployment" + } + } +} +``` + +**Mode 2: Microsoft Entra ID (Azure AD) via `DefaultAzureCredential`** + +Omit `apiKey` (or leave it empty / unset). The provider falls back to [`DefaultAzureCredential`](https://learn.microsoft.com/azure/developer/python/sdk/authentication/credential-chains#defaultazurecredential-overview) and acquires a bearer token scoped to `https://cognitiveservices.azure.com/.default` for every request. The Azure SDK's own MSAL-backed cache returns valid tokens without a network round-trip. + +```json +{ + "providers": { + "azure_openai": { + "apiBase": "https://my-resource.openai.azure.com" + } + }, + "agents": { + "defaults": { + "provider": "azure_openai", + "model": "my-gpt-5-deployment" + } + } +} +``` + +Install the optional dependency: + +```bash +pip install 'nanobot-ai[azure]' +``` + +`DefaultAzureCredential` walks this chain in order and uses the first identity that succeeds: + +1. **EnvironmentCredential** — reads `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and one of `AZURE_CLIENT_SECRET` / `AZURE_CLIENT_CERTIFICATE_PATH` / `AZURE_USERNAME` + `AZURE_PASSWORD`. +2. **WorkloadIdentityCredential** — for AKS workload-identity / federated tokens (`AZURE_FEDERATED_TOKEN_FILE`). +3. **ManagedIdentityCredential** — for Azure VMs, App Service, Functions, Container Apps, etc. +4. **AzureCliCredential** — uses the token from `az login` on your dev machine. +5. **AzurePowerShellCredential** — uses the token from `Connect-AzAccount`. +6. **AzureDeveloperCliCredential** — uses the token from `azd auth login`. +7. **InteractiveBrowserCredential** *(disabled by default)*. + +The identity that ends up signing the request **must be assigned the `Cognitive Services OpenAI User` RBAC role** (or higher) on the Azure OpenAI resource. Without that role you will see `401`/`403` errors at the first request. + +> `apiBase` remains mandatory in both modes — it's your Azure resource endpoint and cannot be inferred. If neither `apiKey` is set nor `azure-identity` is installed, the provider raises a clear error pointing you at `pip install 'nanobot-ai[azure]'`. + +
+
Skywork / APIFree diff --git a/nanobot/providers/azure_openai_provider.py b/nanobot/providers/azure_openai_provider.py index 24a65cdfe..50256a1c5 100644 --- a/nanobot/providers/azure_openai_provider.py +++ b/nanobot/providers/azure_openai_provider.py @@ -3,6 +3,18 @@ Uses ``AsyncOpenAI`` pointed at ``https://{endpoint}/openai/v1/`` which routes to the Responses API (``/responses``). Reuses shared conversion helpers from :mod:`nanobot.providers.openai_responses`. + +Authentication +-------------- +Two modes are supported, selected automatically: + +1. **Static API key** — when ``api_key`` is non-empty it is sent as the + ``api-key`` / ``Authorization: Bearer`` header (existing behavior). +2. **Microsoft Entra ID (AAD)** — when ``api_key`` is empty the provider + falls back to :class:`azure.identity.aio.DefaultAzureCredential` and + acquires a bearer token scoped to + ``https://cognitiveservices.azure.com/.default``. ``azure-identity`` + is an optional dependency installed via ``pip install nanobot-ai[azure]``. """ from __future__ import annotations @@ -21,6 +33,48 @@ from nanobot.providers.openai_responses import ( parse_response_output, ) +_AZURE_OPENAI_SCOPE = "https://cognitiveservices.azure.com/.default" + + +class _AzureTokenProvider: + """Async bearer-token callback for AAD authentication. + + Thin wrapper around :class:`azure.identity.aio.DefaultAzureCredential` + that exposes itself as an async callable returning a fresh bearer + token. The Azure SDK's own MSAL-backed token cache already returns + valid tokens without network calls, so no extra caching is layered on + top here. + + Raises ``RuntimeError`` with a clear install hint if + ``azure-identity`` is not installed. + """ + + def __init__(self, scope: str = _AZURE_OPENAI_SCOPE) -> None: + try: + from azure.identity.aio import DefaultAzureCredential + except ImportError as exc: + raise RuntimeError( + "Azure OpenAI AAD authentication requires the 'azure-identity' package. " + "Install it with: pip install 'nanobot-ai[azure]'" + ) from exc + + self._scope = scope + self._credential = DefaultAzureCredential() + + async def __call__(self) -> str: + """Return a bearer token for the configured scope.""" + access_token = await self._credential.get_token(self._scope) + return access_token.token + + async def aclose(self) -> None: + """Release credential resources. Safe to call multiple times.""" + close = getattr(self._credential, "close", None) + if close is not None: + try: + await close() + except Exception: + pass + class AzureOpenAIProvider(LLMProvider): """Azure OpenAI provider backed by the Responses API. @@ -31,6 +85,8 @@ class AzureOpenAIProvider(LLMProvider): - Calls ``client.responses.create()`` (Responses API) - Reuses shared message/tool/SSE conversion from ``openai_responses`` + - Falls back to :class:`DefaultAzureCredential` (AAD) when ``api_key`` + is empty. See module docstring for details. """ def __init__( @@ -42,8 +98,6 @@ class AzureOpenAIProvider(LLMProvider): super().__init__(api_key, api_base) self.default_model = default_model - if not api_key: - raise ValueError("Azure OpenAI api_key is required") if not api_base: raise ValueError("Azure OpenAI api_base is required") @@ -52,10 +106,22 @@ class AzureOpenAIProvider(LLMProvider): api_base += "/" self.api_base = api_base + # Select auth mode. A truthy api_key wins; otherwise fall back to + # AAD via DefaultAzureCredential. The OpenAI SDK accepts an async + # callable as ``api_key`` and invokes it per request, using the + # returned string as the bearer token. + self._token_provider: _AzureTokenProvider | None = None + client_api_key: str | Callable[[], Awaitable[str]] + if api_key: + client_api_key = api_key + else: + self._token_provider = _AzureTokenProvider() + client_api_key = self._token_provider + # SDK client targeting the Azure Responses API endpoint base_url = f"{api_base.rstrip('/')}/openai/v1/" self._client = AsyncOpenAI( - api_key=api_key, + api_key=client_api_key, base_url=base_url, default_headers={"x-session-affinity": uuid.uuid4().hex}, max_retries=0, diff --git a/nanobot/providers/factory.py b/nanobot/providers/factory.py index a10c0d5cd..e8275f93a 100644 --- a/nanobot/providers/factory.py +++ b/nanobot/providers/factory.py @@ -44,8 +44,8 @@ def _make_provider_core( backend = spec.backend if spec else "openai_compat" if backend == "azure_openai": - if not p or not p.api_key or not p.api_base: - raise ValueError("Azure OpenAI requires api_key and api_base in config.") + if not p or not p.api_base: + raise ValueError("Azure OpenAI requires api_base in config.") elif backend == "openai_compat" and not model.startswith("bedrock/"): needs_key = not (p and p.api_key) exempt = spec and (spec.is_oauth or spec.is_local or spec.is_direct) @@ -60,7 +60,7 @@ def _make_provider_core( from nanobot.providers.azure_openai_provider import AzureOpenAIProvider provider = AzureOpenAIProvider( - api_key=p.api_key, + api_key=p.api_key or "", api_base=p.api_base, default_model=model, ) diff --git a/pyproject.toml b/pyproject.toml index cbfa9f445..7adbb9c51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,9 @@ dependencies = [ api = [ "aiohttp>=3.9.0,<4.0.0", ] +azure = [ + "azure-identity>=1.19.0,<2.0.0", +] wecom = [ "wecom-aibot-sdk-python>=0.1.5", ] diff --git a/tests/providers/test_azure_openai_provider.py b/tests/providers/test_azure_openai_provider.py index 7ce74ee9f..73d078d2e 100644 --- a/tests/providers/test_azure_openai_provider.py +++ b/tests/providers/test_azure_openai_provider.py @@ -1,13 +1,18 @@ """Test Azure OpenAI provider (Responses API via OpenAI SDK).""" -from unittest.mock import AsyncMock, MagicMock +import sys +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from nanobot.providers.azure_openai_provider import AzureOpenAIProvider +from nanobot.providers.azure_openai_provider import ( + AzureOpenAIProvider, + _AzureTokenProvider, +) from nanobot.providers.base import LLMResponse - # --------------------------------------------------------------------------- # Init & validation # --------------------------------------------------------------------------- @@ -25,6 +30,8 @@ def test_init_creates_sdk_client(): assert provider.default_model == "gpt-4o-deployment" # SDK client base_url ends with /openai/v1/ assert str(provider._client.base_url).rstrip("/").endswith("/openai/v1") + # Static-key path must NOT construct an AAD token provider + assert provider._token_provider is None def test_init_base_url_no_trailing_slash(): @@ -42,11 +49,6 @@ def test_init_base_url_with_trailing_slash(): assert str(provider._client.base_url).rstrip("/").endswith("/openai/v1") -def test_init_validation_missing_key(): - with pytest.raises(ValueError, match="Azure OpenAI api_key is required"): - AzureOpenAIProvider(api_key="", api_base="https://test.com") - - def test_init_validation_missing_base(): with pytest.raises(ValueError, match="Azure OpenAI api_base is required"): AzureOpenAIProvider(api_key="test", api_base="") @@ -59,6 +61,96 @@ def test_no_api_version_in_base_url(): assert "api-version" not in base +# --------------------------------------------------------------------------- +# AAD / DefaultAzureCredential fallback +# --------------------------------------------------------------------------- + + +def _install_fake_azure_identity(monkeypatch, credential_factory): + """Install a fake ``azure.identity.aio`` module exposing ``DefaultAzureCredential``.""" + azure_mod = sys.modules.get("azure") or SimpleNamespace() + identity_mod = SimpleNamespace() + aio_mod = SimpleNamespace(DefaultAzureCredential=credential_factory) + identity_mod.aio = aio_mod + azure_mod.identity = identity_mod # type: ignore[attr-defined] + + monkeypatch.setitem(sys.modules, "azure", azure_mod) + monkeypatch.setitem(sys.modules, "azure.identity", identity_mod) + monkeypatch.setitem(sys.modules, "azure.identity.aio", aio_mod) + + +def test_init_missing_key_uses_aad_token_provider(monkeypatch): + """Empty api_key falls back to DefaultAzureCredential via _AzureTokenProvider.""" + credential_instance = MagicMock() + credential_factory = MagicMock(return_value=credential_instance) + _install_fake_azure_identity(monkeypatch, credential_factory) + + provider = AzureOpenAIProvider( + api_key="", api_base="https://res.openai.azure.com", + ) + + assert provider._token_provider is not None + assert isinstance(provider._token_provider, _AzureTokenProvider) + # DefaultAzureCredential must have been instantiated exactly once + credential_factory.assert_called_once_with() + # The SDK client must have received the token provider as its api_key + # (the SDK stores it on the auth wrapper, not directly accessible — so + # we assert the callable was wired in via the provider attribute). + assert provider._token_provider._credential is credential_instance + + +def test_init_explicit_key_does_not_construct_credential(monkeypatch): + """Explicit api_key wins; DefaultAzureCredential must never be touched.""" + credential_factory = MagicMock(side_effect=AssertionError( + "DefaultAzureCredential must not be constructed when api_key is set" + )) + _install_fake_azure_identity(monkeypatch, credential_factory) + + provider = AzureOpenAIProvider( + api_key="real-key", api_base="https://res.openai.azure.com", + ) + + assert provider._token_provider is None + credential_factory.assert_not_called() + + +def test_init_missing_key_without_azure_identity_raises(monkeypatch): + """Clear RuntimeError with pip-install hint when azure-identity is missing.""" + # Force the import inside _AzureTokenProvider to fail. + real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) else __builtins__.__import__ + + def fake_import(name, *args, **kwargs): + if name == "azure.identity.aio": + raise ImportError("No module named 'azure.identity.aio'") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=fake_import): + with pytest.raises(RuntimeError, match=r"pip install 'nanobot-ai\[azure\]'"): + AzureOpenAIProvider(api_key="", api_base="https://res.openai.azure.com") + + +@pytest.mark.asyncio +async def test_token_provider_returns_credential_token(monkeypatch): + """Token callback delegates to DefaultAzureCredential.get_token with the AOAI scope.""" + access_token = SimpleNamespace(token="token-A", expires_on=time.time() + 3600) + credential_instance = MagicMock() + credential_instance.get_token = AsyncMock(return_value=access_token) + + credential_factory = MagicMock(return_value=credential_instance) + _install_fake_azure_identity(monkeypatch, credential_factory) + + tp = _AzureTokenProvider() + + assert await tp() == "token-A" + credential_instance.get_token.assert_awaited_with( + "https://cognitiveservices.azure.com/.default" + ) + # No client-side caching layer — every call delegates to the Azure SDK, + # which has its own MSAL-backed token cache. + assert await tp() == "token-A" + assert credential_instance.get_token.await_count == 2 + + # --------------------------------------------------------------------------- # _supports_temperature # ---------------------------------------------------------------------------