mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
366a7f7d47
commit
9355329dfd
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user