From 9ed8cc5a8fcd831252c42663e1091cbca199389e Mon Sep 17 00:00:00 2001 From: Mark Rossetti Date: Thu, 20 Aug 2026 18:21:23 +0000 Subject: [PATCH 1/3] feat(models) adding python support for Azure foundry models (OpenAI + anthropic types) Signed-off-by: Mark Rossetti --- python/packages/kagent-adk/pyproject.toml | 1 + .../src/kagent/adk/models/__init__.py | 2 + .../src/kagent/adk/models/_anthropic.py | 41 ++- .../src/kagent/adk/models/_azure.py | 179 +++++++++++ .../src/kagent/adk/models/_embedding.py | 81 ++++- .../src/kagent/adk/models/_openai.py | 68 +++- .../kagent-adk/src/kagent/adk/types.py | 36 ++- .../tests/unittests/models/test_anthropic.py | 209 ++++++++++++- .../tests/unittests/models/test_azure.py | 294 ++++++++++++++++++ .../tests/unittests/models/test_openai.py | 3 +- .../tests/unittests/test_embedding.py | 128 +++++++- python/uv.lock | 57 ++++ 12 files changed, 1059 insertions(+), 40 deletions(-) create mode 100644 python/packages/kagent-adk/src/kagent/adk/models/_azure.py create mode 100644 python/packages/kagent-adk/tests/unittests/models/test_azure.py diff --git a/python/packages/kagent-adk/pyproject.toml b/python/packages/kagent-adk/pyproject.toml index 3f906d838..265170227 100644 --- a/python/packages/kagent-adk/pyproject.toml +++ b/python/packages/kagent-adk/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "boto3>=1.28.57", "ollama >=0.3.6", # Ollama SDK "numpy>=2.2.6", + "azure-identity>=1.19.0", # Azure Workload Identity for Azure OpenAI / Foundry ] [tool.uv.sources] diff --git a/python/packages/kagent-adk/src/kagent/adk/models/__init__.py b/python/packages/kagent-adk/src/kagent/adk/models/__init__.py index bad9f56d2..b94ea8ddf 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/__init__.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/__init__.py @@ -4,11 +4,13 @@ from ._gemini import KAgentGeminiLlm, KAgentGeminiVertexAILlm from ._ollama import KAgentOllamaLlm from ._openai import AzureOpenAI, OpenAI +from ._openai import FoundryOpenAI as Foundry from ._sap_ai_core import KAgentSAPAICoreLlm __all__ = [ "OpenAI", "AzureOpenAI", + "Foundry", "KAgentAnthropicLlm", "KAgentBedrockLlm", "KAgentGeminiLlm", diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_anthropic.py b/python/packages/kagent-adk/src/kagent/adk/models/_anthropic.py index b8e9e68cd..098e02440 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_anthropic.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_anthropic.py @@ -10,6 +10,11 @@ from anthropic import AsyncAnthropic from google.adk.models.anthropic_llm import AnthropicLlm +from ._azure import ( + build_foundry_anthropic_client, + resolve_azure_api_key, + resolve_foundry_endpoint_deployment, +) from ._ssl import KAgentTLSMixin logger = logging.getLogger(__name__) @@ -28,10 +33,10 @@ class KAgentAnthropicLlm(KAgentTLSMixin, AnthropicLlm): def set_passthrough_key(self, token: str) -> None: """Forward the Bearer token from the incoming A2A request as the Anthropic API key.""" - self._api_key = token - # Invalidate cached clients so they're recreated with the new key - self.__dict__.pop("_anthropic_client", None) - self.__dict__.pop("_http_client", None) + if self._api_key != token: + self._api_key = token + # The SDK client captures auth at construction, so rebuild it only when the token changes. + self.__dict__.pop("_anthropic_client", None) def _create_http_client(self): """Create HTTP client with custom SSL context using Anthropic SDK defaults. @@ -58,3 +63,31 @@ def _anthropic_client(self) -> AsyncAnthropic: kwargs["http_client"] = http_client return AsyncAnthropic(**kwargs) + + +class FoundryAnthropic(KAgentAnthropicLlm): + """Claude on Azure AI Foundry's Anthropic Messages API.""" + + endpoint: Optional[str] = None + deployment: Optional[str] = None + + def _resolve_model_name(self, model: Optional[str]) -> str: + del model + _, deployment = resolve_foundry_endpoint_deployment(self.endpoint, self.deployment) + return deployment + + @cached_property + def _anthropic_client(self) -> AsyncAnthropic: + endpoint, _ = resolve_foundry_endpoint_deployment(self.endpoint, self.deployment) + api_key = resolve_azure_api_key( + self._api_key, + api_key_passthrough=self.api_key_passthrough, + environment_variable="FOUNDRY_API_KEY", + ) + return build_foundry_anthropic_client( + endpoint=endpoint, + api_key=api_key, + api_key_passthrough=self.api_key_passthrough, + default_headers=self.extra_headers, + http_client=self._create_http_client(), + ) diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_azure.py b/python/packages/kagent-adk/src/kagent/adk/models/_azure.py new file mode 100644 index 000000000..ae11a920e --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/models/_azure.py @@ -0,0 +1,179 @@ +"""Shared configuration and client helpers for Azure AI providers.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional + +import httpx +from anthropic import AsyncAnthropic +from openai import AsyncAzureOpenAI +from openai.lib.azure import API_KEY_SENTINEL + +if TYPE_CHECKING: + from anthropic.lib.credentials import AccessToken, AccessTokenProvider + from azure.identity import DefaultAzureCredential + +COGNITIVE_SERVICES_SCOPE = "https://cognitiveservices.azure.com/.default" +AI_FOUNDRY_SCOPE = "https://ai.azure.com/.default" + +AZURE_OPENAI_DEFAULT_API_VERSION = "2024-02-15-preview" +FOUNDRY_DEFAULT_API_VERSION = "2024-10-21" + +_AUTH_HEADER_NAMES = {"authorization", "api-key", "x-api-key"} + +AsyncTokenProvider = Callable[[], Awaitable[str]] + + +def azure_ad_token_provider(scope: str = COGNITIVE_SERVICES_SCOPE) -> AsyncTokenProvider: + """Return an async bearer-token provider backed by ``DefaultAzureCredential``.""" + from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider + + return get_bearer_token_provider(DefaultAzureCredential(), scope) + + +class _AzureAccessTokenProvider: + """Adapt Azure Identity to Anthropic's synchronous access-token provider.""" + + def __init__(self, scope: str) -> None: + from azure.identity import DefaultAzureCredential + + self._credential: DefaultAzureCredential = DefaultAzureCredential() + self._scope = scope + + def __call__(self, *, force_refresh: bool = False) -> "AccessToken": + del force_refresh + + from anthropic.lib.credentials import AccessToken + + token = self._credential.get_token(self._scope) + return AccessToken(token=token.token, expires_at=token.expires_on) + + def close(self) -> None: + self._credential.close() + + +def azure_access_token_provider(scope: str = AI_FOUNDRY_SCOPE) -> "AccessTokenProvider": + """Return an Azure provider compatible with Anthropic's token cache.""" + return _AzureAccessTokenProvider(scope) + + +def resolve_azure_api_key( + api_key: Optional[str], + *, + api_key_passthrough: Optional[bool], + environment_variable: str, +) -> Optional[str]: + """Resolve an Azure API key without bypassing passthrough mode.""" + if api_key_passthrough: + return api_key + return api_key or os.environ.get(environment_variable) + + +def resolve_azure_openai_config(endpoint: Optional[str], api_version: Optional[str]) -> tuple[str, str]: + """Resolve Azure OpenAI endpoint and API version configuration.""" + resolved_endpoint = endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT") + if not resolved_endpoint: + raise ValueError( + "Azure endpoint must be provided either via azure_endpoint parameter or " + "AZURE_OPENAI_ENDPOINT environment variable" + ) + + resolved_api_version = api_version or os.environ.get("OPENAI_API_VERSION") or AZURE_OPENAI_DEFAULT_API_VERSION + return resolved_endpoint, resolved_api_version + + +def resolve_foundry_endpoint_deployment(endpoint: Optional[str], deployment: Optional[str]) -> tuple[str, str]: + """Resolve the Foundry endpoint and deployment.""" + resolved_endpoint = endpoint or os.environ.get("FOUNDRY_ENDPOINT") + if not resolved_endpoint: + raise ValueError( + "Foundry endpoint must be provided either via endpoint parameter or FOUNDRY_ENDPOINT environment variable" + ) + + resolved_deployment = deployment or os.environ.get("FOUNDRY_DEPLOYMENT") + if not resolved_deployment: + raise ValueError( + "Foundry deployment must be provided either via deployment parameter or " + "FOUNDRY_DEPLOYMENT environment variable" + ) + + return resolved_endpoint, resolved_deployment + + +def resolve_foundry_config( + endpoint: Optional[str], deployment: Optional[str], api_version: Optional[str] +) -> tuple[str, str, str]: + """Resolve Foundry OpenAI-compatible data-plane configuration.""" + resolved_endpoint, resolved_deployment = resolve_foundry_endpoint_deployment(endpoint, deployment) + + resolved_api_version = api_version or os.environ.get("FOUNDRY_API_VERSION") or FOUNDRY_DEFAULT_API_VERSION + return resolved_endpoint, resolved_deployment, resolved_api_version + + +def sanitize_azure_auth_headers(headers: Optional[dict[str, str]]) -> Optional[dict[str, str]]: + """Remove configured headers that could conflict with resolved Azure auth.""" + if not headers: + return None + sanitized = {name: value for name, value in headers.items() if name.lower() not in _AUTH_HEADER_NAMES} + return sanitized or None + + +def build_azure_openai_client( + *, + api_version: str, + azure_endpoint: str, + azure_deployment: Optional[str], + api_key: Optional[str], + api_key_passthrough: Optional[bool], + default_headers: Optional[dict[str, str]], + http_client: Optional[httpx.AsyncClient], + missing_credential_hint: str, +) -> AsyncAzureOpenAI: + """Build an Azure OpenAI client using key, passthrough, or Workload Identity auth.""" + token_provider = None + if not api_key: + if api_key_passthrough: + raise ValueError(missing_credential_hint) + token_provider = azure_ad_token_provider() + + return AsyncAzureOpenAI( + # The sentinel prevents environment-key fallback while the token + # provider authenticates each request. + api_key=API_KEY_SENTINEL if token_provider is not None else api_key, + azure_ad_token_provider=token_provider, + api_version=api_version, + azure_endpoint=azure_endpoint, + azure_deployment=azure_deployment, + default_headers=sanitize_azure_auth_headers(default_headers), + http_client=http_client, + ) + + +def build_foundry_anthropic_client( + *, + endpoint: str, + api_key: Optional[str], + api_key_passthrough: Optional[bool], + default_headers: Optional[dict[str, str]], + http_client: Optional[httpx.AsyncClient], +) -> AsyncAnthropic: + """Build a Foundry Anthropic client using key, passthrough, or Workload Identity auth.""" + if api_key_passthrough and not api_key: + raise ValueError( + "No Azure credential resolved: provide the passthrough token before creating the Foundry Anthropic client" + ) + + kwargs: dict[str, Any] = {"base_url": endpoint.rstrip("/") + "/anthropic"} + if api_key: + kwargs["api_key"] = api_key + else: + kwargs["credentials"] = azure_access_token_provider(AI_FOUNDRY_SCOPE) + + safe_headers = sanitize_azure_auth_headers(default_headers) + if safe_headers: + kwargs["default_headers"] = safe_headers + if http_client is not None: + kwargs["http_client"] = http_client + + return AsyncAnthropic(**kwargs) diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py index deccc2978..611ea9dc8 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py @@ -19,6 +19,13 @@ from kagent.adk._bearer_token import bearer_token from kagent.adk.types import EmbeddingConfig +from ._azure import ( + build_azure_openai_client, + resolve_azure_api_key, + resolve_azure_openai_config, + resolve_foundry_config, +) + logger = logging.getLogger(__name__) @@ -85,6 +92,8 @@ async def _call_provider(self, texts: List[str]) -> List[List[float]]: if provider in ("openai", "azure_openai"): return await self._embed_openai(texts) + if provider == "foundry": + return await self._embed_foundry(texts) if provider == "ollama": return await self._embed_ollama(texts) if provider in ("vertex_ai", "gemini"): @@ -146,8 +155,8 @@ def _normalize_l2(self, x: Union[List[float], np.ndarray]) -> np.ndarray: def _passthrough_api_key(self) -> Optional[str]: """Bearer token to use as the API key when api_key_passthrough is enabled, mirroring BaseOpenAI.set_passthrough_key for chat models. - None falls back to the SDK's own env var lookup (OPENAI_API_KEY / - AZURE_OPENAI_API_KEY). + Azure providers treat a missing token as an error rather than falling + back to a provider environment key. """ if not self.config.api_key_passthrough: return None @@ -159,13 +168,20 @@ async def _embed_openai(self, texts: List[str]) -> List[List[float]]: api_key = self._passthrough_api_key() if provider == "azure_openai": - from openai import AsyncAzureOpenAI - - api_version = os.environ.get("OPENAI_API_VERSION", "2024-02-15-preview") - api_base = self.config.base_url or os.environ.get("AZURE_OPENAI_ENDPOINT") - if not api_base: - raise ValueError("Azure OpenAI endpoint must be set via base_url or AZURE_OPENAI_ENDPOINT env var") - client = AsyncAzureOpenAI(api_version=api_version, azure_endpoint=api_base, api_key=api_key) + api_base, api_version = resolve_azure_openai_config( + self.config.endpoint or self.config.base_url, self.config.api_version + ) + api_key = resolve_azure_api_key( + api_key, + api_key_passthrough=self.config.api_key_passthrough, + environment_variable="AZURE_OPENAI_API_KEY", + ) + client = self._build_azure_client( + api_version=api_version, + endpoint=api_base, + deployment=self.config.deployment, + api_key=api_key, + ) else: from openai import AsyncOpenAI @@ -178,6 +194,53 @@ async def _embed_openai(self, texts: List[str]) -> List[List[float]]: ) return [item.embedding for item in response.data] + async def _embed_foundry(self, texts: List[str]) -> List[List[float]]: + """Embed using the Azure AI Foundry OpenAI-compatible surface.""" + endpoint, deployment, api_version = resolve_foundry_config( + self.config.endpoint, self.config.deployment, self.config.api_version + ) + api_key = resolve_azure_api_key( + self._passthrough_api_key(), + api_key_passthrough=self.config.api_key_passthrough, + environment_variable="FOUNDRY_API_KEY", + ) + + client = self._build_azure_client( + api_version=api_version, + endpoint=endpoint, + deployment=deployment, + api_key=api_key, + ) + response = await client.embeddings.create( + model=self.config.model, + input=texts, + dimensions=self.TARGET_DIMENSION, + ) + return [item.embedding for item in response.data] + + def _build_azure_client( + self, + *, + api_version: str, + endpoint: str, + deployment: Optional[str], + api_key: Optional[str], + ): + """Build an Azure embeddings client with implicit Workload Identity auth.""" + return build_azure_openai_client( + api_version=api_version, + azure_endpoint=endpoint, + azure_deployment=deployment, + api_key=api_key, + api_key_passthrough=self.config.api_key_passthrough, + default_headers=None, + http_client=None, + missing_credential_hint=( + "No Azure credential resolved for embeddings: set an API key, enable " + "api_key_passthrough, or configure Azure Workload Identity" + ), + ) + async def _embed_ollama(self, texts: List[str]) -> List[List[float]]: """Embed using the Ollama SDK.""" import ollama diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_openai.py b/python/packages/kagent-adk/src/kagent/adk/models/_openai.py index b0cca99d6..b6aa938ad 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_openai.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_openai.py @@ -2,7 +2,6 @@ import base64 import json -import os from functools import cached_property from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterable, Literal, Optional, cast @@ -32,6 +31,12 @@ from openai.types.shared_params import FunctionDefinition, FunctionParameters from pydantic import Field +from ._azure import ( + build_azure_openai_client, + resolve_azure_api_key, + resolve_azure_openai_config, + resolve_foundry_config, +) from ._ssl import KAgentTLSMixin from ._token_source import GDCHTokenSource from ._utils import function_declaration_schema @@ -591,26 +596,55 @@ class AzureOpenAI(BaseOpenAI): @cached_property def _client(self) -> AsyncAzureOpenAI: """Get the Azure OpenAI client with optional custom SSL configuration.""" - api_version = self.api_version or os.environ.get("OPENAI_API_VERSION", "2024-02-15-preview") - azure_endpoint = self.azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT") - api_key = self.api_key or os.environ.get("AZURE_OPENAI_API_KEY") + azure_endpoint, api_version = resolve_azure_openai_config(self.azure_endpoint, self.api_version) + api_key = resolve_azure_api_key( + self.api_key, + api_key_passthrough=self.api_key_passthrough, + environment_variable="AZURE_OPENAI_API_KEY", + ) - if not azure_endpoint: - raise ValueError( - "Azure endpoint must be provided either via azure_endpoint parameter or AZURE_OPENAI_ENDPOINT environment variable" - ) + return build_azure_openai_client( + azure_deployment=self.azure_deployment, + api_key=api_key, + api_version=api_version, + azure_endpoint=azure_endpoint, + api_key_passthrough=self.api_key_passthrough, + default_headers=self.default_headers, + http_client=self._create_http_client(), + missing_credential_hint=( + "No Azure credential resolved: set AZURE_OPENAI_API_KEY, enable " + "api_key_passthrough, or configure Azure Workload Identity" + ), + ) - if not api_key: - raise ValueError( - "API key must be provided either via api_key parameter or AZURE_OPENAI_API_KEY environment variable" - ) - http_client = self._create_http_client() +class FoundryOpenAI(BaseOpenAI): + """Azure AI Foundry model using its OpenAI-compatible data plane.""" - return AsyncAzureOpenAI( - api_key=api_key, + type: Literal["foundry"] + endpoint: Optional[str] = None + deployment: Optional[str] = None + api_version: Optional[str] = None + + @cached_property + def _client(self) -> AsyncAzureOpenAI: + endpoint, deployment, api_version = resolve_foundry_config(self.endpoint, self.deployment, self.api_version) + api_key = resolve_azure_api_key( + self.api_key, + api_key_passthrough=self.api_key_passthrough, + environment_variable="FOUNDRY_API_KEY", + ) + + return build_azure_openai_client( api_version=api_version, - azure_endpoint=azure_endpoint, + azure_endpoint=endpoint, + azure_deployment=deployment, + api_key=api_key, + api_key_passthrough=self.api_key_passthrough, default_headers=self.default_headers, - http_client=http_client, + http_client=self._create_http_client(), + missing_credential_hint=( + "No Azure credential resolved: set FOUNDRY_API_KEY, enable " + "api_key_passthrough, or configure Azure Workload Identity" + ), ) diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 9d4a25b28..60763dfa9 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -17,11 +17,12 @@ from kagent.adk._mcp_apps import MCPAppToolNames, make_mcp_app_model_result_callback from kagent.adk._mcp_toolset import KAgentMcpToolset from kagent.adk._remote_a2a_tool import KAgentRemoteA2AToolset -from kagent.adk.models._anthropic import KAgentAnthropicLlm +from kagent.adk.models._anthropic import FoundryAnthropic, KAgentAnthropicLlm from kagent.adk.models._bedrock import KAgentBedrockLlm from kagent.adk.models._gemini import KAgentGeminiLlm, KAgentGeminiVertexAILlm from kagent.adk.models._ollama import create_ollama_llm from kagent.adk.models._openai import AzureOpenAI as OpenAIAzure +from kagent.adk.models._openai import FoundryOpenAI from kagent.adk.models._openai import OpenAI as OpenAINative from kagent.adk.models._ssl import create_ssl_context from kagent.adk.tools.ask_user_tool import AskUserTool @@ -291,6 +292,14 @@ class AzureOpenAI(BaseLLM): type: Literal["azure_openai"] +class Foundry(BaseLLM): + endpoint: str | None = None + deployment: str | None = None + api_version: str | None = None + api_format: Literal["openai", "anthropic"] = "openai" + type: Literal["foundry"] + + class Anthropic(BaseLLM): base_url: str | None = None @@ -348,7 +357,9 @@ class SAPAICore(BaseLLM): type: Literal["sap_ai_core"] -ModelUnion = Union[OpenAI, Anthropic, GeminiVertexAI, GeminiAnthropic, Ollama, AzureOpenAI, Gemini, Bedrock, SAPAICore] +ModelUnion = Union[ + OpenAI, Anthropic, GeminiVertexAI, GeminiAnthropic, Ollama, AzureOpenAI, Foundry, Gemini, Bedrock, SAPAICore +] class ContextCompressionSettings(BaseModel): @@ -371,6 +382,9 @@ class EmbeddingConfig(BaseModel): provider: str base_url: str | None = None api_key_passthrough: bool = False + endpoint: str | None = None + deployment: str | None = None + api_version: str | None = None class MemoryConfig(BaseModel): @@ -705,6 +719,24 @@ def _create_llm_from_model_config(model_config: ModelUnion): default_headers=extra_headers, **_transport_kwargs(model_config), ) + if model_config.type == "foundry": + if model_config.api_format == "anthropic": + return FoundryAnthropic( + model=model_config.deployment or model_config.model, + endpoint=model_config.endpoint, + deployment=model_config.deployment, + extra_headers=extra_headers, + **_transport_kwargs(model_config), + ) + return FoundryOpenAI( + model=model_config.model, + type="foundry", + endpoint=model_config.endpoint, + deployment=model_config.deployment, + api_version=model_config.api_version, + default_headers=extra_headers, + **_transport_kwargs(model_config), + ) if model_config.type == "gemini": return KAgentGeminiLlm( model=model_config.model, diff --git a/python/packages/kagent-adk/tests/unittests/models/test_anthropic.py b/python/packages/kagent-adk/tests/unittests/models/test_anthropic.py index 5f9200b62..24cf0b95e 100644 --- a/python/packages/kagent-adk/tests/unittests/models/test_anthropic.py +++ b/python/packages/kagent-adk/tests/unittests/models/test_anthropic.py @@ -2,11 +2,15 @@ from unittest import mock +import httpx +import pytest from anthropic import AsyncAnthropic +from anthropic.lib.credentials import AccessToken from anthropic.types import ThinkingBlock from google.adk.models.anthropic_llm import content_block_to_part -from kagent.adk.models._anthropic import KAgentAnthropicLlm +from kagent.adk.models._anthropic import FoundryAnthropic, KAgentAnthropicLlm +from kagent.adk.models._azure import AI_FOUNDRY_SCOPE class TestKAgentAnthropicLlm: @@ -24,12 +28,22 @@ def test_set_passthrough_key(self): def test_set_passthrough_key_invalidates_cached_client(self): llm = KAgentAnthropicLlm(model="claude-3-sonnet-20240229") - with mock.patch("anthropic.AsyncAnthropic"): + with mock.patch("kagent.adk.models._anthropic.AsyncAnthropic"): _ = llm._anthropic_client assert "_anthropic_client" in llm.__dict__ llm.set_passthrough_key("new-token") assert "_anthropic_client" not in llm.__dict__ + def test_set_passthrough_key_preserves_cached_client_for_same_token(self): + llm = KAgentAnthropicLlm(model="claude-3-sonnet-20240229", api_key_passthrough=True) + llm.set_passthrough_key("same-token") + with mock.patch("kagent.adk.models._anthropic.AsyncAnthropic"): + cached_client = llm._anthropic_client + + llm.set_passthrough_key("same-token") + + assert llm._anthropic_client is cached_client + def test_client_uses_base_url(self): llm = KAgentAnthropicLlm(model="claude-3-sonnet-20240229", base_url="https://proxy.internal/anthropic") with mock.patch("kagent.adk.models._anthropic.AsyncAnthropic") as mock_anthropic: @@ -67,6 +81,197 @@ def test_create_llm_from_anthropic_model_config(self): assert result.base_url == "https://api.anthropic.com" +class TestFoundryAnthropic: + def test_model_config_dispatches_anthropic_format(self): + from kagent.adk.types import Foundry, _create_llm_from_model_config + + config = Foundry( + type="foundry", + model="claude-haiku-4-5", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + api_format="anthropic", + ) + + result = _create_llm_from_model_config(config) + + assert isinstance(result, FoundryAnthropic) + assert result.model == "claude-haiku-deployment" + assert result._resolve_model_name("wrong-request-model") == "claude-haiku-deployment" + + def test_model_config_defaults_to_openai_format(self): + from kagent.adk.models._openai import FoundryOpenAI + from kagent.adk.types import Foundry, _create_llm_from_model_config + + result = _create_llm_from_model_config( + Foundry( + type="foundry", + model="gpt-4.1", + endpoint="https://example.cognitiveservices.azure.com/", + deployment="gpt-4.1-deployment", + ) + ) + + assert isinstance(result, FoundryOpenAI) + + def test_workload_identity_uses_ai_foundry_scope(self): + token_provider = object() + with ( + mock.patch.dict("os.environ", {}, clear=True), + mock.patch("kagent.adk.models._azure.AsyncAnthropic") as mock_anthropic, + mock.patch( + "kagent.adk.models._azure.azure_access_token_provider", + return_value=token_provider, + ) as mock_provider, + ): + llm = FoundryAnthropic( + model="claude-haiku-deployment", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + extra_headers={"Authorization": "Bearer leaked", "X-Custom": "preserved"}, + ) + _ = llm._anthropic_client + + mock_provider.assert_called_once_with(AI_FOUNDRY_SCOPE) + assert mock_anthropic.call_args.kwargs["credentials"] is token_provider + assert "api_key" not in mock_anthropic.call_args.kwargs + assert mock_anthropic.call_args.kwargs["default_headers"] == {"X-Custom": "preserved"} + assert mock_anthropic.call_args.kwargs["base_url"] == "https://example.services.ai.azure.com/anthropic" + + def test_passthrough_without_token_does_not_fall_back_to_workload_identity(self): + with ( + mock.patch.dict("os.environ", {"FOUNDRY_API_KEY": "must-not-win"}, clear=True), + mock.patch("kagent.adk.models._azure.azure_access_token_provider") as mock_provider, + ): + llm = FoundryAnthropic( + model="claude-haiku-deployment", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + api_key_passthrough=True, + ) + + with pytest.raises(ValueError, match="provide the passthrough token"): + _ = llm._anthropic_client + + mock_provider.assert_not_called() + + def test_passthrough_token_change_rebuilds_foundry_client(self): + llm = FoundryAnthropic( + model="claude-haiku-deployment", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + api_key_passthrough=True, + ) + with mock.patch("kagent.adk.models._azure.AsyncAnthropic") as mock_anthropic: + mock_anthropic.side_effect = [ + mock.MagicMock(spec=AsyncAnthropic), + mock.MagicMock(spec=AsyncAnthropic), + ] + llm.set_passthrough_key("first-token") + first_client = llm._anthropic_client + + llm.set_passthrough_key("second-token") + second_client = llm._anthropic_client + + assert second_client is not first_client + assert mock_anthropic.call_count == 2 + assert mock_anthropic.call_args.kwargs["api_key"] == "second-token" + + @pytest.mark.asyncio + async def test_api_key_uses_messages_path_and_x_api_key(self): + captured_request = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal captured_request + captured_request = request + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-deployment", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + with ( + mock.patch.dict("os.environ", {"FOUNDRY_API_KEY": "foundry-key"}, clear=True), + mock.patch.object(FoundryAnthropic, "_create_http_client", return_value=http_client), + ): + llm = FoundryAnthropic( + model="claude-haiku-deployment", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + extra_headers={"Authorization": "Bearer leaked", "X-Custom": "preserved"}, + ) + await llm._anthropic_client.messages.create( + model=llm._resolve_model_name("wrong-request-model"), + max_tokens=16, + messages=[{"role": "user", "content": "hello"}], + ) + await llm._anthropic_client.close() + + assert captured_request is not None + assert captured_request.url.path == "/anthropic/v1/messages" + assert captured_request.headers["x-api-key"] == "foundry-key" + assert "authorization" not in captured_request.headers + assert captured_request.headers["x-custom"] == "preserved" + + @pytest.mark.asyncio + async def test_workload_identity_uses_bearer_without_x_api_key(self): + captured_request = None + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal captured_request + captured_request = request + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-deployment", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + token_provider = mock.Mock(return_value=AccessToken(token="entra-token", expires_at=4_102_444_800)) + with ( + mock.patch.dict("os.environ", {}, clear=True), + mock.patch.object(FoundryAnthropic, "_create_http_client", return_value=http_client), + mock.patch( + "kagent.adk.models._azure.azure_access_token_provider", + return_value=token_provider, + ), + ): + llm = FoundryAnthropic( + model="claude-haiku-deployment", + endpoint="https://example.services.ai.azure.com/", + deployment="claude-haiku-deployment", + extra_headers={"X-Api-Key": "leaked", "X-Custom": "preserved"}, + ) + await llm._anthropic_client.messages.create( + model=llm._resolve_model_name(None), + max_tokens=16, + messages=[{"role": "user", "content": "hello"}], + ) + await llm._anthropic_client.close() + + assert captured_request is not None + assert captured_request.url.path == "/anthropic/v1/messages" + assert captured_request.headers["authorization"] == "Bearer entra-token" + assert "x-api-key" not in captured_request.headers + assert captured_request.headers["x-custom"] == "preserved" + token_provider.assert_called_once() + + class TestAnthropicThinkingBlock: """Regression guard for the google-adk floor that KAgentAnthropicLlm relies on. diff --git a/python/packages/kagent-adk/tests/unittests/models/test_azure.py b/python/packages/kagent-adk/tests/unittests/models/test_azure.py new file mode 100644 index 000000000..742f94cf8 --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/models/test_azure.py @@ -0,0 +1,294 @@ +from unittest import mock + +import httpx +import pytest +from openai.lib.azure import API_KEY_SENTINEL + +from kagent.adk.models._azure import build_azure_openai_client + +_SENTINEL_TOKEN_PROVIDER = object() + + +def test_azure_openai_workload_identity_when_no_api_key(): + from kagent.adk.models import AzureOpenAI + + with ( + mock.patch.dict("os.environ", {"AZURE_OPENAI_ENDPOINT": "https://test.openai.azure.com"}, clear=True), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch( + "kagent.adk.models._azure.azure_ad_token_provider", return_value=_SENTINEL_TOKEN_PROVIDER + ) as mock_provider, + ): + llm = AzureOpenAI(model="gpt-4o", type="azure_openai") + _ = llm._client + + mock_provider.assert_called_once() + kwargs = mock_azure.call_args.kwargs + assert kwargs["azure_ad_token_provider"] is _SENTINEL_TOKEN_PROVIDER + assert kwargs["api_key"] == API_KEY_SENTINEL + + +def test_azure_openai_api_key_skips_workload_identity(): + from kagent.adk.models import AzureOpenAI + + with ( + mock.patch.dict( + "os.environ", + {"AZURE_OPENAI_ENDPOINT": "https://test.openai.azure.com", "AZURE_OPENAI_API_KEY": "secret"}, + clear=True, + ), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch("kagent.adk.models._azure.azure_ad_token_provider") as mock_provider, + ): + llm = AzureOpenAI(model="gpt-4o", type="azure_openai") + _ = llm._client + + mock_provider.assert_not_called() + kwargs = mock_azure.call_args.kwargs + assert kwargs["api_key"] == "secret" + assert kwargs["azure_ad_token_provider"] is None + + +def test_azure_openai_passthrough_does_not_fall_back_to_workload_identity(): + from kagent.adk.models import AzureOpenAI + + with ( + mock.patch.dict("os.environ", {"AZURE_OPENAI_ENDPOINT": "https://test.openai.azure.com"}, clear=True), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI"), + mock.patch("kagent.adk.models._azure.azure_ad_token_provider") as mock_provider, + ): + llm = AzureOpenAI(model="gpt-4o", type="azure_openai", api_key_passthrough=True) + with pytest.raises(ValueError, match="No Azure credential resolved"): + _ = llm._client + mock_provider.assert_not_called() + + +def test_foundry_workload_identity_when_no_api_key(): + from kagent.adk.models import Foundry + + with ( + mock.patch.dict( + "os.environ", + { + "FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/", + "FOUNDRY_DEPLOYMENT": "gpt-4o", + }, + clear=True, + ), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch( + "kagent.adk.models._azure.azure_ad_token_provider", return_value=_SENTINEL_TOKEN_PROVIDER + ) as mock_provider, + ): + llm = Foundry(model="gpt-4o", type="foundry") + _ = llm._client + + mock_provider.assert_called_once() + kwargs = mock_azure.call_args.kwargs + assert kwargs["azure_ad_token_provider"] is _SENTINEL_TOKEN_PROVIDER + assert kwargs["api_key"] == API_KEY_SENTINEL + assert kwargs["azure_endpoint"] == "https://test.cognitiveservices.azure.com/" + assert kwargs["azure_deployment"] == "gpt-4o" + assert kwargs["api_version"] == "2024-10-21" + + +@pytest.mark.asyncio +async def test_azure_client_workload_identity_uses_bearer_with_real_sdk(): + seen_request: httpx.Request | None = None + + async def token_provider() -> str: + return "workload-token" + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal seen_request + seen_request = request + return httpx.Response( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + }, + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + with ( + mock.patch.dict("os.environ", {"AZURE_OPENAI_API_KEY": "must-not-win"}, clear=True), + mock.patch("kagent.adk.models._azure.azure_ad_token_provider", return_value=token_provider), + ): + client = build_azure_openai_client( + api_version="2024-10-21", + azure_endpoint="https://test.cognitiveservices.azure.com/", + azure_deployment="gpt-4o", + api_key=None, + api_key_passthrough=False, + default_headers=None, + http_client=http_client, + missing_credential_hint="missing credential", + ) + + try: + response = await client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "hello"}], + ) + finally: + await client.close() + + assert response.choices[0].message.content == "ok" + assert seen_request is not None + assert seen_request.headers["authorization"] == "Bearer workload-token" + assert "api-key" not in seen_request.headers + + +def test_foundry_api_key_from_env(): + from kagent.adk.models import Foundry + + with ( + mock.patch.dict( + "os.environ", + { + "FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/", + "FOUNDRY_DEPLOYMENT": "gpt-4o", + "FOUNDRY_API_VERSION": "2025-01-01", + "FOUNDRY_API_KEY": "secret", + }, + clear=True, + ), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch("kagent.adk.models._azure.azure_ad_token_provider") as mock_provider, + ): + llm = Foundry(model="gpt-4o", type="foundry") + _ = llm._client + + mock_provider.assert_not_called() + kwargs = mock_azure.call_args.kwargs + assert kwargs["api_key"] == "secret" + assert kwargs["azure_ad_token_provider"] is None + assert kwargs["api_version"] == "2025-01-01" + + +def test_foundry_sanitizes_default_auth_headers(): + from kagent.adk.models import Foundry + + with ( + mock.patch.dict( + "os.environ", + { + "FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/", + "FOUNDRY_DEPLOYMENT": "gpt-4o", + "FOUNDRY_API_KEY": "must-not-win", + }, + clear=True, + ), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch( + "kagent.adk.models._azure.azure_ad_token_provider", + return_value=_SENTINEL_TOKEN_PROVIDER, + ), + ): + llm = Foundry( + model="gpt-4o", + type="foundry", + default_headers={ + "Authorization": "Bearer leaked", + "api-key": "leaked", + "X-Custom": "preserved", + }, + ) + _ = llm._client + + assert mock_azure.call_args.kwargs["default_headers"] == {"X-Custom": "preserved"} + + +def test_foundry_passthrough_uses_caller_token(): + from kagent.adk.models import Foundry + + with ( + mock.patch.dict( + "os.environ", + {"FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/", "FOUNDRY_DEPLOYMENT": "gpt-4o"}, + clear=True, + ), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + mock.patch("kagent.adk.models._azure.azure_ad_token_provider") as mock_provider, + ): + llm = Foundry(model="gpt-4o", type="foundry", api_key_passthrough=True) + llm.set_passthrough_key("caller-token") + _ = llm._client + + mock_provider.assert_not_called() + assert mock_azure.call_args.kwargs["api_key"] == "caller-token" + + +def test_foundry_passthrough_without_token_does_not_use_workload_identity(): + from kagent.adk.models import Foundry + + with ( + mock.patch.dict( + "os.environ", + {"FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/", "FOUNDRY_DEPLOYMENT": "gpt-4o"}, + clear=True, + ), + mock.patch("kagent.adk.models._azure.azure_ad_token_provider") as mock_provider, + ): + llm = Foundry(model="gpt-4o", type="foundry", api_key_passthrough=True) + with pytest.raises(ValueError, match="No Azure credential resolved"): + _ = llm._client + + mock_provider.assert_not_called() + + +def test_foundry_client_with_tls(): + import ssl + + from kagent.adk.models import Foundry + + with ( + mock.patch("kagent.adk.models._ssl.create_ssl_context") as mock_create_ssl, + mock.patch("kagent.adk.models._openai.DefaultAsyncHttpxClient") as mock_httpx, + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure, + ): + mock_ssl_context = mock.MagicMock(spec=ssl.SSLContext) + mock_create_ssl.return_value = mock_ssl_context + mock_httpx_instance = mock.MagicMock() + mock_httpx.return_value = mock_httpx_instance + + llm = Foundry( + model="gpt-4o", + type="foundry", + endpoint="https://test.cognitiveservices.azure.com/", + deployment="gpt-4o", + api_key="foundry-key", + tls_ca_cert_path="/etc/ssl/certs/custom/corp-ca/ca.crt", + ) + _ = llm._client + + assert mock_httpx.call_args.kwargs["verify"] is mock_ssl_context + assert mock_azure.call_args.kwargs["http_client"] is mock_httpx_instance + + +def test_foundry_missing_endpoint_raises(): + from kagent.adk.models import Foundry + + with mock.patch.dict("os.environ", {"FOUNDRY_DEPLOYMENT": "gpt-4o"}, clear=True): + llm = Foundry(model="gpt-4o", type="foundry") + with pytest.raises(ValueError, match="Foundry endpoint must be provided"): + _ = llm._client + + +def test_foundry_missing_deployment_raises(): + from kagent.adk.models import Foundry + + with mock.patch.dict("os.environ", {"FOUNDRY_ENDPOINT": "https://test.cognitiveservices.azure.com/"}, clear=True): + llm = Foundry(model="gpt-4o", type="foundry") + with pytest.raises(ValueError, match="Foundry deployment must be provided"): + _ = llm._client diff --git a/python/packages/kagent-adk/tests/unittests/models/test_openai.py b/python/packages/kagent-adk/tests/unittests/models/test_openai.py index a0e512fee..78f961a5d 100644 --- a/python/packages/kagent-adk/tests/unittests/models/test_openai.py +++ b/python/packages/kagent-adk/tests/unittests/models/test_openai.py @@ -14,6 +14,7 @@ from unittest import mock +import httpx import pytest from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse @@ -745,7 +746,7 @@ def test_azure_openai_client_with_tls(): with mock.patch("kagent.adk.models._ssl.create_ssl_context") as mock_create_ssl: with mock.patch("kagent.adk.models._openai.DefaultAsyncHttpxClient") as mock_httpx: - with mock.patch("kagent.adk.models._openai.AsyncAzureOpenAI") as mock_azure_openai: + with mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_azure_openai: mock_ssl_context = mock.MagicMock(spec=ssl.SSLContext) mock_create_ssl.return_value = mock_ssl_context mock_httpx_instance = mock.MagicMock() diff --git a/python/packages/kagent-adk/tests/unittests/test_embedding.py b/python/packages/kagent-adk/tests/unittests/test_embedding.py index 0144f6fdf..333c6e71c 100644 --- a/python/packages/kagent-adk/tests/unittests/test_embedding.py +++ b/python/packages/kagent-adk/tests/unittests/test_embedding.py @@ -5,6 +5,7 @@ import numpy as np import pytest +from openai.lib.azure import API_KEY_SENTINEL from kagent.adk._bearer_token import bearer_token from kagent.adk.models import KAgentEmbedding @@ -19,11 +20,23 @@ def _reset_bearer_token(): def make_client( - provider: str, model: str, base_url: str | None = None, api_key_passthrough: bool = False + provider: str, + model: str, + base_url: str | None = None, + api_key_passthrough: bool = False, + endpoint: str | None = None, + deployment: str | None = None, + api_version: str | None = None, ) -> KAgentEmbedding: return KAgentEmbedding( config=EmbeddingConfig( - provider=provider, model=model, base_url=base_url, api_key_passthrough=api_key_passthrough + provider=provider, + model=model, + base_url=base_url, + api_key_passthrough=api_key_passthrough, + endpoint=endpoint, + deployment=deployment, + api_version=api_version, ), ) @@ -103,9 +116,13 @@ async def test_azure_openai_uses_azure_client(self): with ( mock.patch.dict( "os.environ", - {"OPENAI_API_VERSION": "2024-02-01", "AZURE_OPENAI_ENDPOINT": "https://myazure.openai.azure.com"}, + { + "OPENAI_API_VERSION": "2024-02-01", + "AZURE_OPENAI_ENDPOINT": "https://myazure.openai.azure.com", + "AZURE_OPENAI_API_KEY": "test-key", + }, ), - mock.patch("openai.AsyncAzureOpenAI") as mock_cls, + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_cls, ): instance = mock.AsyncMock() instance.embeddings.create = mock.AsyncMock(return_value=mock_response) @@ -114,6 +131,73 @@ async def test_azure_openai_uses_azure_client(self): assert result == vec assert mock_cls.called + @pytest.mark.asyncio + async def test_foundry_embed_uses_config_and_api_key(self): + client = make_client( + provider="foundry", + model="text-embedding-3-small", + endpoint="https://example.cognitiveservices.azure.com/", + deployment="embedding-deployment", + api_version="2025-01-01", + ) + mock_response = make_openai_embedding_response([[0.5] * 768]) + with ( + mock.patch.dict("os.environ", {"FOUNDRY_API_KEY": "foundry-key"}, clear=True), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_cls, + ): + instance = mock.AsyncMock() + instance.embeddings.create = mock.AsyncMock(return_value=mock_response) + mock_cls.return_value = instance + result = await client.generate("hello") + + assert result == [0.5] * 768 + assert mock_cls.call_args.kwargs["api_key"] == "foundry-key" + assert mock_cls.call_args.kwargs["azure_endpoint"] == "https://example.cognitiveservices.azure.com/" + assert mock_cls.call_args.kwargs["azure_deployment"] == "embedding-deployment" + assert mock_cls.call_args.kwargs["api_version"] == "2025-01-01" + + @pytest.mark.asyncio + async def test_foundry_embed_uses_workload_identity(self): + client = make_client( + provider="foundry", + model="text-embedding-3-small", + endpoint="https://example.cognitiveservices.azure.com/", + deployment="embedding-deployment", + ) + mock_response = make_openai_embedding_response([[0.5] * 768]) + token_provider = object() + with ( + mock.patch.dict("os.environ", {}, clear=True), + mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_cls, + mock.patch( + "kagent.adk.models._azure.azure_ad_token_provider", + return_value=token_provider, + ), + ): + instance = mock.AsyncMock() + instance.embeddings.create = mock.AsyncMock(return_value=mock_response) + mock_cls.return_value = instance + await client.generate("hello") + + assert mock_cls.call_args.kwargs["api_key"] == API_KEY_SENTINEL + assert mock_cls.call_args.kwargs["azure_ad_token_provider"] is token_provider + + @pytest.mark.asyncio + async def test_foundry_embed_requires_endpoint_and_deployment(self): + missing_endpoint = make_client(provider="foundry", model="text-embedding-3-small") + with mock.patch.dict("os.environ", {"FOUNDRY_DEPLOYMENT": "embedding-deployment"}, clear=True): + with pytest.raises(ValueError, match="Foundry endpoint must be provided"): + await missing_endpoint._embed_foundry(["hello"]) + + missing_deployment = make_client( + provider="foundry", + model="text-embedding-3-small", + endpoint="https://example.cognitiveservices.azure.com/", + ) + with mock.patch.dict("os.environ", {}, clear=True): + with pytest.raises(ValueError, match="Foundry deployment must be provided"): + await missing_deployment._embed_foundry(["hello"]) + @pytest.mark.asyncio async def test_ollama_embed(self): client = make_client(provider="ollama", model="nomic-embed-text") @@ -251,12 +335,46 @@ async def test_azure_openai_uses_bearer_token(self): ) bearer_token.set("the-callers-token") mock_response = make_openai_embedding_response([[0.1] * 768]) - with mock.patch("openai.AsyncAzureOpenAI") as mock_cls: + with mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_cls: + instance = mock.AsyncMock() + instance.embeddings.create = mock.AsyncMock(return_value=mock_response) + mock_cls.return_value = instance + await client.generate("hello") + assert mock_cls.call_args.kwargs["api_key"] == "the-callers-token" + + @pytest.mark.asyncio + async def test_foundry_uses_bearer_token(self): + client = make_client( + provider="foundry", + model="text-embedding-3-small", + endpoint="https://example.cognitiveservices.azure.com/", + deployment="embedding-deployment", + api_key_passthrough=True, + ) + bearer_token.set("the-callers-token") + mock_response = make_openai_embedding_response([[0.1] * 768]) + with mock.patch("kagent.adk.models._azure.AsyncAzureOpenAI") as mock_cls: instance = mock.AsyncMock() instance.embeddings.create = mock.AsyncMock(return_value=mock_response) mock_cls.return_value = instance await client.generate("hello") + assert mock_cls.call_args.kwargs["api_key"] == "the-callers-token" + assert mock_cls.call_args.kwargs["azure_ad_token_provider"] is None + + @pytest.mark.asyncio + async def test_foundry_passthrough_without_token_ignores_environment_key(self): + client = make_client( + provider="foundry", + model="text-embedding-3-small", + endpoint="https://example.cognitiveservices.azure.com/", + deployment="embedding-deployment", + api_key_passthrough=True, + ) + + with mock.patch.dict("os.environ", {"FOUNDRY_API_KEY": "must-not-win"}, clear=True): + with pytest.raises(ValueError, match="No Azure credential resolved"): + await client._embed_foundry(["hello"]) @pytest.mark.asyncio async def test_disabled_ignores_bearer_token(self): diff --git a/python/uv.lock b/python/uv.lock index 15d5d48b8..29f0991ed 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -393,6 +393,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" }, ] +[[package]] +name = "azure-core" +version = "1.41.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a6/f3/b416179e408990df5db0d516283022dde0f5d0111d98c1a848e41853e81c/azure_core-1.41.0.tar.gz", hash = "sha256:f46ff5dfcd230f25cf1c19e8a34b8dc08a337b2503e268bb600a16c00db8ad5a", size = 381042, upload-time = "2026-05-07T23:30:54.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/db/325c6d7312d2200251c52323878281045aaffcb5586612296484e4280eaa/azure_core-1.41.0-py3-none-any.whl", hash = "sha256:522b4011e8180b1a3dcd2024396a4e7fe9ac37fb8597db47163d230b5efe892d", size = 220920, upload-time = "2026-05-07T23:30:56.357Z" }, +] + +[[package]] +name = "azure-identity" +version = "1.25.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "msal" }, + { name = "msal-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/0e/3a63efb48aa4a5ae2cfca61ee152fbcb668092134d3eb8bfda472dd5c617/azure_identity-1.25.3.tar.gz", hash = "sha256:ab23c0d63015f50b630ef6c6cf395e7262f439ce06e5d07a64e874c724f8d9e6", size = 286304, upload-time = "2026-03-13T01:12:20.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" }, +] + [[package]] name = "backoff" version = "2.2.1" @@ -1795,6 +1824,7 @@ dependencies = [ { name = "aiofiles" }, { name = "anthropic", extra = ["vertex"] }, { name = "anyio" }, + { name = "azure-identity" }, { name = "boto3" }, { name = "fastapi" }, { name = "filelock" }, @@ -1839,6 +1869,7 @@ requires-dist = [ { name = "aiofiles", specifier = ">=24.1.0" }, { name = "anthropic", extras = ["vertex"], specifier = ">=0.120.0" }, { name = "anyio", specifier = ">=4.14.2" }, + { name = "azure-identity", specifier = ">=1.19.0" }, { name = "boto3", specifier = ">=1.28.57" }, { name = "fastapi", specifier = ">=0.115.1" }, { name = "filelock", specifier = ">=3.20.3" }, @@ -2572,6 +2603,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] +[[package]] +name = "msal" +version = "1.37.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9a/99/d840198ecf6e8057bbc937f129ae940404485d736cda73253bbff9537f01/msal-1.37.0.tar.gz", hash = "sha256:1b1672a33ee467c1d70b341bb16cafd51bb3c817147a95b93263794b03971bec", size = 182444, upload-time = "2026-05-29T19:49:05.561Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/b0/d807279f4b55d16d1f120d5ac4344c6e39b56732e2a224d40bded7fd67ad/msal-1.37.0-py3-none-any.whl", hash = "sha256:dd17e95a7c71bce75e8108113438ba7c4a086b3bcad4f57a8c09b7af3d753c2d", size = 123725, upload-time = "2026-05-29T19:49:04.335Z" }, +] + +[[package]] +name = "msal-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/99/5d239b6156eddf761a636bded1118414d161bd6b7b37a9335549ed159396/msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4", size = 23315, upload-time = "2025-03-14T23:51:03.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, +] + [[package]] name = "multidict" version = "6.6.3" From 41dce82355073b9a07b956af41b0e705e16ce01d Mon Sep 17 00:00:00 2001 From: Mark Rossetti Date: Thu, 20 Aug 2026 12:19:53 -0700 Subject: [PATCH 2/3] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mark Rossetti --- python/packages/kagent-adk/tests/unittests/models/test_openai.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/packages/kagent-adk/tests/unittests/models/test_openai.py b/python/packages/kagent-adk/tests/unittests/models/test_openai.py index 78f961a5d..83f1cd7f6 100644 --- a/python/packages/kagent-adk/tests/unittests/models/test_openai.py +++ b/python/packages/kagent-adk/tests/unittests/models/test_openai.py @@ -14,7 +14,6 @@ from unittest import mock -import httpx import pytest from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse From 6c0076062c1c52a21d0e5837f5e04b1d3430e25b Mon Sep 17 00:00:00 2001 From: Mark Rossetti Date: Wed, 26 Aug 2026 20:09:17 +0000 Subject: [PATCH 3/3] drop embedding size for foundry Signed-off-by: Mark Rossetti --- .../packages/kagent-adk/src/kagent/adk/models/_embedding.py | 1 - python/packages/kagent-adk/tests/unittests/test_embedding.py | 4 ++++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py index 611ea9dc8..26f41a90a 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py @@ -214,7 +214,6 @@ async def _embed_foundry(self, texts: List[str]) -> List[List[float]]: response = await client.embeddings.create( model=self.config.model, input=texts, - dimensions=self.TARGET_DIMENSION, ) return [item.embedding for item in response.data] diff --git a/python/packages/kagent-adk/tests/unittests/test_embedding.py b/python/packages/kagent-adk/tests/unittests/test_embedding.py index 333c6e71c..a0b8819f6 100644 --- a/python/packages/kagent-adk/tests/unittests/test_embedding.py +++ b/python/packages/kagent-adk/tests/unittests/test_embedding.py @@ -155,6 +155,10 @@ async def test_foundry_embed_uses_config_and_api_key(self): assert mock_cls.call_args.kwargs["azure_endpoint"] == "https://example.cognitiveservices.azure.com/" assert mock_cls.call_args.kwargs["azure_deployment"] == "embedding-deployment" assert mock_cls.call_args.kwargs["api_version"] == "2025-01-01" + instance.embeddings.create.assert_called_once_with( + model="text-embedding-3-small", + input=["hello"], + ) @pytest.mark.asyncio async def test_foundry_embed_uses_workload_identity(self):