diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.py index 950bfaf73334..8b68c0515aec 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.py @@ -24,6 +24,86 @@ logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Shared helpers used by both the sync and async AIProjectClient.get_openai_client() +# implementations. Defined at module level so the async client can import and reuse +# them without duplicating the logic. +# --------------------------------------------------------------------------- + + +def _resolve_openai_base_url(config: Any, agent_name: Optional[str], kwargs: dict) -> str: + """Resolve the base URL for the (Async)OpenAI client. + + :param config: Generated client configuration carrying ``endpoint`` and ``allow_preview``. + :type config: Any + :param agent_name: Optional hosted-agent name. + :type agent_name: str or None + :param kwargs: Caller keyword arguments; ``base_url`` is popped when present. + :type kwargs: dict + :return: The base URL to use for the (Async)OpenAI client. + :rtype: str + :raises ValueError: If ``agent_name`` is provided but ``allow_preview=True`` was not set. + """ + if "base_url" in kwargs: + return kwargs.pop("base_url") + if agent_name is not None: + if config.allow_preview: + return config.endpoint.rstrip("/") + f"/agents/{agent_name}/endpoint/protocols/openai" + raise ValueError( + "Calling `get_openai_client` method with an `agent_name` requires you to set `allow_preview=True`" + "\nwhen constructing the AIProjectClient. Note that preview features are under development and " + "\nsubject to change. They should not be used in production environments." + ) + return config.endpoint.rstrip("/") + "/openai/v1" + + +def _resolve_openai_query_params(config: Any, agent_name: Optional[str], kwargs: dict) -> dict: + """Build the ``default_query`` dict for the (Async)OpenAI client. + + :param config: Generated client configuration carrying ``api_version``. + :type config: Any + :param agent_name: Optional hosted-agent name. + :type agent_name: str or None + :param kwargs: Caller keyword arguments; ``default_query`` is popped when present. + :type kwargs: dict + :return: Query parameters to forward to the (Async)OpenAI client. + :rtype: dict + """ + default_query = dict[str, str](kwargs.pop("default_query", None) or {}) + if agent_name is not None and "api-version" not in default_query: + default_query["api-version"] = config.api_version + return default_query + + +def _resolve_openai_default_headers(agent_name: Optional[str], kwargs: dict) -> dict: + """Build the ``default_headers`` dict for the (Async)OpenAI client. + + :param agent_name: Optional hosted-agent name. + :type agent_name: str or None + :param kwargs: Caller keyword arguments; ``default_headers`` is popped when present. + :type kwargs: dict + :return: Headers to forward to the (Async)OpenAI client. + :rtype: dict + """ + default_headers = dict[str, str](kwargs.pop("default_headers", None) or {}) + if agent_name is not None and not _has_header_case_insensitive(default_headers, _FOUNDRY_FEATURES_HEADER_NAME): + default_headers[_FOUNDRY_FEATURES_HEADER_NAME] = _BETA_OPERATION_FEATURE_HEADERS["agents"] + return default_headers + + +def _build_openai_user_agent(custom_user_agent: Optional[str], openai_default_user_agent: str) -> str: + """Build the SDK-prefixed User-Agent string for the (Async)OpenAI client. + + :param custom_user_agent: Caller-supplied user_agent kwarg captured at construction time. + :type custom_user_agent: str or None + :param openai_default_user_agent: The OpenAI client's own default user-agent. + :type openai_default_user_agent: str + :return: Combined User-Agent string. + :rtype: str + """ + return "-".join(ua for ua in [custom_user_agent, "AIProjectClient"] if ua) + " " + openai_default_user_agent + + class AIProjectClient(AIProjectClientGenerated): # pylint: disable=too-many-instance-attributes """AIProjectClient. @@ -101,6 +181,35 @@ def __init__( self.telemetry = TelemetryOperations(self) # type: ignore + def _get_openai_api_key(self, kwargs: dict): + """Resolve the API key for the OpenAI client. + + :param kwargs: Caller keyword arguments; ``api_key`` is popped when present. + :type kwargs: dict + :return: The API key string or a bearer-token-provider callable. + :rtype: str or Callable + """ + if "api_key" in kwargs: + return kwargs.pop("api_key") + return get_bearer_token_provider( + self._config.credential, # pylint: disable=protected-access + "https://ai.azure.com/.default", + ) + + def _get_openai_http_client(self, kwargs: dict): + """Resolve the HTTP transport client for the OpenAI client. + + :param kwargs: Caller keyword arguments; ``http_client`` is popped when present. + :type kwargs: dict + :return: An httpx.Client instance configured with logging transport, or ``None``. + :rtype: httpx.Client or None + """ + if "http_client" in kwargs: + return kwargs.pop("http_client") + if self._console_logging_enabled: + return httpx.Client(transport=_OpenAILoggingTransport()) + return None + @distributed_trace def get_openai_client(self, *, agent_name: Optional[str] = None, **kwargs: Any) -> OpenAI: """Get an authenticated OpenAI client from the `openai` package. @@ -131,51 +240,17 @@ def get_openai_client(self, *, agent_name: Optional[str] = None, **kwargs: Any) kwargs = kwargs.copy() if kwargs else {} - # Allow caller to override base_url - if "base_url" in kwargs: - base_url = kwargs.pop("base_url") - elif agent_name is not None: - if self._config.allow_preview: - base_url = ( - self._config.endpoint.rstrip("/") + f"/agents/{agent_name}/endpoint/protocols/openai" - ) # pylint: disable=protected-access - else: - raise ValueError( - "Calling `get_openai_client` method with an `agent_name` requires you to set `allow_preview=True`" - "\nwhen constructing the AIProjectClient. Note that preview features are under development and " - "\nsubject to change. They should not be used in production environments." - ) - else: - base_url = self._config.endpoint.rstrip("/") + "/openai/v1" # pylint: disable=protected-access - - default_query = dict[str, str](kwargs.pop("default_query", None) or {}) - if agent_name is not None and "api-version" not in default_query: - default_query["api-version"] = self._config.api_version # pylint: disable=protected-access + base_url = _resolve_openai_base_url(self._config, agent_name, kwargs) + default_query = _resolve_openai_query_params(self._config, agent_name, kwargs) logger.debug( # pylint: disable=specify-parameter-names-in-call "[get_openai_client] Creating OpenAI client using Entra ID authentication, base_url = `%s`", # pylint: disable=line-too-long base_url, ) - # Allow caller to override api_key, otherwise use token provider - if "api_key" in kwargs: - api_key = kwargs.pop("api_key") - else: - api_key = get_bearer_token_provider( - self._config.credential, # pylint: disable=protected-access - "https://ai.azure.com/.default", - ) - - if "http_client" in kwargs: - http_client = kwargs.pop("http_client") - elif self._console_logging_enabled: - http_client = httpx.Client(transport=_OpenAILoggingTransport()) - else: - http_client = None - - default_headers = dict[str, str](kwargs.pop("default_headers", None) or {}) - if agent_name is not None and not _has_header_case_insensitive(default_headers, _FOUNDRY_FEATURES_HEADER_NAME): - default_headers[_FOUNDRY_FEATURES_HEADER_NAME] = _BETA_OPERATION_FEATURE_HEADERS["agents"] + api_key = self._get_openai_api_key(kwargs) + http_client = self._get_openai_http_client(kwargs) + default_headers = _resolve_openai_default_headers(agent_name, kwargs) openai_custom_user_agent = default_headers.get("User-Agent", None) @@ -195,11 +270,7 @@ def _create_openai_client(**kwargs) -> OpenAI: if openai_custom_user_agent: final_user_agent = openai_custom_user_agent else: - final_user_agent = ( - "-".join(ua for ua in [self._custom_user_agent, "AIProjectClient"] if ua) - + " " - + openai_default_user_agent - ) + final_user_agent = _build_openai_user_agent(self._custom_user_agent, openai_default_user_agent) default_headers["User-Agent"] = final_user_agent diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.pyi b/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.pyi index c40454fc9320..42009c1c227e 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.pyi +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/_patch.pyi @@ -108,6 +108,11 @@ class AIProjectClient(AIProjectClientGenerated): # To make mypy happy... otherwise imports of the below result in mypy "attr-defined" error class _AuthSecretsFilter(logging.Filter): ... +def _resolve_openai_base_url(config: Any, agent_name: Optional[str], kwargs: dict) -> str: ... +def _resolve_openai_query_params(config: Any, agent_name: Optional[str], kwargs: dict) -> dict: ... +def _resolve_openai_default_headers(agent_name: Optional[str], kwargs: dict) -> dict: ... +def _build_openai_user_agent(custom_user_agent: Optional[str], openai_default_user_agent: str) -> str: ... + __all__: List[str] = ["AIProjectClient"] def patch_sdk() -> None: ... diff --git a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/_patch.py b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/_patch.py index 34528053a55c..6eafa71053e6 100644 --- a/sdk/ai/azure-ai-projects/azure/ai/projects/aio/_patch.py +++ b/sdk/ai/azure-ai-projects/azure/ai/projects/aio/_patch.py @@ -16,8 +16,13 @@ from azure.core.tracing.decorator import distributed_trace from azure.core.credentials_async import AsyncTokenCredential from azure.identity.aio import get_bearer_token_provider -from .._patch import _AuthSecretsFilter -from ..models._patch import _BETA_OPERATION_FEATURE_HEADERS, _FOUNDRY_FEATURES_HEADER_NAME, _has_header_case_insensitive +from .._patch import ( + _AuthSecretsFilter, + _build_openai_user_agent, + _resolve_openai_base_url, + _resolve_openai_default_headers, + _resolve_openai_query_params, +) from ._client import AIProjectClient as AIProjectClientGenerated from .operations import TelemetryOperations @@ -101,6 +106,35 @@ def __init__( self.telemetry = TelemetryOperations(self) # type: ignore + def _get_openai_api_key(self, kwargs: dict): + """Resolve the API key for the AsyncOpenAI client. + + :param kwargs: Caller keyword arguments; ``api_key`` is popped when present. + :type kwargs: dict + :return: The API key string or a bearer-token-provider callable. + :rtype: str or Callable + """ + if "api_key" in kwargs: + return kwargs.pop("api_key") + return get_bearer_token_provider( + self._config.credential, # pylint: disable=protected-access + "https://ai.azure.com/.default", + ) + + def _get_openai_http_client(self, kwargs: dict): + """Resolve the HTTP transport client for the AsyncOpenAI client. + + :param kwargs: Caller keyword arguments; ``http_client`` is popped when present. + :type kwargs: dict + :return: An httpx.AsyncClient instance configured with logging transport, or ``None``. + :rtype: httpx.AsyncClient or None + """ + if "http_client" in kwargs: + return kwargs.pop("http_client") + if self._console_logging_enabled: + return httpx.AsyncClient(transport=_OpenAILoggingTransport()) + return None + @distributed_trace def get_openai_client(self, *, agent_name: Optional[str] = None, **kwargs: Any) -> AsyncOpenAI: """Get an authenticated AsyncOpenAI client from the `openai` package. @@ -131,51 +165,17 @@ def get_openai_client(self, *, agent_name: Optional[str] = None, **kwargs: Any) kwargs = kwargs.copy() if kwargs else {} - # Allow caller to override base_url - if "base_url" in kwargs: - base_url = kwargs.pop("base_url") - elif agent_name is not None: - if self._config.allow_preview: - base_url = ( - self._config.endpoint.rstrip("/") + f"/agents/{agent_name}/endpoint/protocols/openai" - ) # pylint: disable=protected-access - else: - raise ValueError( - "Calling `get_openai_client` method with an `agent_name` requires you to set `allow_preview=True`" - "\nwhen constructing the AIProjectClient. Note that preview features are under development and " - "\nsubject to change. They should not be used in production environments." - ) - else: - base_url = self._config.endpoint.rstrip("/") + "/openai/v1" # pylint: disable=protected-access - - default_query = dict[str, str](kwargs.pop("default_query", None) or {}) - if agent_name is not None and "api-version" not in default_query: - default_query["api-version"] = self._config.api_version # pylint: disable=protected-access + base_url = _resolve_openai_base_url(self._config, agent_name, kwargs) + default_query = _resolve_openai_query_params(self._config, agent_name, kwargs) logger.debug( # pylint: disable=specify-parameter-names-in-call "[get_openai_client] Creating OpenAI client using Entra ID authentication, base_url = `%s`", # pylint: disable=line-too-long base_url, ) - # Allow caller to override api_key, otherwise use token provider - if "api_key" in kwargs: - api_key = kwargs.pop("api_key") - else: - api_key = get_bearer_token_provider( - self._config.credential, # pylint: disable=protected-access - "https://ai.azure.com/.default", - ) - - if "http_client" in kwargs: - http_client = kwargs.pop("http_client") - elif self._console_logging_enabled: - http_client = httpx.AsyncClient(transport=_OpenAILoggingTransport()) - else: - http_client = None - - default_headers = dict[str, str](kwargs.pop("default_headers", None) or {}) - if agent_name is not None and not _has_header_case_insensitive(default_headers, _FOUNDRY_FEATURES_HEADER_NAME): - default_headers[_FOUNDRY_FEATURES_HEADER_NAME] = _BETA_OPERATION_FEATURE_HEADERS["agents"] + api_key = self._get_openai_api_key(kwargs) + http_client = self._get_openai_http_client(kwargs) + default_headers = _resolve_openai_default_headers(agent_name, kwargs) openai_custom_user_agent = default_headers.get("User-Agent", None) @@ -195,11 +195,7 @@ def _create_openai_client(**kwargs) -> AsyncOpenAI: if openai_custom_user_agent: final_user_agent = openai_custom_user_agent else: - final_user_agent = ( - "-".join(ua for ua in [self._custom_user_agent, "AIProjectClient"] if ua) - + " " - + openai_default_user_agent - ) + final_user_agent = _build_openai_user_agent(self._custom_user_agent, openai_default_user_agent) default_headers["User-Agent"] = final_user_agent diff --git a/sdk/ai/azure-ai-projects/tests/responses/openai_test_helpers.py b/sdk/ai/azure-ai-projects/tests/responses/openai_test_helpers.py new file mode 100644 index 000000000000..fbf3637291ea --- /dev/null +++ b/sdk/ai/azure-ai-projects/tests/responses/openai_test_helpers.py @@ -0,0 +1,69 @@ +# pylint: disable=line-too-long,useless-suppression +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +"""Shared helpers for unit-testing AIProjectClient.get_openai_client (sync and async). + +These helpers build lightweight client stubs that bypass the real ``__init__`` so unit +tests can target individual branches of ``get_openai_client`` without making any +network calls. +""" + +from typing import Optional +from unittest.mock import MagicMock + +from azure.ai.projects import AIProjectClient +from azure.ai.projects.aio import AIProjectClient as AsyncAIProjectClient + +ENDPOINT = "https://myaccount.services.ai.azure.com/api/projects/myproject" +API_VERSION = "2025-01-01" + +# Patch targets used by tests to swap in mocked OpenAI/AsyncOpenAI constructors +# and bearer-token providers. +SYNC_OPENAI_PATCH = "azure.ai.projects._patch.OpenAI" +ASYNC_OPENAI_PATCH = "azure.ai.projects.aio._patch.AsyncOpenAI" +SYNC_TOKEN_PROVIDER_PATCH = "azure.ai.projects._patch.get_bearer_token_provider" +ASYNC_TOKEN_PROVIDER_PATCH = "azure.ai.projects.aio._patch.get_bearer_token_provider" + + +def make_sync_client( + allow_preview: bool = True, + console_logging: bool = False, + custom_user_agent: Optional[str] = None, +) -> AIProjectClient: + """Return a minimal sync AIProjectClient stub suitable for unit-testing get_openai_client.""" + client = AIProjectClient.__new__(AIProjectClient) + client._config = MagicMock() + client._config.endpoint = ENDPOINT + client._config.allow_preview = allow_preview + client._config.api_version = API_VERSION + client._config.credential = MagicMock() + client._console_logging_enabled = console_logging + client._custom_user_agent = custom_user_agent + return client + + +def make_async_client( + allow_preview: bool = True, + console_logging: bool = False, + custom_user_agent: Optional[str] = None, +) -> AsyncAIProjectClient: + """Return a minimal async AIProjectClient stub suitable for unit-testing get_openai_client.""" + client = AsyncAIProjectClient.__new__(AsyncAIProjectClient) + client._config = MagicMock() + client._config.endpoint = ENDPOINT + client._config.allow_preview = allow_preview + client._config.api_version = API_VERSION + client._config.credential = MagicMock() + client._console_logging_enabled = console_logging + client._custom_user_agent = custom_user_agent + return client + + +def mock_openai(user_agent: str = "openai/1.0"): + """Return ``(mock_class, mock_instance)`` where ``mock_class`` acts as the OpenAI constructor.""" + instance = MagicMock() + instance.user_agent = user_agent + mock_cls = MagicMock(return_value=instance) + return mock_cls, instance diff --git a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint.py b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint.py index 87aecf0a94fa..be5d6f429040 100644 --- a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint.py +++ b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint.py @@ -4,13 +4,26 @@ # Licensed under the MIT License. # ------------------------------------ """ -Unit tests for verifying the base_url of the OpenAI client returned by AIProjectClient.get_openai_client(). +Unit tests for verifying the base_url, default_query (api-version) and Foundry feature +header behavior of the OpenAI client returned by AIProjectClient.get_openai_client(). No network calls are made. """ +from unittest.mock import patch + import pytest from azure.core.credentials import AccessToken from azure.ai.projects import AIProjectClient +from azure.ai.projects.models._patch import _FOUNDRY_FEATURES_HEADER_NAME + +from openai_test_helpers import ( + ENDPOINT, + API_VERSION, + SYNC_OPENAI_PATCH, + SYNC_TOKEN_PROVIDER_PATCH, + make_sync_client, + mock_openai, +) FAKE_ENDPOINT = "https://fake-account.services.ai.azure.com/api/projects/fake-project" AGENT_NAME = "fake-agent-name" @@ -59,3 +72,106 @@ def test_get_openai_client_with_agent_name_and_allow_preview(self): expected_base_url = FAKE_ENDPOINT.rstrip("/") + f"/agents/{AGENT_NAME}/endpoint/protocols/openai" assert str(openai_client.base_url).rstrip("/") == expected_base_url + + def test_trailing_slash_on_endpoint_is_stripped(self): + """Trailing slash on endpoint must not produce a double slash in the base URL.""" + client = make_sync_client() + client._config.endpoint = ENDPOINT + "/" + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + host = c.kwargs["base_url"].replace("https://", "") + assert "//" not in host + + +# =========================================================================== +# default_query / api-version injection branches +# =========================================================================== + + +class TestDefaultQueryBranches: + def test_no_agent_no_api_version_injected(self): + """Branch: no agent_name -> api-version is NOT injected into default_query.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + assert "api-version" not in c.kwargs.get("default_query", {}) + + def test_agent_injects_api_version(self): + """Branch: agent_name + no caller api-version -> inject SDK api_version.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent") + for c in mock_cls.call_args_list: + assert c.kwargs["default_query"]["api-version"] == API_VERSION + + def test_agent_does_not_override_caller_api_version(self): + """Branch: agent_name + caller-provided api-version -> keep caller value.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent", default_query={"api-version": "caller-v"}) + for c in mock_cls.call_args_list: + assert c.kwargs["default_query"]["api-version"] == "caller-v" + + def test_caller_default_query_values_preserved(self): + """Caller-supplied default_query keys other than api-version are preserved.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(default_query={"foo": "bar"}) + for c in mock_cls.call_args_list: + assert c.kwargs["default_query"]["foo"] == "bar" + + +# =========================================================================== +# Foundry feature header injection branches +# =========================================================================== + + +class TestFoundryFeatureHeaderBranches: + def test_no_agent_no_feature_header_injected(self): + """Branch: no agent_name -> Foundry feature header is NOT injected.""" + + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + real_call = mock_cls.call_args_list[-1] + assert _FOUNDRY_FEATURES_HEADER_NAME not in real_call.kwargs.get("default_headers", {}) + + def test_agent_injects_foundry_feature_header(self): + """Branch: agent_name + header not present -> inject the Foundry feature header.""" + + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent") + real_call = mock_cls.call_args_list[-1] + assert _FOUNDRY_FEATURES_HEADER_NAME in real_call.kwargs.get("default_headers", {}) + + def test_agent_does_not_override_existing_feature_header(self): + """Branch: agent_name + header already in caller headers -> keep caller value.""" + + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client( + agent_name="my-agent", + default_headers={_FOUNDRY_FEATURES_HEADER_NAME: "caller-value"}, + ) + real_call = mock_cls.call_args_list[-1] + assert real_call.kwargs["default_headers"][_FOUNDRY_FEATURES_HEADER_NAME] == "caller-value" + + def test_caller_headers_other_keys_preserved(self): + """Caller-supplied non-feature headers are passed through to the OpenAI client.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(default_headers={"X-My-Header": "hello"}) + real_call = mock_cls.call_args_list[-1] + assert real_call.kwargs["default_headers"]["X-My-Header"] == "hello" diff --git a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint_async.py b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint_async.py index f489a4f982b6..8335531787af 100644 --- a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint_async.py +++ b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_endpoint_async.py @@ -4,14 +4,27 @@ # Licensed under the MIT License. # ------------------------------------ """ -Unit tests for verifying the base_url of the AsyncOpenAI client returned by AIProjectClient.get_openai_client(). +Unit tests for verifying the base_url, default_query (api-version) and Foundry feature +header behavior of the AsyncOpenAI client returned by AIProjectClient.get_openai_client(). No network calls are made. """ -import pytest from typing import Any +from unittest.mock import patch + +import pytest from azure.core.credentials_async import AsyncTokenCredential from azure.ai.projects.aio import AIProjectClient +from azure.ai.projects.models._patch import _FOUNDRY_FEATURES_HEADER_NAME + +from openai_test_helpers import ( + ENDPOINT, + API_VERSION, + ASYNC_OPENAI_PATCH, + ASYNC_TOKEN_PROVIDER_PATCH, + make_async_client, + mock_openai, +) FAKE_ENDPOINT = "https://fake-account.services.ai.azure.com/api/projects/fake-project" AGENT_NAME = "fake-agent-name" @@ -68,3 +81,95 @@ async def test_get_openai_client_with_agent_name_and_allow_preview_async(self): expected_base_url = FAKE_ENDPOINT.rstrip("/") + f"/agents/{AGENT_NAME}/endpoint/protocols/openai" assert str(openai_client.base_url).rstrip("/") == expected_base_url + + @pytest.mark.asyncio + async def test_trailing_slash_on_endpoint_is_stripped_async(self): + """Trailing slash on endpoint must not produce a double slash in the base URL.""" + client = make_async_client() + client._config.endpoint = ENDPOINT + "/" + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + host = c.kwargs["base_url"].replace("https://", "") + assert "//" not in host + + +# =========================================================================== +# default_query / api-version injection branches (async) +# =========================================================================== + + +class TestDefaultQueryBranchesAsync: + @pytest.mark.asyncio + async def test_no_agent_no_api_version_injected(self): + """Branch: no agent_name -> api-version is NOT injected into default_query.""" + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + assert "api-version" not in c.kwargs.get("default_query", {}) + + @pytest.mark.asyncio + async def test_agent_injects_api_version(self): + """Branch: agent_name + no caller api-version -> inject SDK api_version.""" + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent") + for c in mock_cls.call_args_list: + assert c.kwargs["default_query"]["api-version"] == API_VERSION + + @pytest.mark.asyncio + async def test_agent_does_not_override_caller_api_version(self): + """Branch: agent_name + caller-provided api-version -> keep caller value.""" + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent", default_query={"api-version": "caller-v"}) + for c in mock_cls.call_args_list: + assert c.kwargs["default_query"]["api-version"] == "caller-v" + + +# =========================================================================== +# Foundry feature header injection branches (async) +# =========================================================================== + + +class TestFoundryFeatureHeaderBranchesAsync: + @pytest.mark.asyncio + async def test_no_agent_no_feature_header_injected(self): + """Branch: no agent_name -> Foundry feature header is NOT injected.""" + + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + real_call = mock_cls.call_args_list[-1] + assert _FOUNDRY_FEATURES_HEADER_NAME not in real_call.kwargs.get("default_headers", {}) + + @pytest.mark.asyncio + async def test_agent_injects_foundry_feature_header(self): + """Branch: agent_name + header not present -> inject the Foundry feature header.""" + + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client(agent_name="my-agent") + real_call = mock_cls.call_args_list[-1] + assert _FOUNDRY_FEATURES_HEADER_NAME in real_call.kwargs.get("default_headers", {}) + + @pytest.mark.asyncio + async def test_agent_does_not_override_existing_feature_header(self): + """Branch: agent_name + header already in caller headers -> keep caller value.""" + + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client( + agent_name="my-agent", + default_headers={_FOUNDRY_FEATURES_HEADER_NAME: "caller-value"}, + ) + real_call = mock_cls.call_args_list[-1] + assert real_call.kwargs["default_headers"][_FOUNDRY_FEATURES_HEADER_NAME] == "caller-value" diff --git a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides.py b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides.py index 782c5024f43a..e0895ae552ca 100644 --- a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides.py +++ b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides.py @@ -4,17 +4,27 @@ # Licensed under the MIT License. # ------------------------------------ """ -Tests to verify that a custom http_client can be passed to get_openai_client() -and that the returned OpenAI client uses it instead of the default one. +Tests covering caller-side overrides (http_client, api_key, base_url, default_headers) +and the user-agent / token-provider / logging-transport branches of +AIProjectClient.get_openai_client() (sync). """ import os from typing import Any +from unittest.mock import patch + import pytest import httpx from azure.core.credentials import TokenCredential from azure.ai.projects import AIProjectClient +from openai_test_helpers import ( + SYNC_OPENAI_PATCH, + SYNC_TOKEN_PROVIDER_PATCH, + make_sync_client, + mock_openai, +) + class DummyTokenCredential(TokenCredential): """A dummy credential that returns None for testing purposes.""" @@ -134,3 +144,60 @@ def test_api_key_and_base_url_overrides(self): assert ( str(openai_client.base_url) == custom_base_url + "/" ), f"Expected base_url '{custom_base_url}/', got '{openai_client.base_url}'" + + +# =========================================================================== +# api_key resolution branches +# =========================================================================== + + +class TestApiKeyBranches: + def test_token_provider_used_when_no_api_key(self): + """Branch: no 'api_key' kwarg -> get_bearer_token_provider() is invoked.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="provider") as mock_tp: + client.get_openai_client() + mock_tp.assert_called_once_with(client._config.credential, "https://ai.azure.com/.default") + for c in mock_cls.call_args_list: + assert c.kwargs["api_key"] == "provider" + + def test_caller_api_key_skips_token_provider(self): + """Branch: 'api_key' in kwargs -> token provider is NOT called.""" + client = make_sync_client() + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH) as mock_tp: + client.get_openai_client(api_key="my-secret-key") + mock_tp.assert_not_called() + for c in mock_cls.call_args_list: + assert c.kwargs["api_key"] == "my-secret-key" + + +# =========================================================================== +# http_client resolution branches +# =========================================================================== + + +class TestHttpClientBranches: + def test_http_client_is_none_by_default(self): + """Branch: no override + console logging off -> http_client is None.""" + client = make_sync_client(console_logging=False) + mock_cls, _ = mock_openai() + with patch(SYNC_OPENAI_PATCH, mock_cls), patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + assert c.kwargs["http_client"] is None + + def test_console_logging_creates_logging_transport(self): + """Branch: no override + _console_logging_enabled=True -> httpx.Client with logging transport.""" + client = make_sync_client(console_logging=True) + mock_cls, _ = mock_openai() + with ( + patch(SYNC_OPENAI_PATCH, mock_cls), + patch(SYNC_TOKEN_PROVIDER_PATCH, return_value="tok"), + patch("azure.ai.projects._patch.httpx") as mock_httpx, + patch("azure.ai.projects._patch._OpenAILoggingTransport"), + ): + mock_httpx.Client.return_value = object() + client.get_openai_client() + mock_httpx.Client.assert_called_once() diff --git a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides_async.py b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides_async.py index de8c484fd9f6..220dc961d22a 100644 --- a/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides_async.py +++ b/sdk/ai/azure-ai-projects/tests/responses/test_openai_client_overrides_async.py @@ -4,17 +4,27 @@ # Licensed under the MIT License. # ------------------------------------ """ -Tests to verify that a custom http_client can be passed to get_openai_client() -and that the returned AsyncOpenAI client uses it instead of the default one. +Tests covering caller-side overrides (http_client, api_key, base_url, default_headers) +and the user-agent / token-provider / logging-transport branches of +AIProjectClient.get_openai_client() (async). """ import os from typing import Any +from unittest.mock import patch + import pytest import httpx from azure.core.credentials_async import AsyncTokenCredential from azure.ai.projects.aio import AIProjectClient +from openai_test_helpers import ( + ASYNC_OPENAI_PATCH, + ASYNC_TOKEN_PROVIDER_PATCH, + make_async_client, + mock_openai, +) + class DummyAsyncTokenCredential(AsyncTokenCredential): """A dummy async credential that returns None for testing purposes.""" @@ -138,3 +148,67 @@ async def test_api_key_and_base_url_overrides_async(self): assert ( str(openai_client.base_url) == custom_base_url + "/" ), f"Expected base_url '{custom_base_url}/', got '{openai_client.base_url}'" + + +# =========================================================================== +# api_key resolution branches (async) +# =========================================================================== + + +class TestApiKeyBranchesAsync: + @pytest.mark.asyncio + async def test_token_provider_used_when_no_api_key(self): + """Branch: no 'api_key' kwarg -> get_bearer_token_provider() is invoked.""" + client = make_async_client() + mock_cls, _ = mock_openai() + with ( + patch(ASYNC_OPENAI_PATCH, mock_cls), + patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="async-provider") as mock_tp, + ): + client.get_openai_client() + mock_tp.assert_called_once_with(client._config.credential, "https://ai.azure.com/.default") + for c in mock_cls.call_args_list: + assert c.kwargs["api_key"] == "async-provider" + + @pytest.mark.asyncio + async def test_caller_api_key_skips_token_provider(self): + """Branch: 'api_key' in kwargs -> token provider is NOT called.""" + client = make_async_client() + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH) as mock_tp: + client.get_openai_client(api_key="async-secret") + mock_tp.assert_not_called() + for c in mock_cls.call_args_list: + assert c.kwargs["api_key"] == "async-secret" + + +# =========================================================================== +# http_client resolution branches (async) +# =========================================================================== + + +class TestHttpClientBranchesAsync: + @pytest.mark.asyncio + async def test_http_client_is_none_by_default(self): + """Branch: no override + console logging off -> http_client is None.""" + client = make_async_client(console_logging=False) + mock_cls, _ = mock_openai() + with patch(ASYNC_OPENAI_PATCH, mock_cls), patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"): + client.get_openai_client() + for c in mock_cls.call_args_list: + assert c.kwargs["http_client"] is None + + @pytest.mark.asyncio + async def test_console_logging_creates_async_logging_transport(self): + """Branch: no override + _console_logging_enabled=True -> httpx.AsyncClient with logging transport.""" + client = make_async_client(console_logging=True) + mock_cls, _ = mock_openai() + with ( + patch(ASYNC_OPENAI_PATCH, mock_cls), + patch(ASYNC_TOKEN_PROVIDER_PATCH, return_value="tok"), + patch("azure.ai.projects.aio._patch.httpx") as mock_httpx, + patch("azure.ai.projects.aio._patch._OpenAILoggingTransport"), + ): + mock_httpx.AsyncClient.return_value = object() + client.get_openai_client() + mock_httpx.AsyncClient.assert_called_once()