Python: [BREAKING] updated structure and samples (#875)

* updated structure and samples

* updated names and removed cross tests

* updated projects etc

* updated tests

* updated test

* test fixes

* removed devui for now

* updated all-tests task

* removed old style configs

* remove coverage from tests

* updated to unit tests with all-tests

* updated foundry everywhere

* fix azure ai tests

* fix merge tests

* fix mypy
This commit is contained in:
Eduard van Valkenburg
2025-09-25 09:02:53 +02:00
committed by GitHub
Unverified
parent 366a7f7d47
commit 9355329dfd
169 changed files with 1159 additions and 1761 deletions
+10 -12
View File
@@ -5,7 +5,7 @@ Highlights
- Flexible Agent Framework: build, orchestrate, and deploy AI agents and multi-agent systems
- Multi-Agent Orchestration: Group chat, sequential, concurrent, and handoff patterns
- Plugin Ecosystem: Extend with native functions, OpenAPI, Model Context Protocol (MCP), and more
- LLM Support: OpenAI, Azure OpenAI, Azure AI Foundry, and more
- LLM Support: OpenAI, Azure OpenAI, Azure AI, and more
- Runtime Support: In-process and distributed agent execution
- Multimodal: Text, vision, and function calling
- Cross-Platform: .NET and Python implementations
@@ -13,13 +13,11 @@ Highlights
## Quick Install
```bash
pip install agent-framework
# Optional: Add Azure integration
pip install agent-framework[azure]
# Optional: Add Foundry integration
pip install agent-framework[foundry]
pip install agent-framework[all]
# Optional: Add Azure AI integration
pip install agent-framework-azure-ai
# Optional: Both
pip install agent-framework[azure,foundry]
pip install agent-framework-azure-ai agent-framework-copilotstudio
```
Supported Platforms:
@@ -40,16 +38,16 @@ AZURE_OPENAI_API_KEY=...
AZURE_OPENAI_ENDPOINT=...
AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=...
...
FOUNDRY_PROJECT_ENDPOINT=...
FOUNDRY_MODEL_DEPLOYMENT_NAME=...
AZURE_AI_PROJECT_ENDPOINT=...
AZURE_AI_MODEL_DEPLOYMENT_NAME=...
```
You can also override environment variables by explicitly passing configuration parameters to the chat client constructor:
```python
from agent_framework.azure import AzureChatClient
from agent_framework.azure import AzureOpenAIChatClient
chat_client = AzureChatClient(
chat_client = AzureOpenAIChatClient(
api_key="",
endpoint="",
deployment_name="",
@@ -223,7 +221,7 @@ if __name__ == "__main__":
- [Getting Started with Agents](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/agents): Basic agent creation and tool usage
- [Chat Client Examples](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/chat_client): Direct chat client usage patterns
- [Azure Integration](https://github.com/microsoft/agent-framework/tree/main/python/packages/azure): Azure OpenAI and AI Foundry integration
- [Azure AI Integration](https://github.com/microsoft/agent-framework/tree/main/python/packages/azure-ai): Azure AI integration
- [.NET Orchestration Samples](https://github.com/microsoft/agent-framework/tree/main/dotnet/samples/GettingStarted/Orchestration): Advanced multi-agent patterns (.NET)
## Agent Framework Documentation
@@ -1,31 +1,33 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
from typing import Any
PACKAGE_NAME = "agent_framework_azure"
PACKAGE_EXTRA = "azure"
_IMPORTS = [
"AzureAssistantsClient",
"AzureChatClient",
"AzureOpenAISettings",
"AzureResponsesClient",
"get_entra_auth_token",
"__version__",
]
_IMPORTS: dict[str, tuple[str, list[str]]] = {
"AzureAIAgentClient": ("agent_framework_azure_ai", ["azure_ai", "azure"]),
"AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", []),
"AzureOpenAIChatClient": ("agent_framework.azure._chat_client", []),
"AzureAISettings": ("agent_framework_azure_ai", ["azure_ai", "azure"]),
"AzureOpenAISettings": ("agent_framework.azure._shared", []),
"AzureOpenAIResponsesClient": ("agent_framework.azure._responses_client", []),
"get_entra_auth_token": ("agent_framework.azure._entra_id_authentication", []),
}
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
package_name, package_extra = _IMPORTS[name]
try:
return getattr(importlib.import_module(PACKAGE_NAME), name)
return getattr(importlib.import_module(package_name), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The '{PACKAGE_EXTRA}' extra is not installed, "
f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
f"The {' or '.join(package_extra)} extra is not installed, "
f"please use `pip install agent-framework[{package_extra[0]}]`, "
"or update your requirements.txt or pyproject.toml file."
) from exc
raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
raise AttributeError(f"Module `azure` has no attribute {name}.")
def __dir__() -> list[str]:
return _IMPORTS
return list(_IMPORTS.keys())
@@ -1,19 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework_azure import (
AzureAssistantsClient,
AzureChatClient,
AzureOpenAISettings,
AzureResponsesClient,
__version__,
get_entra_auth_token,
)
__all__ = [
"AzureAssistantsClient",
"AzureChatClient",
"AzureOpenAISettings",
"AzureResponsesClient",
"__version__",
"get_entra_auth_token",
]
@@ -0,0 +1,135 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, ClassVar
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
from ..exceptions import ServiceInitializationError
from ..openai import OpenAIAssistantsClient
from ._shared import AzureOpenAISettings
if TYPE_CHECKING:
from azure.core.credentials import TokenCredential
__all__ = ["AzureOpenAIAssistantsClient"]
class AzureOpenAIAssistantsClient(OpenAIAssistantsClient):
"""Azure OpenAI Assistants client."""
DEFAULT_AZURE_API_VERSION: ClassVar[str] = "2024-05-01-preview"
def __init__(
self,
deployment_name: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
thread_id: str | None = None,
api_key: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
credential: "TokenCredential | None" = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an Azure OpenAI Assistants client.
Args:
deployment_name: The Azure OpenAI deployment name for the model to use.
assistant_id: The ID of an Azure OpenAI assistant to use.
If not provided, a new assistant will be created (and deleted after the request).
assistant_name: The name to use when creating new assistants.
thread_id: Default thread ID to use for conversations. Can be overridden by
conversation_id property, when making a request.
If not provided, a new thread will be created (and deleted after the request).
api_key: The optional API key to use. If provided will override,
the env vars or .env file value.
endpoint: The optional deployment endpoint. If provided will override the value
in the env vars or .env file.
base_url: The optional deployment base_url. If provided will override the value
in the env vars or .env file.
api_version: The optional deployment api version. If provided will override the value
in the env vars or .env file.
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
credential: The Azure credential to use for authentication. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
env_file_path: Use the environment settings file as a fallback
to environment variables. (Optional)
env_file_encoding: The encoding of the environment settings file. (Optional)
"""
try:
azure_openai_settings = AzureOpenAISettings(
api_key=SecretStr(api_key) if api_key else None,
base_url=AnyUrl(base_url) if base_url else None,
endpoint=AnyUrl(endpoint) if endpoint else None,
chat_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
default_api_version=self.DEFAULT_AZURE_API_VERSION,
)
except ValidationError as ex:
raise ServiceInitializationError("Failed to create Azure OpenAI settings.", ex) from ex
if not azure_openai_settings.chat_deployment_name:
raise ServiceInitializationError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
# Handle authentication: try API key first, then AD token, then Entra ID
if (
not async_client
and not azure_openai_settings.api_key
and not ad_token
and not ad_token_provider
and azure_openai_settings.token_endpoint
and credential
):
ad_token = azure_openai_settings.get_azure_auth_token(credential)
if not async_client and not azure_openai_settings.api_key and not ad_token and not ad_token_provider:
raise ServiceInitializationError("The Azure OpenAI API key, ad_token, or ad_token_provider is required.")
# Create Azure client if not provided
if not async_client:
client_params: dict[str, Any] = {
"api_version": azure_openai_settings.api_version,
"default_headers": default_headers,
}
if azure_openai_settings.api_key:
client_params["api_key"] = azure_openai_settings.api_key.get_secret_value()
elif ad_token:
client_params["azure_ad_token"] = ad_token
elif ad_token_provider:
client_params["azure_ad_token_provider"] = ad_token_provider
if azure_openai_settings.base_url:
client_params["base_url"] = str(azure_openai_settings.base_url)
elif azure_openai_settings.endpoint:
client_params["azure_endpoint"] = str(azure_openai_settings.endpoint)
async_client = AsyncAzureOpenAI(**client_params)
super().__init__(
ai_model_id=azure_openai_settings.chat_deployment_name,
assistant_id=assistant_id,
assistant_name=assistant_name,
thread_id=thread_id,
async_client=async_client, # type: ignore[reportArgumentType]
)
@@ -0,0 +1,191 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
import sys
from collections.abc import Mapping
from typing import Any, TypeVar
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
from .._tools import use_function_invocation
from .._types import (
ChatResponse,
ChatResponseUpdate,
CitationAnnotation,
TextContent,
)
from ..exceptions import ServiceInitializationError
from ..observability import use_observability
from ..openai._chat_client import OpenAIBaseChatClient
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
)
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
logger: logging.Logger = logging.getLogger(__name__)
TChatResponse = TypeVar("TChatResponse", ChatResponse, ChatResponseUpdate)
TAzureOpenAIChatClient = TypeVar("TAzureOpenAIChatClient", bound="AzureOpenAIChatClient")
@use_function_invocation
@use_observability
class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient):
"""Azure OpenAI Chat completion class."""
def __init__(
self,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
credential: TokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
) -> None:
"""Initialize an AzureChatCompletion service.
Args:
api_key: The optional api key. If provided, will override the value in the
env vars or .env file.
deployment_name: The optional deployment. If provided, will override the value
(chat_deployment_name) in the env vars or .env file.
endpoint: The optional deployment endpoint. If provided will override the value
in the env vars or .env file.
base_url: The optional deployment base_url. If provided will override the value
in the env vars or .env file.
api_version: The optional deployment api version. If provided will override the value
in the env vars or .env file.
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
credential: The Azure credential for authentication. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
env_file_path: Use the environment settings file as a fallback to using env vars.
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`. (Optional)
"""
try:
# Filter out any None values from the arguments
azure_openai_settings = AzureOpenAISettings(
api_key=SecretStr(api_key) if api_key else None,
base_url=AnyUrl(base_url) if base_url else None,
endpoint=AnyUrl(endpoint) if endpoint else None,
chat_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
except ValidationError as exc:
raise ServiceInitializationError(f"Failed to validate settings: {exc}") from exc
if not azure_openai_settings.chat_deployment_name:
raise ServiceInitializationError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
super().__init__(
deployment_name=azure_openai_settings.chat_deployment_name,
endpoint=azure_openai_settings.endpoint,
base_url=azure_openai_settings.base_url,
api_version=azure_openai_settings.api_version, # type: ignore
api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None,
ad_token=ad_token,
ad_token_provider=ad_token_provider,
token_endpoint=azure_openai_settings.token_endpoint,
credential=credential,
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls: type[TAzureOpenAIChatClient], settings: dict[str, Any]) -> TAzureOpenAIChatClient:
"""Initialize an Azure OpenAI service from a dictionary of settings.
Args:
settings: A dictionary of settings for the service.
should contain keys: service_id, and optionally:
ad_auth, ad_token_provider, default_headers
"""
return cls(
api_key=settings.get("api_key"),
deployment_name=settings.get("deployment_name"),
endpoint=settings.get("endpoint"),
base_url=settings.get("base_url"),
api_version=settings.get("api_version"),
ad_token=settings.get("ad_token"),
ad_token_provider=settings.get("ad_token_provider"),
default_headers=settings.get("default_headers"),
env_file_path=settings.get("env_file_path"),
)
@override
def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
"""Parse the choice into a TextContent object.
Overwritten from OpenAIBaseChatClient to deal with Azure On Your Data function.
For docs see:
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
"""
message = choice.message if isinstance(choice, Choice) else choice.delta
if hasattr(message, "refusal") and message.refusal:
return TextContent(text=message.refusal, raw_representation=choice)
if not message.content:
return None
text_content = TextContent(text=message.content, raw_representation=choice)
if not message.model_extra or "context" not in message.model_extra:
return text_content
context: dict[str, Any] | str = message.context # type: ignore[assignment, union-attr]
if isinstance(context, str):
try:
context = json.loads(context)
except json.JSONDecodeError:
logger.warning("Context is not a valid JSON string, ignoring context.")
return text_content
if not isinstance(context, dict):
logger.warning("Context is not a valid dictionary, ignoring context.")
return text_content
# `all_retrieved_documents` is currently not used, but can be retrieved
# through the raw_representation in the text content.
if intent := context.get("intent"):
text_content.additional_properties = {"intent": intent}
if citations := context.get("citations"):
text_content.annotations = []
for citation in citations:
text_content.annotations.append(
CitationAnnotation(
title=citation.get("title", ""),
url=citation.get("url", ""),
snippet=citation.get("content", ""),
file_id=citation.get("filepath", ""),
tool_name="Azure-on-your-Data",
additional_properties={"chunk_id": citation.get("chunk_id", "")},
raw_representation=citation,
)
)
return text_content
@@ -0,0 +1,76 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from typing import TYPE_CHECKING, Any
from azure.core.exceptions import ClientAuthenticationError
from ..exceptions import ServiceInvalidAuthError
if TYPE_CHECKING:
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
logger: logging.Logger = logging.getLogger(__name__)
def get_entra_auth_token(
credential: "TokenCredential",
token_endpoint: str,
**kwargs: Any,
) -> str | None:
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint.
The token endpoint may be specified as an environment variable, via the .env
file or as an argument. If the token endpoint is not provided, the default is None.
Args:
credential: The Azure credential to use for authentication.
token_endpoint: The token endpoint to use to retrieve the authentication token.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
"""
if not token_endpoint:
raise ServiceInvalidAuthError(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
try:
auth_token = credential.get_token(token_endpoint, **kwargs)
except ClientAuthenticationError as ex:
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
return None
return auth_token.token if auth_token else None
async def get_entra_auth_token_async(
credential: "AsyncTokenCredential", token_endpoint: str, **kwargs: Any
) -> str | None:
"""Retrieve a async Microsoft Entra Auth Token for a given token endpoint.
The token endpoint may be specified as an environment variable, via the .env
file or as an argument. If the token endpoint is not provided, the default is None.
Args:
credential: The async Azure credential to use for authentication.
token_endpoint: The token endpoint to use to retrieve the authentication token.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
"""
if not token_endpoint:
raise ServiceInvalidAuthError(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
try:
auth_token = await credential.get_token(token_endpoint, **kwargs)
except ClientAuthenticationError as ex:
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
return None
return auth_token.token if auth_token else None
@@ -0,0 +1,134 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Mapping
from typing import Any, TypeVar
from urllib.parse import urljoin
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
from .._tools import use_function_invocation
from ..exceptions import ServiceInitializationError
from ..observability import use_observability
from ..openai._responses_client import OpenAIBaseResponsesClient
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
)
TAzureOpenAIResponsesClient = TypeVar("TAzureOpenAIResponsesClient", bound="AzureOpenAIResponsesClient")
@use_observability
@use_function_invocation
class AzureOpenAIResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClient):
"""Azure Responses completion class."""
def __init__(
self,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
credential: TokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
) -> None:
"""Initialize an AzureResponses service.
Args:
api_key: The optional api key. If provided, will override the value in the
env vars or .env file.
deployment_name: The optional deployment. If provided, will override the value
(responses_deployment_name) in the env vars or .env file.
endpoint: The optional deployment endpoint. If provided will override the value
in the env vars or .env file.
base_url: The optional deployment base_url. If provided will override the value
in the env vars or .env file. Currently, the base_url must end with "/openai/v1/"
api_version: The optional deployment api version. If provided will override the value
in the env vars or .env file. Currently, the api_version must be "preview".
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
credential: The Azure credential for authentication. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
env_file_path: Use the environment settings file as a fallback to using env vars.
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`. (Optional)
"""
try:
# Filter out any None values from the arguments
azure_openai_settings = AzureOpenAISettings(
api_key=SecretStr(api_key) if api_key else None,
base_url=AnyUrl(base_url) if base_url else None,
endpoint=AnyUrl(endpoint) if endpoint else None,
responses_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
default_api_version="preview",
)
# TODO(peterychang): This is a temporary hack to ensure that the base_url is set correctly
# while this feature is in preview.
# But we should only do this if we're on azure. Private deployments may not need this.
if (
not azure_openai_settings.base_url
and azure_openai_settings.endpoint
and str(azure_openai_settings.endpoint).rstrip("/").endswith("openai.azure.com")
):
azure_openai_settings.base_url = AnyUrl(urljoin(str(azure_openai_settings.endpoint), "/openai/v1/"))
except ValidationError as exc:
raise ServiceInitializationError(f"Failed to validate settings: {exc}") from exc
if not azure_openai_settings.responses_deployment_name:
raise ServiceInitializationError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
)
super().__init__(
deployment_name=azure_openai_settings.responses_deployment_name,
endpoint=azure_openai_settings.endpoint,
base_url=azure_openai_settings.base_url,
api_version=azure_openai_settings.api_version, # type: ignore
api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None,
ad_token=ad_token,
ad_token_provider=ad_token_provider,
token_endpoint=azure_openai_settings.token_endpoint,
credential=credential,
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls: type[TAzureOpenAIResponsesClient], settings: dict[str, Any]) -> TAzureOpenAIResponsesClient:
"""Initialize an Open AI service from a dictionary of settings.
Args:
settings: A dictionary of settings for the service.
"""
return cls(
api_key=settings.get("api_key"),
deployment_name=settings.get("deployment_name"),
endpoint=settings.get("endpoint"),
base_url=settings.get("base_url"),
api_version=settings.get("api_version"),
ad_token=settings.get("ad_token"),
ad_token_provider=settings.get("ad_token_provider"),
default_headers=settings.get("default_headers"),
env_file_path=settings.get("env_file_path"),
)
@@ -0,0 +1,242 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from collections.abc import Awaitable, Callable, Mapping
from copy import copy
from typing import Any, ClassVar, Final
from azure.core.credentials import TokenCredential
from openai.lib.azure import AsyncAzureOpenAI
from pydantic import ConfigDict, SecretStr, model_validator, validate_call
from .._pydantic import AFBaseSettings, HTTPsUrl
from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from ..exceptions import ServiceInitializationError
from ..openai._shared import OpenAIBase
from ._entra_id_authentication import get_entra_auth_token
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
logger: logging.Logger = logging.getLogger(__name__)
DEFAULT_AZURE_API_VERSION: Final[str] = "2024-10-21"
DEFAULT_AZURE_TOKEN_ENDPOINT: Final[str] = "https://cognitiveservices.azure.com/.default" # noqa: S105
class AzureOpenAISettings(AFBaseSettings):
"""AzureOpenAI model settings.
The settings are first loaded from environment variables with the prefix 'AZURE_OPENAI_'.
If the environment variables are not found, the settings can be loaded from a .env file
with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
are ignored; however, validation will fail alerting that the settings are missing.
Args:
endpoint: The endpoint of the Azure deployment. This value
can be found in the Keys & Endpoint section when examining
your resource from the Azure portal, the endpoint should end in openai.azure.com.
If both base_url and endpoint are supplied, base_url will be used.
(Env var AZURE_OPENAI_ENDPOINT)
chat_deployment_name: The name of the Azure Chat deployment. This value
will correspond to the custom name you chose for your deployment
when you deployed a model. This value can be found under
Resource Management > Deployments in the Azure portal or, alternatively,
under Management > Deployments in Azure AI Foundry.
(Env var AZURE_OPENAI_CHAT_DEPLOYMENT_NAME)
responses_deployment_name: The name of the Azure Responses deployment. This value
will correspond to the custom name you chose for your deployment
when you deployed a model. This value can be found under
Resource Management > Deployments in the Azure portal or, alternatively,
under Management > Deployments in Azure AI Foundry.
(Env var AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME)
api_key: The API key for the Azure deployment. This value can be
found in the Keys & Endpoint section when examining your resource in
the Azure portal. You can use either KEY1 or KEY2.
(Env var AZURE_OPENAI_API_KEY)
api_version: The API version to use. The default value is `default_api_version`.
(Env var AZURE_OPENAI_API_VERSION)
base_url: The url of the Azure deployment. This value
can be found in the Keys & Endpoint section when examining
your resource from the Azure portal, the base_url consists of the endpoint,
followed by /openai/deployments/{deployment_name}/,
use endpoint if you only want to supply the endpoint.
(Env var AZURE_OPENAI_BASE_URL)
token_endpoint: The token endpoint to use to retrieve the authentication token.
The default value is `default_token_endpoint`.
(Env var AZURE_OPENAI_TOKEN_ENDPOINT)
default_api_version: The default API version to use if not specified.
The default value is "2024-10-21".
default_token_endpoint: The default token endpoint to use if not specified.
The default value is "https://cognitiveservices.azure.com/.default".
env_file_path: The path to the .env file to load settings from.
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
"""
env_prefix: ClassVar[str] = "AZURE_OPENAI_"
chat_deployment_name: str | None = None
responses_deployment_name: str | None = None
endpoint: HTTPsUrl | None = None
base_url: HTTPsUrl | None = None
api_key: SecretStr | None = None
api_version: str | None = None
token_endpoint: str | None = None
default_api_version: str = DEFAULT_AZURE_API_VERSION
default_token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT
def get_azure_auth_token(
self, credential: "TokenCredential", token_endpoint: str | None = None, **kwargs: Any
) -> str | None:
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint for the use with Azure OpenAI.
The required role for the token is `Cognitive Services OpenAI Contributor`.
The token endpoint may be specified as an environment variable, via the .env
file or as an argument. If the token endpoint is not provided, the default is None.
The `token_endpoint` argument takes precedence over the `token_endpoint` attribute.
Args:
credential: The Azure AD credential to use.
token_endpoint: The token endpoint to use. Defaults to `https://cognitiveservices.azure.com/.default`.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
Raises:
ServiceInitializationError: If the token endpoint is not provided.
"""
endpoint_to_use = token_endpoint or self.token_endpoint or self.default_token_endpoint
return get_entra_auth_token(credential, endpoint_to_use, **kwargs)
@model_validator(mode="after")
def _validate_fields(self) -> Self:
self.api_version = self.api_version or self.default_api_version
self.token_endpoint = self.token_endpoint or self.default_token_endpoint
return self
class AzureOpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an Azure OpenAI service."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure_openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def __init__(
self,
deployment_name: str,
endpoint: HTTPsUrl | None = None,
base_url: HTTPsUrl | None = None,
api_version: str = DEFAULT_AZURE_API_VERSION,
api_key: str | None = None,
ad_token: str | None = None,
ad_token_provider: Callable[[], str | Awaitable[str]] | None = None,
token_endpoint: str | None = None,
credential: TokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncAzureOpenAI | None = None,
instruction_role: str | None = None,
**kwargs: Any,
) -> None:
"""Internal class for configuring a connection to an Azure OpenAI service.
The `validate_call` decorator is used with a configuration that allows arbitrary types.
This is necessary for types like `HTTPsUrl` and `OpenAIModelTypes`.
Args:
deployment_name: Name of the deployment.
ai_model_type: The type of OpenAI model to deploy.
endpoint: The specific endpoint URL for the deployment.
base_url: The base URL for Azure services.
api_version: Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
api_key: API key for Azure services.
ad_token: Azure AD token for authentication.
ad_token_provider: A callable or coroutine function providing Azure AD tokens.
token_endpoint: Azure AD token endpoint use to get the token.
credential: Azure credential for authentication.
default_headers: Default headers for HTTP requests.
client: An existing client to use.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
kwargs: Additional keyword arguments.
"""
# Merge APP_INFO into the headers if it exists
merged_headers = dict(copy(default_headers)) if default_headers else {}
if APP_INFO:
merged_headers.update(APP_INFO)
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
if not client:
# If the client is None, the api_key is none, the ad_token is none, and the ad_token_provider is none,
# then we will attempt to get the ad_token using the default endpoint specified in the Azure OpenAI
# settings.
if not api_key and not ad_token_provider and not ad_token and token_endpoint and credential:
ad_token = get_entra_auth_token(credential, token_endpoint)
if not api_key and not ad_token and not ad_token_provider:
raise ServiceInitializationError(
"Please provide either api_key, ad_token or ad_token_provider or a client."
)
if not endpoint and not base_url:
raise ServiceInitializationError("Please provide an endpoint or a base_url")
args: dict[str, Any] = {
"default_headers": merged_headers,
}
if api_version:
args["api_version"] = api_version
if ad_token:
args["azure_ad_token"] = ad_token
if ad_token_provider:
args["azure_ad_token_provider"] = ad_token_provider
if api_key:
args["api_key"] = api_key
if base_url:
args["base_url"] = str(base_url)
if endpoint and not base_url:
args["azure_endpoint"] = str(endpoint)
if deployment_name:
args["azure_deployment"] = deployment_name
if "websocket_base_url" in kwargs:
args["websocket_base_url"] = kwargs.pop("websocket_base_url")
client = AsyncAzureOpenAI(**args)
args = {
"ai_model_id": deployment_name,
"client": client,
}
if instruction_role:
args["instruction_role"] = instruction_role
super().__init__(**args, **kwargs)
def to_dict(self) -> dict[str, Any]:
"""Convert the configuration to a dictionary."""
client_settings = {
"base_url": str(self.client.base_url),
"api_version": self.client._custom_query["api-version"], # type: ignore
"api_key": self.client.api_key,
"ad_token": getattr(self.client, "_azure_ad_token", None),
"ad_token_provider": getattr(self.client, "_azure_ad_token_provider", None),
"default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY},
}
base = self.model_dump(
exclude={
"prompt_tokens",
"completion_tokens",
"total_tokens",
"api_type",
"org_id",
"service_id",
"client",
},
by_alias=True,
exclude_none=True,
)
base.update(client_settings)
return base
@@ -1,24 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
from typing import Any
PACKAGE_NAME = "agent_framework_copilotstudio"
PACKAGE_EXTRA = "copilotstudio"
_IMPORTS = ["CopilotStudioAgent", "__version__", "acquire_token"]
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
try:
return getattr(importlib.import_module(PACKAGE_NAME), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The '{PACKAGE_EXTRA}' extra is not installed, "
f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
) from exc
raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
def __dir__() -> list[str]:
return _IMPORTS
@@ -1,24 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
from typing import Any
PACKAGE_NAME = "agent_framework_foundry"
PACKAGE_EXTRA = "foundry"
_IMPORTS = ["__version__", "FoundryChatClient", "FoundrySettings"]
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
try:
return getattr(importlib.import_module(PACKAGE_NAME), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The '{PACKAGE_EXTRA}' extra is not installed, "
f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
) from exc
raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
def __dir__() -> list[str]:
return _IMPORTS
@@ -1,5 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework_foundry import FoundryChatClient, FoundrySettings, __version__
__all__ = ["FoundryChatClient", "FoundrySettings", "__version__"]
@@ -0,0 +1,30 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
from typing import Any
PACKAGE_NAME = "agent_framework_copilotstudio"
PACKAGE_EXTRA = ["microsoft-copilotstudio", "copilotstudio"]
_IMPORTS: dict[str, tuple[str, list[str]]] = {
"CopilotStudioAgent": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
"__version__": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
"acquire_token": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
}
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
package_name, package_extra = _IMPORTS[name]
try:
return getattr(importlib.import_module(package_name), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The {' or '.join(package_extra)} extra is not installed, "
f"please use `pip install agent-framework[{package_extra[0]}]`, "
"or update your requirements.txt or pyproject.toml file."
) from exc
raise AttributeError(f"Module `azure` has no attribute {name}.")
def __dir__() -> list[str]:
return list(_IMPORTS.keys())
@@ -82,7 +82,7 @@ class OpenAISettings(AFBaseSettings):
encoding 'utf-8'. If the settings are not found in the .env file, the settings are ignored;
however, validation will fail alerting that the settings are missing.
Attributes:
Args:
api_key: OpenAI API key, see https://platform.openai.com/account/api-keys
(Env var OPENAI_API_KEY)
base_url: The base URL for the OpenAI API.
@@ -93,21 +93,6 @@ class OpenAISettings(AFBaseSettings):
(Env var OPENAI_CHAT_MODEL_ID)
responses_model_id: The OpenAI responses model ID to use, for example, gpt-4o or o1.
(Env var OPENAI_RESPONSES_MODEL_ID)
text_model_id: The OpenAI text model ID to use, for example, gpt-3.5-turbo-instruct.
(Env var OPENAI_TEXT_MODEL_ID)
embedding_model_id: The OpenAI embedding model ID to use, for example, text-embedding-ada-002.
(Env var OPENAI_EMBEDDING_MODEL_ID)
text_to_image_model_id: The OpenAI text to image model ID to use, for example, dall-e-3.
(Env var OPENAI_TEXT_TO_IMAGE_MODEL_ID)
audio_to_text_model_id: The OpenAI audio to text model ID to use, for example, whisper-1.
(Env var OPENAI_AUDIO_TO_TEXT_MODEL_ID)
text_to_audio_model_id: The OpenAI text to audio model ID to use, for example, jukebox-1.
(Env var OPENAI_TEXT_TO_AUDIO_MODEL_ID)
realtime_model_id: The OpenAI realtime model ID to use,
for example, gpt-4o-realtime-preview-2024-12-17.
(Env var OPENAI_REALTIME_MODEL_ID)
Parameters:
env_file_path: The path to the .env file to load settings from.
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
"""
@@ -119,12 +104,6 @@ class OpenAISettings(AFBaseSettings):
org_id: str | None = None
chat_model_id: str | None = None
responses_model_id: str | None = None
text_model_id: str | None = None
embedding_model_id: str | None = None
text_to_image_model_id: str | None = None
audio_to_text_model_id: str | None = None
text_to_audio_model_id: str | None = None
realtime_model_id: str | None = None
class OpenAIBase(AFBaseModel):
+26 -19
View File
@@ -1,10 +1,10 @@
[project]
name = "agent-framework"
description = "Microsoft Agent Framework for building AI Agents with Python."
authors = [{ name = "Microsoft", email = "SK-Support@microsoft.com"}]
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "0.1.0b1"
version = "0.1.0b1" # TODO: decide on initial version and versioning strategy
license-files = ["LICENSE"]
urls.homepage = "https://learn.microsoft.com/en-us/semantic-kernel/overview/"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -24,31 +24,38 @@ classifiers = [
]
dependencies = [
"openai>=1.99.0",
"pydantic>=2.11.7",
"pydantic-settings>=2.10.1",
"typing-extensions>=4.14.0",
"opentelemetry-api ~= 1.24",
"opentelemetry-sdk ~= 1.24",
"pydantic>=2,<3",
"pydantic-settings>=2,<3",
"typing-extensions",
"opentelemetry-api>=1.24",
"opentelemetry-sdk>=1.24",
"mcp[ws]>=1.13",
"azure-monitor-opentelemetry>=1.7.0",
"azure-monitor-opentelemetry-exporter>=1.0.0b41",
"opentelemetry-exporter-otlp-proto-grpc>=1.36.0",
"opentelemetry-semantic-conventions-ai>=0.4.13",
"aiofiles>=24.1.0"
"aiofiles>=24.1.0",
"azure-identity>=1,<2"
]
[project.optional-dependencies]
azure = [
"agent-framework-azure"
azure-ai = [
"agent-framework-azure-ai"
]
foundry = [
"agent-framework-foundry"
azure = [
"agent-framework-azure-ai"
]
microsoft-copilotstudio = [
"agent-framework-copilotstudio"
]
microsoft = [
"agent-framework-copilotstudio"
]
redis = [
"agent-framework-redis"
]
viz = [
"graphviz>=0.20.0",
"graphviz>=0.20.0"
]
runtime = [
"agent-framework-runtime"
@@ -60,12 +67,13 @@ devui = [
"agent-framework-devui"
]
all = [
"agent-framework-azure",
"agent-framework-foundry",
"agent_framework_copilotstudio",
"agent-framework-azure-ai",
"agent-framework-runtime",
"agent-framework-mem0",
"agent-framework-redis",
"agent-framework-devui"
"agent-framework-devui",
"graphviz>=0.20.0"
]
[tool.uv]
@@ -81,10 +89,9 @@ fallback-version = "0.0.0"
[tool.pytest.ini_options]
testpaths = [
'tests',
'tests',
'packages/main/tests',
'packages/azure/tests',
'packages/foundry/tests',
'packages/azure-ai/tests',
'packages/copilotstudio/tests',
'packages/mem0/tests',
'packages/runtime/tests'
@@ -0,0 +1,62 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from pytest import fixture
from agent_framework import ChatMessage
# region: Connector Settings fixtures
@fixture
def exclude_list(request: Any) -> list[str]:
"""Fixture that returns a list of environment variables to exclude."""
return request.param if hasattr(request, "param") else []
@fixture
def override_env_param_dict(request: Any) -> dict[str, str]:
"""Fixture that returns a dict of environment variables to override."""
return request.param if hasattr(request, "param") else {}
# These two fixtures are used for multiple things, also non-connector tests
@fixture()
def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
"""Fixture to set environment variables for AzureOpenAISettings."""
if exclude_list is None:
exclude_list = []
if override_env_param_dict is None:
override_env_param_dict = {}
env_vars = {
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.com",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_TEXT_DEPLOYMENT_NAME": "test_text_deployment",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
"AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME": "test_text_to_image_deployment",
"AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME": "test_audio_to_text_deployment",
"AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME": "test_text_to_audio_deployment",
"AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME": "test_realtime_deployment",
"AZURE_OPENAI_API_KEY": "test_api_key",
"AZURE_OPENAI_API_VERSION": "2023-03-15-preview",
"AZURE_OPENAI_BASE_URL": "https://test_text_deployment.test-base-url.com",
"AZURE_OPENAI_TOKEN_ENDPOINT": "https://test-token-endpoint.com",
}
env_vars.update(override_env_param_dict) # type: ignore
for key, value in env_vars.items():
if key in exclude_list:
monkeypatch.delenv(key, raising=False) # type: ignore
continue
monkeypatch.setenv(key, value) # type: ignore
return env_vars
@fixture(scope="function")
def chat_history() -> list[ChatMessage]:
return []
@@ -0,0 +1,723 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from azure.identity import AzureCliCredential
from pydantic import Field
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
TextContent,
)
from agent_framework.azure import AzureOpenAIAssistantsClient
from agent_framework.exceptions import ServiceInitializationError
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
def create_test_azure_assistants_client(
mock_async_azure_openai: MagicMock,
deployment_name: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
thread_id: str | None = None,
should_delete_assistant: bool = False,
) -> AzureOpenAIAssistantsClient:
"""Helper function to create AzureOpenAIAssistantsClient instances for testing, bypassing Pydantic validation."""
return AzureOpenAIAssistantsClient.model_construct(
ai_model_id=deployment_name or "test_chat_deployment",
assistant_id=assistant_id,
assistant_name=assistant_name,
thread_id=thread_id,
api_key="test-api-key",
endpoint="https://test-endpoint.com",
client=mock_async_azure_openai,
_should_delete_assistant=should_delete_assistant,
)
@pytest.fixture
def mock_async_azure_openai() -> MagicMock:
"""Mock AsyncAzureOpenAI client."""
mock_client = MagicMock()
# Mock beta.assistants
mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
mock_client.beta.assistants.delete = AsyncMock()
# Mock beta.threads
mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
mock_client.beta.threads.delete = AsyncMock()
# Mock beta.threads.runs
mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
mock_client.beta.threads.runs.retrieve = AsyncMock()
mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
# Mock beta.threads.messages
mock_client.beta.threads.messages.create = AsyncMock()
mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
return mock_client
def test_azure_assistants_client_init_with_client(mock_async_azure_openai: MagicMock) -> None:
"""Test AzureOpenAIAssistantsClient initialization with existing client."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai,
deployment_name="test_chat_deployment",
assistant_id="existing-assistant-id",
thread_id="test-thread-id",
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.ai_model_id == "test_chat_deployment"
assert chat_client.assistant_id == "existing-assistant-id"
assert chat_client.thread_id == "test-thread-id"
assert not chat_client._should_delete_assistant # type: ignore
assert isinstance(chat_client, ChatClientProtocol)
def test_azure_assistants_client_init_auto_create_client(
azure_openai_unit_test_env: dict[str, str],
mock_async_azure_openai: MagicMock,
) -> None:
"""Test AzureOpenAIAssistantsClient initialization with auto-created client."""
chat_client = AzureOpenAIAssistantsClient.model_construct(
ai_model_id=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
assistant_id=None,
assistant_name="TestAssistant",
thread_id=None,
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
client=mock_async_azure_openai,
_should_delete_assistant=False,
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert chat_client.assistant_id is None
assert chat_client.assistant_name == "TestAssistant"
assert not chat_client._should_delete_assistant # type: ignore
def test_azure_assistants_client_init_validation_fail() -> None:
"""Test AzureOpenAIAssistantsClient initialization with validation failure."""
with pytest.raises(ServiceInitializationError):
# Force failure by providing invalid deployment name type - this should cause validation to fail
AzureOpenAIAssistantsClient(deployment_name=123, api_key="valid-key") # type: ignore
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
def test_azure_assistants_client_init_missing_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with missing deployment name."""
with pytest.raises(ServiceInitializationError):
AzureOpenAIAssistantsClient(
api_key=azure_openai_unit_test_env.get("AZURE_OPENAI_API_KEY", "test-key"), env_file_path="nonexistent.env"
)
def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with default headers."""
default_headers = {"X-Unit-Test": "test-guid"}
chat_client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
assert chat_client.ai_model_id == "test_chat_deployment"
assert isinstance(chat_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in chat_client.client.default_headers
assert chat_client.client.default_headers[key] == value
async def test_azure_assistants_client_get_assistant_id_or_create_existing_assistant(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
chat_client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "existing-assistant-id"
assert not chat_client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_not_called()
async def test_azure_assistants_client_get_assistant_id_or_create_create_new(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when creating a new assistant."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
)
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "test-assistant-id"
assert chat_client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_called_once()
async def test_azure_assistants_client_aclose_should_not_delete(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test close when assistant should not be deleted."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
)
await chat_client.close() # type: ignore
# Verify assistant deletion was not called
mock_async_azure_openai.beta.assistants.delete.assert_not_called()
assert not chat_client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_aclose_should_delete(mock_async_azure_openai: MagicMock) -> None:
"""Test close method calls cleanup."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
await chat_client.close()
# Verify assistant deletion was called
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
assert not chat_client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_async_context_manager(mock_async_azure_openai: MagicMock) -> None:
"""Test async context manager functionality."""
chat_client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
# Test context manager
async with chat_client:
pass # Just test that we can enter and exit
# Verify cleanup was called on exit
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test serialization of AzureOpenAIAssistantsClient."""
default_headers = {"X-Unit-Test": "test-guid"}
# Test basic initialization and to_dict
chat_client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
assistant_id="test-assistant-id",
assistant_name="TestAssistant",
thread_id="test-thread-id",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
dumped_settings = chat_client.to_dict()
assert dumped_settings["ai_model_id"] == "test_chat_deployment"
assert dumped_settings["assistant_id"] == "test-assistant-id"
assert dumped_settings["assistant_name"] == "TestAssistant"
assert dumped_settings["thread_id"] == "test-thread-id"
assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
return f"The weather in {location} is sunny with a high of 25°C."
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response() -> None:
"""Test Azure Assistants Client response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response_tools() -> None:
"""Test Azure Assistants Client response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming() -> None:
"""Test Azure Assistants Client streaming response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_streaming_response(messages=messages)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming_tools() -> None:
"""Test Azure Assistants Client streaming response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_with_existing_assistant() -> None:
"""Test Azure Assistants Client with existing assistant ID."""
# First create an assistant to use in the test
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
# Now test using the existing assistant
async with AzureOpenAIAssistantsClient(
assistant_id=assistant_id, credential=AzureCliCredential()
) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert azure_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert len(response.text) > 0
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run():
"""Test ChatAgent basic run functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "Hello World" in response.text
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run_streaming():
"""Test ChatAgent basic streaming functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert chunk is not None
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_message += chunk.text
# Validate streaming response
assert len(full_message) > 0
assert "streaming response test" in full_message.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_thread_persistence():
"""Test ChatAgent thread persistence across runs with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", thread=thread
)
assert isinstance(first_response, AgentRunResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", thread=thread
)
assert isinstance(second_response, AgentRunResponse)
assert "42" in second_response.text
# Verify thread has been populated with conversation ID
assert thread.service_thread_id is not None
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_existing_thread_id():
"""Test ChatAgent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Start a conversation and get the thread ID
thread = agent.get_new_thread()
response1 = await agent.run("What's the weather in Paris?", thread=thread)
# Validate first response
assert isinstance(response1, AgentRunResponse)
assert response1.text is not None
assert any(word in response1.text.lower() for word in ["weather", "paris"])
# The thread ID is set after the first response
existing_thread_id = thread.service_thread_id
assert existing_thread_id is not None
# Now continue with the same thread ID in a new agent instance
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Create a thread with the existing ID
thread = AgentThread(service_thread_id=existing_thread_id)
# Ask about the previous conversation
response2 = await agent.run("What was the last city I asked about?", thread=thread)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentRunResponse)
assert response2.text is not None
# Should reference Paris from the previous conversation
assert "paris" in response2.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_code_interpreter():
"""Test ChatAgent with code interpreter through AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Request code execution
response = await agent.run("Write Python code to calculate the factorial of 5 and show the result.")
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
# Factorial of 5 is 120
assert "120" in response.text or "factorial" in response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Assistants Client."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
def test_azure_assistants_client_entra_id_authentication() -> None:
"""Test Entra ID authentication path with credential."""
mock_credential = MagicMock()
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key to trigger Entra ID path
mock_settings.token_endpoint = "https://login.microsoftonline.com/test"
mock_settings.get_azure_auth_token.return_value = "entra-token-12345"
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
api_key="placeholder-key",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_credential,
token_endpoint="https://login.microsoftonline.com/test",
)
# Verify Entra ID token was requested
mock_settings.get_azure_auth_token.assert_called_once_with(mock_credential)
# Verify client was created with the token
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token"] == "entra-token-12345"
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_no_authentication_error() -> None:
"""Test authentication validation error when no auth provided."""
with patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class:
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.token_endpoint = None # No token endpoint
mock_settings_class.return_value = mock_settings
# Test missing authentication raises error
with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"):
AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
# No authentication provided at all
)
def test_azure_assistants_client_ad_token_authentication() -> None:
"""Test ad_token authentication client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
ad_token="test-ad-token-12345",
)
# ad_token path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token"] == "test-ad-token-12345"
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_ad_token_provider_authentication() -> None:
"""Test ad_token_provider authentication client parameter path."""
from openai.lib.azure import AsyncAzureADTokenProvider
mock_token_provider = MagicMock(spec=AsyncAzureADTokenProvider)
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key = None # No API key
mock_settings.api_version = "2024-05-01-preview"
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.base_url = None
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
ad_token_provider=mock_token_provider,
)
# ad_token_provider path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_token_provider
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_base_url_configuration() -> None:
"""Test base_url client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key.get_secret_value.return_value = "test-api-key"
mock_settings.base_url = "https://custom-base-url.com"
mock_settings.endpoint = None # No endpoint, should use base_url
mock_settings.api_version = "2024-05-01-preview"
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment", api_key="test-api-key", base_url="https://custom-base-url.com"
)
# base_url path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["base_url"] == "https://custom-base-url.com"
assert "azure_endpoint" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_azure_endpoint_configuration() -> None:
"""Test azure_endpoint client parameter path."""
with (
patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_settings = MagicMock()
mock_settings.chat_deployment_name = "test-deployment"
mock_settings.api_key.get_secret_value.return_value = "test-api-key"
mock_settings.base_url = None # No base_url
mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
mock_settings.api_version = "2024-05-01-preview"
mock_settings_class.return_value = mock_settings
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
api_key="test-api-key",
endpoint="https://test-endpoint.openai.azure.com",
)
# azure_endpoint path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_endpoint"] == "https://test-endpoint.openai.azure.com"
assert "base_url" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
@@ -0,0 +1,835 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
import openai
import pytest
from azure.identity import AzureCliCredential
from httpx import Request, Response
from openai import AsyncAzureOpenAI, AsyncStream
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
BaseChatClient,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
TextContent,
ai_function,
)
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.openai import (
ContentFilterResultSeverity,
OpenAIContentFilterException,
)
# region Service Setup
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_chat_client = AzureOpenAIChatClient()
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
def test_init_client(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization with client
client = MagicMock(spec=AsyncAzureOpenAI)
azure_chat_client = AzureOpenAIChatClient(async_client=client)
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
# Custom header for testing
default_headers = {"X-Unit-Test": "test-guid"}
azure_chat_client = AzureOpenAIChatClient(
default_headers=default_headers,
)
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
for key, value in default_headers.items():
assert key in azure_chat_client.client.default_headers
assert azure_chat_client.client.default_headers[key] == value
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
azure_chat_client = AzureOpenAIChatClient()
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, BaseChatClient)
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
def test_init_with_empty_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient(
env_file_path="test.env",
)
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient(
env_file_path="test.env",
)
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
def test_init_with_invalid_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIChatClient()
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Test": "test"}
settings = {
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
"default_headers": default_headers,
"env_file_path": "test.env",
}
azure_chat_client = AzureOpenAIChatClient.from_dict(settings)
dumped_settings = azure_chat_client.to_dict()
assert dumped_settings["ai_model_id"] == settings["deployment_name"]
assert str(settings["endpoint"]) in str(dumped_settings["base_url"])
assert str(settings["deployment_name"]) in str(dumped_settings["base_url"])
assert settings["api_key"] == dumped_settings["api_key"]
assert settings["api_version"] == dumped_settings["api_version"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
assert USER_AGENT_KEY not in dumped_settings["default_headers"]
# endregion
# region CMC
@pytest.fixture
def mock_chat_completion_response() -> ChatCompletion:
return ChatCompletion(
id="test_id",
choices=[
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
],
created=0,
model="test",
object="chat.completion",
)
@pytest.fixture
def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
content = ChatCompletionChunk(
id="test_id",
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
created=0,
model="test",
object="chat.completion.chunk",
)
stream = MagicMock(spec=AsyncStream)
stream.__aiter__.return_value = [content]
return stream
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(
messages=chat_history,
)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
stream=False,
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_with_logit_bias(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
token_bias: dict[str | int, float] = {"1": -100}
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, logit_bias=token_bias)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
stream=False,
logit_bias=token_bias,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_with_stop(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
stop = ["!"]
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(messages=chat_history, stop=stop)
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
stream=False,
stop=stop,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context={ # type: ignore
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
},
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
chat_history.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data_string(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context=json.dumps({ # type: ignore
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
}),
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data_fail(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context="not a dictionary", # type: ignore
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
{
"type": "AzureCognitiveSearch",
"parameters": {
"indexName": "test_index",
"endpoint": "https://test-endpoint-search.com",
"key": "test_key",
},
}
]
}
azure_chat_client = AzureOpenAIChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
stream=False,
extra_body=expected_data_settings,
)
CONTENT_FILTERED_ERROR_MESSAGE = (
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy. Please "
"modify your prompt and retry. To learn more about our content filtering policies please read our "
"documentation: https://go.microsoft.com/fwlink/?linkid=2198766"
)
CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
"Error code: 400 - {'error': {'message': \"%s\", 'type': null, 'param': 'prompt', 'code': 'content_filter', "
"'status': 400, 'innererror': {'code': 'ResponsibleAIPolicyViolation', 'content_filter_result': {'hate': "
"{'filtered': True, 'severity': 'high'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': "
"{'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}}}"
) % CONTENT_FILTERED_ERROR_MESSAGE
@patch.object(AsyncChatCompletions, "create")
async def test_content_filtering_raises_correct_exception(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
response=Response(400, request=Request("POST", test_endpoint)),
body={
"message": CONTENT_FILTERED_ERROR_MESSAGE,
"type": None,
"param": "prompt",
"code": "content_filter",
"status": 400,
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {
"hate": {"filtered": True, "severity": "high"},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": False, "severity": "safe"},
},
},
},
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error") as exc_info:
await azure_chat_client.get_response(
messages=chat_history,
)
content_filter_exc = exc_info.value
assert content_filter_exc.param == "prompt"
assert content_filter_exc.content_filter_result["hate"].filtered
assert content_filter_exc.content_filter_result["hate"].severity == ContentFilterResultSeverity.HIGH
@patch.object(AsyncChatCompletions, "create")
async def test_content_filtering_without_response_code_raises_with_default_code(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
response=Response(400, request=Request("POST", test_endpoint)),
body={
"message": CONTENT_FILTERED_ERROR_MESSAGE,
"type": None,
"param": "prompt",
"code": "content_filter",
"status": 400,
"innererror": {
"content_filter_result": {
"hate": {"filtered": True, "severity": "high"},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": False, "severity": "safe"},
},
},
},
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error"):
await azure_chat_client.get_response(
messages=chat_history,
)
@patch.object(AsyncChatCompletions, "create")
async def test_bad_request_non_content_filter(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
mock_create.side_effect = openai.BadRequestError(
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
)
azure_chat_client = AzureOpenAIChatClient()
with pytest.raises(ServiceResponseException, match="service failed to complete the prompt"):
await azure_chat_client.get_response(
messages=chat_history,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_get_streaming(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
) -> None:
mock_create.return_value = mock_streaming_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
async for msg in azure_chat_client.get_streaming_response(
messages=chat_history,
):
assert msg is not None
assert msg.message_id is not None
assert msg.response_id is not None
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
stream=True,
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
# NOTE: The `stream_options={"include_usage": True}` is explicitly enforced in
# `OpenAIChatCompletionBase._inner_get_streaming_response`.
# To ensure consistency, we align the arguments here accordingly.
stream_options={"include_usage": True},
)
@ai_function
def get_story_text() -> str:
"""Returns a story about Emily and David."""
return (
"Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change."
)
@ai_function
def get_weather(location: str) -> str:
"""Get the current weather for a location."""
return f"The weather in {location} is sunny and 72°F."
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
# Check for any relevant keywords that indicate the AI understood the context
assert any(
word in response.text.lower() for word in ["scientists", "research", "antarctica", "glaciology", "climate"]
)
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_chat_client.get_streaming_response(messages=messages)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
assert chunk.message_id is not None
assert chunk.response_id is not None
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_chat_client.get_streaming_response(
messages=messages,
tools=[get_story_text],
tool_choice="auto",
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run():
"""Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
) as agent:
# Test basic run
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "hello world" in response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run_streaming():
"""Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_thread_persistence():
"""Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
response1 = await agent.run("My name is Alice. Remember this.", thread=thread)
assert isinstance(response1, AgentRunResponse)
assert response1.text is not None
# Second interaction - test memory
response2 = await agent.run("What is my name?", thread=thread)
assert isinstance(response2, AgentRunResponse)
assert response2.text is not None
assert "alice" in response2.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_existing_thread():
"""Test Azure OpenAI chat client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my name?", thread=preserved_thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
assert "alice" in second_response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Chat Client."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@@ -0,0 +1,624 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated
import pytest
from azure.identity import AzureCliCredential
from pydantic import BaseModel
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
HostedFileSearchTool,
HostedMCPTool,
HostedVectorStoreContent,
TextContent,
ai_function,
)
from agent_framework.azure import AzureOpenAIResponsesClient
from agent_framework.exceptions import ServiceInitializationError
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str
@ai_function
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
return f"The weather in {location} is sunny and 72°F."
async def create_vector_store(client: AzureOpenAIResponsesClient) -> tuple[str, HostedVectorStoreContent]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="assistants"
)
vector_store = await client.client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
if result.last_error is not None:
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
async def delete_vector_store(client: AzureOpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after tests."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient()
assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, ChatClientProtocol)
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ServiceInitializationError):
AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"}) # type: ignore
def test_init_ai_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
ai_model_id = "test_model_id"
azure_responses_client = AzureOpenAIResponsesClient(deployment_name=ai_model_id)
assert azure_responses_client.ai_model_id == ai_model_id
assert isinstance(azure_responses_client, ChatClientProtocol)
def test_init_with_default_header(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient(
default_headers=default_headers,
)
assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, ChatClientProtocol)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in azure_responses_client.client.default_headers
assert azure_responses_client.client.default_headers[key] == value
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]], indirect=True)
def test_init_with_empty_model_id(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ServiceInitializationError):
AzureOpenAIResponsesClient(
env_file_path="test.env",
)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
settings = {
"ai_model_id": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"],
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"default_headers": default_headers,
}
azure_responses_client = AzureOpenAIResponsesClient.from_dict(settings)
dumped_settings = azure_responses_client.to_dict()
assert dumped_settings["ai_model_id"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response() -> None:
"""Test azure responses client responses."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_responses_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert "scientists" in response.text
messages.clear()
messages.append(ChatMessage(role="user", text="The weather in New York is sunny"))
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
# Test that the client can be used to get a structured response
structured_response = await azure_responses_client.get_response( # type: ignore[reportAssignmentType]
messages=messages,
response_format=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert structured_response.value.location == "New York"
assert "sunny" in structured_response.value.weather.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response_tools() -> None:
"""Test azure responses client tools."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
# Test that the client can be used to get a response
response = await azure_responses_client.get_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
assert response is not None
assert isinstance(response, ChatResponse)
assert "sunny" in response.text
messages.clear()
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
# Test that the client can be used to get a response
structured_response: ChatResponse = await azure_responses_client.get_response( # type: ignore[reportAssignmentType]
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert "Seattle" in structured_response.value.location
assert "sunny" in structured_response.value.weather.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming() -> None:
"""Test Azure azure responses client streaming responses."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(
ChatMessage(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_responses_client.get_streaming_response(messages=messages)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "scientists" in full_message
messages.clear()
messages.append(ChatMessage(role="user", text="The weather in Seattle is sunny"))
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
structured_response = await ChatResponse.from_chat_response_generator(
azure_responses_client.get_streaming_response(
messages=messages,
response_format=OutputStruct,
),
output_format_type=OutputStruct,
)
assert structured_response is not None
assert isinstance(structured_response, ChatResponse)
assert isinstance(structured_response.value, OutputStruct)
assert "Seattle" in structured_response.value.location
assert "sunny" in structured_response.value.weather.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming_tools() -> None:
"""Test azure responses client streaming tools."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
messages: list[ChatMessage] = [ChatMessage(role="user", text="What is the weather in Seattle?")]
# Test that the client can be used to get a response
response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
assert "sunny" in full_message
messages.clear()
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
structured_response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
full_message = ""
async for chunk in structured_response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_basic_run():
"""Test Azure Responses Client agent basic run functionality with AzureOpenAIResponsesClient."""
agent = AzureOpenAIResponsesClient(credential=AzureCliCredential()).create_agent(
instructions="You are a helpful assistant.",
)
# Test basic run
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
assert "hello world" in response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_basic_run_streaming():
"""Test Azure Responses Client agent basic streaming functionality with AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentRunResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_thread_persistence():
"""Test Azure Responses Client agent thread persistence across runs with AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
first_response = await agent.run("My favorite programming language is Python. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Second interaction - test memory
second_response = await agent.run("What is my favorite programming language?", thread=thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_thread_storage_with_store_true():
"""Test Azure Responses Client agent with store=True to verify service_thread_id is returned."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
) as agent:
# Create a new thread
thread = AgentThread()
# Initially, service_thread_id should be None
assert thread.service_thread_id is None
# Run with store=True to store messages on Azure/OpenAI side
response = await agent.run(
"Hello! Please remember that my name is Alex.",
thread=thread,
store=True,
)
# Validate response
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# After store=True, service_thread_id should be populated
assert thread.service_thread_id is not None
assert isinstance(thread.service_thread_id, str)
assert len(thread.service_thread_id) > 0
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_existing_thread():
"""Test Azure Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread)
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with HostedCodeInterpreterTool through AzureOpenAIResponsesClient."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Test code interpreter functionality
response = await agent.run("Calculate the sum of numbers from 1 to 10 using Python code.")
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Responses Client."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentRunResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentRunResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_chat_options_run_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
max_tokens=100,
temperature=0.7,
top_p=0.9,
seed=123,
user="comprehensive-test-user",
tools=[get_weather],
tool_choice="auto",
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_chat_options_agent_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
max_tokens=100,
temperature=0.7,
top_p=0.9,
seed=123,
user="comprehensive-test-user",
tools=[get_weather],
tool_choice="auto",
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP."""
mcp_tool = HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
)
async with ChatAgent(
chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
max_tokens=200,
)
assert isinstance(response, AgentRunResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@skip_if_azure_integration_tests_disabled
@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
async def test_azure_responses_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the web search tool
response = await azure_responses_client.get_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
tools=[HostedFileSearchTool(inputs=vector_store)],
tool_choice="auto",
)
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
assert "sunny" in response.text.lower()
assert "75" in response.text
@skip_if_azure_integration_tests_disabled
@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
async def test_azure_responses_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert isinstance(azure_responses_client, ChatClientProtocol)
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the web search tool
response = azure_responses_client.get_streaming_response(
messages=[
ChatMessage(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
tools=[HostedFileSearchTool(inputs=vector_store)],
tool_choice="auto",
)
assert response is not None
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
assert "sunny" in full_message.lower()
assert "75" in full_message
@@ -0,0 +1,156 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import AsyncMock, MagicMock
import pytest
from azure.core.exceptions import ClientAuthenticationError
from agent_framework.azure._entra_id_authentication import (
get_entra_auth_token,
get_entra_auth_token_async,
)
from agent_framework.exceptions import ServiceInvalidAuthError
@pytest.fixture
def mock_credential() -> MagicMock:
"""Mock synchronous TokenCredential."""
mock_cred = MagicMock()
# Create a mock token object with a .token attribute
mock_token = MagicMock()
mock_token.token = "test-access-token-12345"
mock_cred.get_token.return_value = mock_token
return mock_cred
@pytest.fixture
def mock_async_credential() -> MagicMock:
"""Mock asynchronous AsyncTokenCredential."""
mock_cred = MagicMock()
# Create a mock token object with a .token attribute
mock_token = MagicMock()
mock_token.token = "test-async-access-token-12345"
mock_cred.get_token = AsyncMock(return_value=mock_token)
return mock_cred
def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None:
"""Test successful token retrieval with sync function."""
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert - check the results
assert result == "test-access-token-12345"
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None:
"""Test successful token retrieval with async function."""
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert - check the results
assert result == "test-async-access-token-12345"
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None:
"""Test that missing token endpoint raises ServiceInvalidAuthError."""
# Test with empty string
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
get_entra_auth_token(mock_credential, "")
# Test with None
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
get_entra_auth_token(mock_credential, None) # type: ignore
async def test_get_entra_auth_token_async_missing_endpoint(mock_async_credential: MagicMock) -> None:
"""Test that missing token endpoint raises ServiceInvalidAuthError in async function."""
# Test with empty string
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
await get_entra_auth_token_async(mock_async_credential, "")
# Test with None
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
await get_entra_auth_token_async(mock_async_credential, None) # type: ignore
def test_get_entra_auth_token_auth_failure(mock_credential: MagicMock) -> None:
"""Test that Azure authentication failure returns None."""
mock_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert - should return None on auth failure
assert result is None
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_auth_failure(mock_async_credential: MagicMock) -> None:
"""Test that Azure authentication failure returns None in async function."""
mock_async_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert - should return None on auth failure
assert result is None
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_none_token_response(mock_credential: MagicMock) -> None:
"""Test that None token response returns None."""
mock_credential.get_token.return_value = None
token_endpoint = "https://test-endpoint.com/.default"
result = get_entra_auth_token(mock_credential, token_endpoint)
# Assert
assert result is None
mock_credential.get_token.assert_called_once_with(token_endpoint)
async def test_get_entra_auth_token_async_none_token_response(mock_async_credential: MagicMock) -> None:
"""Test that None token response returns None in async function."""
mock_async_credential.get_token.return_value = None
token_endpoint = "https://test-endpoint.com/.default"
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
# Assert
assert result is None
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
def test_get_entra_auth_token_with_kwargs(mock_credential: MagicMock) -> None:
"""Test that kwargs are passed through to get_token."""
token_endpoint = "https://test-endpoint.com/.default"
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
result = get_entra_auth_token(mock_credential, token_endpoint, **extra_kwargs)
# Assert
assert result == "test-access-token-12345"
mock_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
async def test_get_entra_auth_token_async_with_kwargs(mock_async_credential: MagicMock) -> None:
"""Test that kwargs are passed through to async get_token."""
token_endpoint = "https://test-endpoint.com/.default"
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint, **extra_kwargs)
# Assert
assert result == "test-async-access-token-12345"
mock_async_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
@@ -1560,7 +1560,7 @@ async def test_openai_responses_client_agent_chat_options_agent_level() -> None:
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with OpenAI Response Agent using Microsoft Learn MCP."""
# Use the same MCP server as the Foundry example
mcp_tool = HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
@@ -1573,7 +1573,6 @@ async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
# Use the same query as the Foundry example
response = await agent.run(
"How to create an Azure storage account using az cli?",
max_tokens=200,