agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] Python: Provider-leading client design & OpenAI package extraction (#4818)

Browse Source 
* Python: Provider-leading client design & OpenAI package extraction

Major refactoring of the Python Agent Framework client architecture:

- Extract OpenAI clients into new `agent-framework-openai` package
- Core package no longer depends on openai, azure-identity, azure-ai-projects
- Rename clients for discoverability: OpenAIResponsesClient → OpenAIChatClient,
  OpenAIChatClient → OpenAIChatCompletionClient
- Unify `model_id`/`deployment_name`/`model_deployment_name` → `model` param
- New FoundryChatClient for Azure AI Foundry Responses API
- New FoundryAgent/FoundryAgentClient for connecting to pre-configured Foundry agents
- Remove OpenAIBase/OpenAIConfigMixin from non-deprecated client MRO
- Deprecate AzureOpenAI* clients, AzureAIClient, OpenAIAssistantsClient
- Reorganize samples: azure_openai+azure_ai+azure_ai_agent → azure/
- ADR-0020: Provider-Leading Client Design

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: missing Agent imports in samples, .model_id → .model in foundry_local sample

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: CI failures — mypy errors, coverage targets, sample imports

- azure-ai mypy: add type ignores for TypedDict total=, model arg, forward ref
- Coverage: replace core.azure/openai targets with openai package target
- project_provider: add type annotation for opts dict

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: populate openai .pyi stub, fix broken README links, coverage targets

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fixes

* updated observabilitty

* reset azure init.pyi

* fix errors

* updated adr number

* fix foundry local

* fixed not renamed docstrings and comments, and added deprecated markers to old classes

* fix tests and pyprojects

* fix test vars

* updated function tests

* update durable

* updated test setup for functions

* Fix Foundry auth in workflow samples

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Stabilize Python integration workflows

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Update hosting samples for Foundry

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Trigger full CI rerun

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Trigger CI rerun again

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* trigger rerun

* trigger rerun

* fix for litellm

* undo durabletask changes

* Move Foundry APIs into foundry namespace

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix Foundry pyproject formatting

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Split provider samples by Foundry surface

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Restore hosting sample requirements

Also fix the Foundry Local sample link after the provider sample move.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* updated tests

* udpated foundry integration tests

* removed dist from azurefunctions tests

* Use separate Foundry clients for concurrent agents

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix client setup in azfunc and durable

* disabled two tests

* updated setup for some function and durable tests

* improved azure openai setup with new clients

* ignore deprecated

* fixes

* skip 11

* remove openai assistants int tests

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-25 09:56:29 +00:00
committed by GitHub
co-authored by Copilot
parent 4b533608b6
commit 5e056b672e
485 changed files with 9784 additions and 12084 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure-ai/agent_framework_azure_ai/__init__.py
+30 -7
View File
@@ -2,23 +2,35 @@
import importlib.metadata
from ._agent_provider import AzureAIAgentsProvider
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions
from ._client import AzureAIClient, AzureAIProjectAgentOptions, RawAzureAIClient
from ._agent_provider import AzureAIAgentsProvider # pyright: ignore[reportDeprecated]
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions # pyright: ignore[reportDeprecated]
from ._client import AzureAIClient, AzureAIProjectAgentOptions, RawAzureAIClient # pyright: ignore[reportDeprecated]
from ._deprecated_azure_openai import (
AzureOpenAIAssistantsClient, # pyright: ignore[reportDeprecated]
AzureOpenAIAssistantsOptions,
AzureOpenAIChatClient, # pyright: ignore[reportDeprecated]
AzureOpenAIChatOptions,
AzureOpenAIConfigMixin,
AzureOpenAIEmbeddingClient, # pyright: ignore[reportDeprecated]
AzureOpenAIResponsesClient, # pyright: ignore[reportDeprecated]
AzureOpenAIResponsesOptions,
AzureOpenAISettings,
AzureUserSecurityContext,
)
from ._embedding_client import (
AzureAIInferenceEmbeddingClient,
AzureAIInferenceEmbeddingOptions,
AzureAIInferenceEmbeddingSettings,
RawAzureAIInferenceEmbeddingClient,
)
from ._foundry_memory_provider import FoundryMemoryProvider
from ._project_provider import AzureAIProjectAgentProvider
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._project_provider import AzureAIProjectAgentProvider # pyright: ignore[reportDeprecated]
from ._shared import AzureAISettings
try:
__version__ = importlib.metadata.version(__name__)
except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
__version__ = "0.0.0"
__all__ = [
"AzureAIAgentClient",
@@ -31,7 +43,18 @@ __all__ = [
"AzureAIProjectAgentOptions",
"AzureAIProjectAgentProvider",
"AzureAISettings",
"FoundryMemoryProvider",
"AzureCredentialTypes",
"AzureOpenAIAssistantsClient",
"AzureOpenAIAssistantsOptions",
"AzureOpenAIChatClient",
"AzureOpenAIChatOptions",
"AzureOpenAIConfigMixin",
"AzureOpenAIEmbeddingClient",
"AzureOpenAIResponsesClient",
"AzureOpenAIResponsesOptions",
"AzureOpenAISettings",
"AzureTokenProvider",
"AzureUserSecurityContext",
"RawAzureAIClient",
"RawAzureAIInferenceEmbeddingClient",
"__version__",
python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py
+12 -3
View File
@@ -18,19 +18,23 @@ from agent_framework import (
from agent_framework._mcp import MCPTool
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from azure.ai.agents.aio import AgentsClient
from azure.ai.agents.models import Agent as AzureAgent
from azure.ai.agents.models import ResponseFormatJsonSchema, ResponseFormatJsonSchemaType
from pydantic import BaseModel
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions # pyright: ignore[reportDeprecated]
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import AzureAISettings, to_azure_ai_agent_tools
if sys.version_info >= (3, 13):
from typing import Self, TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import Self, TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 13):
from warnings import deprecated # type: ignore # pragma: no cover
else:
from typing_extensions import deprecated # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
@@ -47,6 +51,11 @@ OptionsCoT = TypeVar(
)
@deprecated(
"AzureAIAgentClient and the AzureAIAgentsProvider are deprecated. "
"They target the V1 Agents Service API and have no direct replacement; "
"for new Foundry projects, use FoundryAgent."
)
class AzureAIAgentsProvider(Generic[OptionsCoT]):
"""Provider for Azure AI Agent Service V1 (Persistent Agents API).
@@ -426,7 +435,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
context_providers: Context providers to include during agent invocation.
"""
# Create the underlying client
client = AzureAIAgentClient(
client = AzureAIAgentClient( # pyright: ignore[reportDeprecated]
agents_client=self._agents_client,
agent_id=agent.id,
agent_name=agent.name,
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+22 -17
View File
@@ -36,7 +36,6 @@ from agent_framework import (
)
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import (
ChatClientException,
ChatClientInvalidRequestException,
@@ -92,12 +91,14 @@ from azure.ai.agents.models import (
)
from pydantic import BaseModel
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import AzureAISettings, resolve_file_ids, to_azure_ai_agent_tools
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
from warnings import deprecated # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
from typing_extensions import TypeVar, deprecated # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
@@ -210,6 +211,11 @@ AzureAIAgentOptionsT = TypeVar(
# endregion
@deprecated(
"AzureAIAgentClient is deprecated. "
"It targets the V1 Agents Service API and has no direct replacement; "
"for new Foundry projects, use FoundryAgent."
)
class AzureAIAgentClient(
FunctionInvocationLayer[AzureAIAgentOptionsT],
ChatMiddlewareLayer[AzureAIAgentOptionsT],
@@ -221,7 +227,8 @@ class AzureAIAgentClient(
.. deprecated::
AzureAIAgentClient is deprecated and will be removed in a future release.
Use :class:`AzureAIClient` instead for the V2 (Projects/Responses) API.
It targets the V1 Agents Service API and has no direct replacement.
For new Foundry projects, use :class:`FoundryAgent`.
"""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -239,7 +246,8 @@ class AzureAIAgentClient(
.. deprecated::
This method is deprecated and will be removed in a future release.
Use :meth:`AzureAIClient.get_code_interpreter_tool` instead.
For new Foundry projects, configure hosted tools on the Foundry agent definition
in the service instead.
Keyword Args:
file_ids: List of uploaded file IDs or Content objects to make available to
@@ -272,7 +280,7 @@ class AzureAIAgentClient(
"""
warnings.warn(
"AzureAIAgentClient.get_code_interpreter_tool() is deprecated and will be removed in a future release; "
"use AzureAIClient.get_code_interpreter_tool() instead.",
"for new Foundry projects, configure hosted tools on the Foundry agent definition in the service instead.",
DeprecationWarning,
stacklevel=2,
)
@@ -288,7 +296,8 @@ class AzureAIAgentClient(
.. deprecated::
This method is deprecated and will be removed in a future release.
Use :meth:`AzureAIClient.get_file_search_tool` instead.
For new Foundry projects, configure hosted tools on the Foundry agent definition
in the service instead.
Keyword Args:
vector_store_ids: List of vector store IDs to search within.
@@ -308,7 +317,7 @@ class AzureAIAgentClient(
"""
warnings.warn(
"AzureAIAgentClient.get_file_search_tool() is deprecated and will be removed in a future release; "
"use AzureAIClient.get_file_search_tool() instead.",
"for new Foundry projects, configure hosted tools on the Foundry agent definition in the service instead.",
DeprecationWarning,
stacklevel=2,
)
@@ -325,7 +334,8 @@ class AzureAIAgentClient(
.. deprecated::
This method is deprecated and will be removed in a future release.
Use :meth:`AzureAIClient.get_web_search_tool` instead.
For new Foundry projects, configure hosted tools on the Foundry agent definition
in the service instead.
For Azure AI Agents, web search uses Bing Grounding or Bing Custom Search.
If no arguments are provided, attempts to read from environment variables.
@@ -369,7 +379,7 @@ class AzureAIAgentClient(
"""
warnings.warn(
"AzureAIAgentClient.get_web_search_tool() is deprecated and will be removed in a future release; "
"use AzureAIClient.get_web_search_tool() instead.",
"for new Foundry projects, configure hosted tools on the Foundry agent definition in the service instead.",
DeprecationWarning,
stacklevel=2,
)
@@ -410,7 +420,8 @@ class AzureAIAgentClient(
.. deprecated::
This method is deprecated and will be removed in a future release.
Use :meth:`AzureAIClient.get_mcp_tool` instead.
For new Foundry projects, configure hosted tools on the Foundry agent definition
in the service instead.
This configures an MCP (Model Context Protocol) server that will be called
by Azure AI's service. The tools from this MCP server are executed remotely
@@ -446,7 +457,7 @@ class AzureAIAgentClient(
"""
warnings.warn(
"AzureAIAgentClient.get_mcp_tool() is deprecated and will be removed in a future release; "
"use AzureAIClient.get_mcp_tool() instead.",
"for new Foundry projects, configure hosted tools on the Foundry agent definition in the service instead.",
DeprecationWarning,
stacklevel=2,
)
@@ -561,12 +572,6 @@ class AzureAIAgentClient(
client: AzureAIAgentClient[MyOptions] = AzureAIAgentClient(credential=credential)
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
warnings.warn(
"AzureAIAgentClient is deprecated and will be removed in a future release; "
"use AzureAIClient instead for the V2 (Projects/Responses) API.",
DeprecationWarning,
stacklevel=2,
)
azure_ai_settings = load_settings(
AzureAISettings,
env_prefix="AZURE_AI_",
python/packages/azure-ai/agent_framework_azure_ai/_client.py
+27 -20
View File
@@ -30,10 +30,9 @@ from agent_framework import (
)
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai import OpenAIResponsesOptions
from agent_framework.openai._responses_client import RawOpenAIResponsesClient
from agent_framework_openai._chat_client import RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
ApproximateLocation,
@@ -50,12 +49,14 @@ from azure.ai.projects.models import (
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.core.exceptions import ResourceNotFoundError
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import AzureAISettings, create_text_format_config, resolve_file_ids
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
from warnings import deprecated # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
from typing_extensions import TypeVar, deprecated # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
@@ -68,7 +69,7 @@ else:
logger = logging.getLogger("agent_framework.azure")
class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False):
class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False): # type: ignore[misc, call-arg]
"""Azure AI Project Agent options."""
rai_config: RaiConfig
@@ -88,8 +89,13 @@ AzureAIClientOptionsT = TypeVar(
_DOC_INDEX_PATTERN = re.compile(r"doc_(\d+)")
class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[AzureAIClientOptionsT]):
"""Raw Azure AI client without middleware, telemetry, or function invocation layers.
@deprecated(
"RawAzureAIClient is deprecated. "
"Use RawFoundryAgentChatClient for low-level Foundry agent client customization, "
"or FoundryAgent for the recommended production API."
)
class RawAzureAIClient(RawOpenAIChatClient[AzureAIClientOptionsT], Generic[AzureAIClientOptionsT]):
"""Deprecated raw Azure AI client without middleware, telemetry, or function invocation layers.
Warning:
**This class should not normally be used directly.** It does not include middleware,
@@ -101,7 +107,8 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry
3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry
Use ``AzureAIClient`` instead for a fully-featured client with all layers applied.
Use ``RawFoundryAgentChatClient`` for low-level Foundry agent customization, or
``FoundryAgent`` for the recommended production API.
"""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
@@ -215,8 +222,10 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
project_client = AIProjectClient(**project_client_kwargs)
should_close_client = True
# Initialize parent
super().__init__(
# Initialize parent with OpenAI client from project
super().__init__( # type: ignore
async_client=project_client.get_openai_client(),
model=azure_ai_settings.get("model"), # type: ignore[arg-type]
additional_properties=additional_properties,
)
@@ -680,10 +689,6 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
return result, instructions
async def _initialize_client(self) -> None:
"""Initialize OpenAI client."""
self.client = self.project_client.get_openai_client() # type: ignore
def _update_agent_name_and_description(self, agent_name: str | None, description: str | None = None) -> None:
"""Update the agent name in the chat client.
@@ -842,7 +847,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
if not stream:
async def _enrich_response() -> ChatResponse:
response = await super(RawAzureAIClient, self)._inner_get_response(
response = await super(RawAzureAIClient, self)._inner_get_response( # pyright: ignore[reportDeprecated]
messages=messages, options=options, stream=False, **kwargs
)
get_urls = self._extract_azure_search_urls(response.raw_representation.output) # type: ignore[union-attr]
@@ -1182,8 +1187,8 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
It does NOT create an agent on the Azure AI service - the actual agent
will be created on the server during the first invocation (run).
For creating and managing persistent agents on the server, use
:class:`~agent_framework_azure_ai.AzureAIProjectAgentProvider` instead.
For working with pre-configured persistent agents on the server, use
:class:`~agent_framework_azure_ai.FoundryAgent` instead.
Keyword Args:
id: The unique identifier for the agent. Will be created automatically if not provided.
@@ -1213,21 +1218,23 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
)
@deprecated("AzureAIClient is deprecated. Use FoundryAgent instead.")
class AzureAIClient(
FunctionInvocationLayer[AzureAIClientOptionsT],
ChatMiddlewareLayer[AzureAIClientOptionsT],
ChatTelemetryLayer[AzureAIClientOptionsT],
RawAzureAIClient[AzureAIClientOptionsT],
RawAzureAIClient[AzureAIClientOptionsT], # pyright: ignore[reportDeprecated]
Generic[AzureAIClientOptionsT],
):
"""Azure AI client with middleware, telemetry, and function invocation support.
"""Deprecated Azure AI client with middleware, telemetry, and function invocation support.
This is the recommended client for most use cases. It includes:
This class is deprecated. Use ``FoundryAgent`` instead for connecting to
pre-configured agents in Foundry. It includes:
- Chat middleware support for request/response interception
- OpenTelemetry-based telemetry for observability
- Automatic function/tool invocation handling
For a minimal implementation without these features, use :class:`RawAzureAIClient`.
For a minimal implementation without these features, use :class:`RawFoundryAgentChatClient`.
"""
def __init__(
python/packages/azure-ai/agent_framework_azure_ai/_deprecated_azure_openai.py
+897
View File
@@ -0,0 +1,897 @@
# Copyright (c) Microsoft. All rights reserved.
"""Deprecated Azure OpenAI client classes.
All classes in this module are deprecated and will be removed in a future release.
Migrate to the ``agent_framework_openai`` package equivalents with an ``AsyncAzureOpenAI`` client,
or use ``FoundryChatClient`` for Azure AI Foundry projects.
"""
from __future__ import annotations
import json
import logging
import sys
from collections.abc import Mapping, Sequence
from copy import copy
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, cast
from urllib.parse import urljoin, urlparse
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT, APP_INFO, prepend_agent_framework_to_user_agent
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
from agent_framework._types import Annotation, Content
from agent_framework.observability import ChatTelemetryLayer, EmbeddingTelemetryLayer
from agent_framework_openai._assistants_client import OpenAIAssistantsClient, OpenAIAssistantsOptions
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from agent_framework_openai._chat_completion_client import OpenAIChatCompletionOptions, RawOpenAIChatCompletionClient
from agent_framework_openai._embedding_client import OpenAIEmbeddingOptions, RawOpenAIEmbeddingClient
from agent_framework_openai._shared import OpenAIBase
from azure.ai.projects.aio import AIProjectClient
from openai import AsyncOpenAI
from openai.lib.azure import AsyncAzureOpenAI
from pydantic import BaseModel
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
from warnings import deprecated # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar, deprecated # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import MiddlewareTypes
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
logger: logging.Logger = logging.getLogger(__name__)
# region Constants and Settings
DEFAULT_AZURE_API_VERSION: Final[str] = "2024-10-21"
DEFAULT_AZURE_TOKEN_ENDPOINT: Final[str] = "https://cognitiveservices.azure.com/.default" # noqa: S105
class AzureOpenAISettings(TypedDict, total=False):
"""AzureOpenAI model settings.
Settings are resolved in this order: explicit keyword arguments, values from an
explicitly provided .env file, then environment variables with the prefix
'AZURE_OPENAI_'. If settings are missing after resolution, validation will fail.
Keyword Args:
endpoint: The endpoint of the Azure deployment.
chat_deployment_name: The name of the Azure Chat deployment.
responses_deployment_name: The name of the Azure Responses deployment.
embedding_deployment_name: The name of the Azure Embedding deployment.
api_key: The API key for the Azure deployment.
api_version: The API version to use.
base_url: The url of the Azure deployment.
token_endpoint: The token endpoint to use to retrieve the authentication token.
"""
chat_deployment_name: str | None
responses_deployment_name: str | None
embedding_deployment_name: str | None
endpoint: str | None
base_url: str | None
api_key: SecretString | None
api_version: str | None
token_endpoint: str | None
def _apply_azure_defaults(
settings: AzureOpenAISettings,
default_api_version: str = DEFAULT_AZURE_API_VERSION,
default_token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT,
) -> None:
"""Apply default values for api_version and token_endpoint after loading settings.
Args:
settings: The loaded Azure OpenAI settings dict.
default_api_version: The default API version to use if not set.
default_token_endpoint: The default token endpoint to use if not set.
"""
if not settings.get("api_version"):
settings["api_version"] = default_api_version
if not settings.get("token_endpoint"):
settings["token_endpoint"] = default_token_endpoint
# endregion
# region AzureOpenAIConfigMixin
class AzureOpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an Azure OpenAI service."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai"
def __init__(
self,
deployment_name: str,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str = DEFAULT_AZURE_API_VERSION,
api_key: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncOpenAI | None = None,
instruction_role: str | None = None,
**kwargs: Any,
) -> None:
"""Configure a connection to an Azure OpenAI service.
Args:
deployment_name: Name of the deployment.
endpoint: The specific endpoint URL for the deployment.
base_url: The base URL for Azure services.
api_version: Azure API version.
api_key: API key for Azure services.
token_endpoint: Azure AD token scope.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
client: An existing client to use.
instruction_role: The role to use for 'instruction' messages.
kwargs: Additional keyword arguments.
"""
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:
ad_token_provider = None
if not api_key and credential:
ad_token_provider = resolve_credential_to_token_provider(credential, token_endpoint)
if not api_key and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
if not endpoint and not base_url:
raise ValueError("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_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)
self.endpoint = str(endpoint)
self.base_url = str(base_url)
self.api_version = api_version
self.deployment_name = deployment_name
self.instruction_role = instruction_role
if default_headers:
from agent_framework._telemetry import USER_AGENT_KEY
def_headers = {k: v for k, v in default_headers.items() if k != USER_AGENT_KEY}
else:
def_headers = None
self.default_headers = def_headers
super().__init__(model_id=deployment_name, client=client, **kwargs)
# endregion
# region AzureOpenAIResponsesClient
AzureOpenAIResponsesOptionsT = TypeVar(
"AzureOpenAIResponsesOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIChatOptions",
covariant=True,
)
AzureOpenAIResponsesOptions = OpenAIChatOptions
@deprecated(
"AzureOpenAIResponsesClient is deprecated. "
"Use OpenAIChatClient with an AsyncAzureOpenAI client, or FoundryChatClient for Foundry projects."
)
class AzureOpenAIResponsesClient( # type: ignore[misc]
FunctionInvocationLayer[AzureOpenAIResponsesOptionsT],
ChatMiddlewareLayer[AzureOpenAIResponsesOptionsT],
ChatTelemetryLayer[AzureOpenAIResponsesOptionsT],
RawOpenAIChatClient[AzureOpenAIResponsesOptionsT],
Generic[AzureOpenAIResponsesOptionsT],
):
"""Deprecated Azure Responses client. Use OpenAIChatClient with an AsyncAzureOpenAI client instead."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai"
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,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
project_client: Any | None = None,
project_endpoint: str | None = None,
allow_preview: bool | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure OpenAI Responses client.
Keyword Args:
api_key: The API key.
deployment_name: The deployment name.
endpoint: The deployment endpoint.
base_url: The deployment base URL.
api_version: The deployment API version.
token_endpoint: The token endpoint to request an Azure token.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
project_client: An existing AIProjectClient to use.
project_endpoint: The Azure AI Foundry project endpoint URL.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
instruction_role: The role to use for 'instruction' messages.
middleware: Optional sequence of middleware.
function_invocation_configuration: Optional function invocation configuration.
kwargs: Additional keyword arguments.
"""
if (model_id := kwargs.pop("model_id", None)) and not deployment_name:
deployment_name = str(model_id)
if async_client is None and (project_client is not None or project_endpoint is not None):
async_client = self._create_client_from_project(
project_client=project_client,
project_endpoint=project_endpoint,
credential=credential,
allow_preview=allow_preview,
)
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
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,
)
_apply_azure_defaults(azure_openai_settings, default_api_version="preview")
endpoint_value = azure_openai_settings.get("endpoint")
if (
not azure_openai_settings.get("base_url")
and endpoint_value
and (hostname := urlparse(str(endpoint_value)).hostname)
and hostname.endswith(".openai.azure.com")
):
azure_openai_settings["base_url"] = urljoin(str(endpoint_value), "/openai/v1/")
responses_deployment_name = azure_openai_settings.get("responses_deployment_name")
if not responses_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
)
if not async_client:
# Create the Azure OpenAI client directly
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)
api_key_secret = azure_openai_settings.get("api_key")
ad_token_provider = None
if not api_key_secret and credential:
ad_token_provider = resolve_credential_to_token_provider(
credential, azure_openai_settings.get("token_endpoint")
)
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
client_endpoint = azure_openai_settings.get("endpoint")
client_base_url = azure_openai_settings.get("base_url")
if not client_endpoint and not client_base_url:
raise ValueError("Please provide an endpoint or a base_url")
client_args: dict[str, Any] = {"default_headers": merged_headers}
if resolved_api_version := azure_openai_settings.get("api_version"):
client_args["api_version"] = resolved_api_version
if ad_token_provider:
client_args["azure_ad_token_provider"] = ad_token_provider
if api_key_secret:
client_args["api_key"] = api_key_secret.get_secret_value()
if client_base_url:
client_args["base_url"] = str(client_base_url)
if client_endpoint and not client_base_url:
client_args["azure_endpoint"] = str(client_endpoint)
if responses_deployment_name:
client_args["azure_deployment"] = responses_deployment_name
if "websocket_base_url" in kwargs:
client_args["websocket_base_url"] = kwargs.pop("websocket_base_url")
async_client = AsyncAzureOpenAI(**client_args)
# Store Azure-specific attributes for serialization
self.endpoint = str(endpoint_value) if endpoint_value else None
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = responses_deployment_name
super().__init__(
async_client=async_client,
model=responses_deployment_name,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
@staticmethod
def _create_client_from_project(
*,
project_client: AIProjectClient | None,
project_endpoint: str | None,
credential: AzureCredentialTypes | AzureTokenProvider | None,
allow_preview: bool | None = None,
) -> AsyncOpenAI:
"""Create an AsyncOpenAI client from an Azure AI Foundry project."""
if project_client is not None:
return project_client.get_openai_client()
if not project_endpoint:
raise ValueError("Azure AI project endpoint is required when project_client is not provided.")
if not credential:
raise ValueError("Azure credential is required when using project_endpoint without a project_client.")
project_client_kwargs: dict[str, Any] = {
"endpoint": project_endpoint,
"credential": credential, # type: ignore[arg-type]
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
project_client = AIProjectClient(**project_client_kwargs)
return project_client.get_openai_client()
@override
def _check_model_presence(self, options: dict[str, Any]) -> None:
if not options.get("model"):
if not self.model:
raise ValueError("deployment_name must be a non-empty string")
options["model"] = self.model
# endregion
# region AzureOpenAIChatClient
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
class AzureUserSecurityContext(TypedDict, total=False):
"""User security context for Azure AI applications.
These fields help security operations teams investigate and mitigate security
incidents by providing context about the application and end user.
"""
application_name: str
"""Name of the application making the request."""
end_user_id: str
"""Unique identifier for the end user (recommend hashing username/email)."""
end_user_tenant_id: str
"""Microsoft 365 tenant ID the end user belongs to. Required for multi-tenant apps."""
source_ip: str
"""The original client's IP address."""
class AzureOpenAIChatOptions(OpenAIChatCompletionOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""Azure OpenAI-specific chat options dict.
Extends OpenAIChatCompletionOptions with Azure-specific options including
the "On Your Data" feature and enhanced security context.
"""
data_sources: list[dict[str, Any]]
"""Azure "On Your Data" data sources for retrieval-augmented generation."""
user_security_context: AzureUserSecurityContext
"""Enhanced security context for Azure Defender integration."""
n: int
"""Number of chat completion choices to generate for each input message."""
AzureOpenAIChatOptionsT = TypeVar(
"AzureOpenAIChatOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AzureOpenAIChatOptions",
covariant=True,
)
@deprecated("AzureOpenAIChatClient is deprecated. Use OpenAIChatCompletionClient with an AsyncAzureOpenAI client.")
class AzureOpenAIChatClient( # type: ignore[misc]
FunctionInvocationLayer[AzureOpenAIChatOptionsT],
ChatMiddlewareLayer[AzureOpenAIChatOptionsT],
ChatTelemetryLayer[AzureOpenAIChatOptionsT],
RawOpenAIChatCompletionClient[AzureOpenAIChatOptionsT],
Generic[AzureOpenAIChatOptionsT],
):
"""Deprecated Azure OpenAI Chat client. Use OpenAIChatCompletionClient with AsyncAzureOpenAI instead."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai"
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,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
additional_properties: dict[str, Any] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None:
"""Initialize an Azure OpenAI Chat completion client.
Keyword Args:
api_key: The API key.
deployment_name: The deployment name.
endpoint: The deployment endpoint.
base_url: The deployment base URL.
api_version: The deployment API version.
token_endpoint: The token endpoint to request an Azure token.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
additional_properties: Additional properties stored on the client instance.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
instruction_role: The role to use for 'instruction' messages.
middleware: Optional sequence of middleware.
function_invocation_configuration: Optional function invocation configuration.
"""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
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,
)
_apply_azure_defaults(azure_openai_settings)
chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
if not chat_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
if not async_client:
# Create the Azure OpenAI client directly
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)
api_key_secret = azure_openai_settings.get("api_key")
ad_token_provider = None
if not api_key_secret and credential:
ad_token_provider = resolve_credential_to_token_provider(
credential, azure_openai_settings.get("token_endpoint")
)
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not endpoint_value and not base_url_value:
raise ValueError("Please provide an endpoint or a base_url")
client_args: dict[str, Any] = {"default_headers": merged_headers}
if resolved_api_version := azure_openai_settings.get("api_version"):
client_args["api_version"] = resolved_api_version
if ad_token_provider:
client_args["azure_ad_token_provider"] = ad_token_provider
if api_key_secret:
client_args["api_key"] = api_key_secret.get_secret_value()
if base_url_value:
client_args["base_url"] = str(base_url_value)
if endpoint_value and not base_url_value:
client_args["azure_endpoint"] = str(endpoint_value)
if chat_deployment_name:
client_args["azure_deployment"] = chat_deployment_name
async_client = AsyncAzureOpenAI(**client_args)
# Store Azure-specific attributes for serialization
self.endpoint = str(azure_openai_settings.get("endpoint") or "")
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = chat_deployment_name
super().__init__(
async_client=async_client,
model=chat_deployment_name,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
additional_properties=additional_properties,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
)
@override
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
"""Parse the choice into a Content object with type='text'.
Overwritten from RawOpenAIChatCompletionClient to deal with Azure On Your Data function.
"""
message = getattr(choice, "message", None)
if message is None:
message = getattr(choice, "delta", None)
if message is None: # type: ignore
return None
if hasattr(message, "refusal") and message.refusal:
return Content.from_text(text=message.refusal, raw_representation=choice)
if not message.content:
return None
text_content = Content.from_text(text=message.content, raw_representation=choice)
if not message.model_extra or "context" not in message.model_extra:
return text_content
context_raw: object = cast(object, message.context) # type: ignore[union-attr]
if isinstance(context_raw, str):
try:
context_raw = json.loads(context_raw)
except json.JSONDecodeError:
logger.warning("Context is not a valid JSON string, ignoring context.")
return text_content
if not isinstance(context_raw, dict):
logger.warning("Context is not a valid dictionary, ignoring context.")
return text_content
context = cast(dict[str, Any], context_raw)
if intent := context.get("intent"):
text_content.additional_properties = {"intent": intent}
citations = context.get("citations")
if isinstance(citations, list) and citations:
annotations: list[Annotation] = []
for citation_raw in cast(list[object], citations):
if not isinstance(citation_raw, dict):
continue
citation = cast(dict[str, Any], citation_raw)
annotations.append(
Annotation(
type="citation",
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,
)
)
text_content.annotations = annotations
return text_content
# endregion
# region AzureOpenAIAssistantsClient
AzureOpenAIAssistantsOptionsT = TypeVar(
"AzureOpenAIAssistantsOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIAssistantsOptions",
covariant=True,
)
AzureOpenAIAssistantsOptions = OpenAIAssistantsOptions
@deprecated(
"AzureOpenAIAssistantsClient is deprecated. "
"Use OpenAIAssistantsClient (also deprecated) or migrate to OpenAIChatClient."
)
class AzureOpenAIAssistantsClient(
OpenAIAssistantsClient[AzureOpenAIAssistantsOptionsT], Generic[AzureOpenAIAssistantsOptionsT]
):
"""Deprecated Azure OpenAI Assistants client. Use OpenAIAssistantsClient or migrate to OpenAIChatClient."""
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,
assistant_description: 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,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | 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.
Keyword Args:
deployment_name: The Azure OpenAI deployment name.
assistant_id: The ID of an Azure OpenAI assistant to use.
assistant_name: The name to use when creating new assistants.
assistant_description: The description to use when creating new assistants.
thread_id: Default thread ID to use for conversations.
api_key: The API key to use.
endpoint: The deployment endpoint.
base_url: The deployment base URL.
api_version: The deployment API version.
token_endpoint: The token endpoint to request an Azure token.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
"""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
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,
)
_apply_azure_defaults(azure_openai_settings, default_api_version=self.DEFAULT_AZURE_API_VERSION)
chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
if not chat_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
api_key_secret = azure_openai_settings.get("api_key")
token_scope = azure_openai_settings.get("token_endpoint")
ad_token_provider = None
if not async_client and not api_key_secret and credential:
ad_token_provider = resolve_credential_to_token_provider(credential, token_scope)
if not async_client and not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
if not async_client:
client_params: dict[str, Any] = {
"default_headers": default_headers,
}
if resolved_api_version := azure_openai_settings.get("api_version"):
client_params["api_version"] = resolved_api_version
if api_key_secret:
client_params["api_key"] = api_key_secret.get_secret_value()
elif ad_token_provider:
client_params["azure_ad_token_provider"] = ad_token_provider
if resolved_base_url := azure_openai_settings.get("base_url"):
client_params["base_url"] = str(resolved_base_url)
elif resolved_endpoint := azure_openai_settings.get("endpoint"):
client_params["azure_endpoint"] = str(resolved_endpoint)
async_client = AsyncAzureOpenAI(**client_params)
super().__init__(
model_id=chat_deployment_name,
assistant_id=assistant_id,
assistant_name=assistant_name,
assistant_description=assistant_description,
thread_id=thread_id,
async_client=async_client, # type: ignore[reportArgumentType]
default_headers=default_headers,
)
# endregion
# region AzureOpenAIEmbeddingClient
AzureOpenAIEmbeddingOptionsT = TypeVar(
"AzureOpenAIEmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIEmbeddingOptions",
covariant=True,
)
@deprecated("AzureOpenAIEmbeddingClient is deprecated. Use OpenAIEmbeddingClient with an AsyncAzureOpenAI client.")
class AzureOpenAIEmbeddingClient(
EmbeddingTelemetryLayer[str, list[float], AzureOpenAIEmbeddingOptionsT],
RawOpenAIEmbeddingClient[AzureOpenAIEmbeddingOptionsT],
Generic[AzureOpenAIEmbeddingOptionsT],
):
"""Deprecated Azure OpenAI embedding client. Use OpenAIEmbeddingClient with AsyncAzureOpenAI instead."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai"
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,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
otel_provider_name: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an Azure OpenAI embedding client.
Keyword Args:
api_key: The API key.
deployment_name: The deployment name.
endpoint: The deployment endpoint.
base_url: The deployment base URL.
api_version: The deployment API version.
token_endpoint: The token endpoint to request an Azure token.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
otel_provider_name: Override the OpenTelemetry provider name.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
"""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
embedding_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings)
embedding_deployment_name = azure_openai_settings.get("embedding_deployment_name")
if not embedding_deployment_name:
raise ValueError(
"Azure OpenAI embedding deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME' environment variable."
)
if not async_client:
# Create the Azure OpenAI client directly
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)
api_key_secret = azure_openai_settings.get("api_key")
ad_token_provider = None
if not api_key_secret and credential:
ad_token_provider = resolve_credential_to_token_provider(
credential, azure_openai_settings.get("token_endpoint")
)
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not endpoint_value and not base_url_value:
raise ValueError("Please provide an endpoint or a base_url")
client_args: dict[str, Any] = {"default_headers": merged_headers}
if resolved_api_version := azure_openai_settings.get("api_version"):
client_args["api_version"] = resolved_api_version
if ad_token_provider:
client_args["azure_ad_token_provider"] = ad_token_provider
if api_key_secret:
client_args["api_key"] = api_key_secret.get_secret_value()
if base_url_value:
client_args["base_url"] = str(base_url_value)
if endpoint_value and not base_url_value:
client_args["azure_endpoint"] = str(endpoint_value)
if embedding_deployment_name:
client_args["azure_deployment"] = embedding_deployment_name
async_client = AsyncAzureOpenAI(**client_args)
# Store Azure-specific attributes for serialization
self.endpoint = str(azure_openai_settings.get("endpoint") or "")
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = embedding_deployment_name
super().__init__(
async_client=async_client,
model=embedding_deployment_name,
default_headers=default_headers,
)
if otel_provider_name is not None:
self.OTEL_PROVIDER_NAME = otel_provider_name # type: ignore[misc]
# endregion
python/packages/azure-ai/agent_framework_azure_ai/_entra_id_authentication.py
+67
View File
@@ -0,0 +1,67 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from typing import Union
from agent_framework.exceptions import ChatClientInvalidAuthException
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
logger: logging.Logger = logging.getLogger(__name__)
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]
"""A callable that returns a bearer token string, either synchronously or asynchronously."""
AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
"""Union of Azure credential types.
Accepts:
- ``TokenCredential`` — synchronous Azure credential (e.g. ``DefaultAzureCredential()``)
- ``AsyncTokenCredential`` — asynchronous Azure credential (e.g. ``azure.identity.aio.DefaultAzureCredential()``)
"""
def resolve_credential_to_token_provider(
credential: AzureCredentialTypes | AzureTokenProvider,
token_endpoint: str | None,
) -> AzureTokenProvider:
"""Convert an Azure credential or token provider into an ``ad_token_provider`` callable.
If the credential is already a callable token provider, it is returned as-is
(``token_endpoint`` is not required in this case).
If it is a ``TokenCredential`` or ``AsyncTokenCredential``, it is wrapped using
``azure.identity.get_bearer_token_provider`` (sync or async variant) which
handles token caching and automatic refresh.
Args:
credential: An Azure credential or token provider callable.
token_endpoint: The token scope/endpoint
(e.g. ``"https://cognitiveservices.azure.com/.default"``).
Required when ``credential`` is a ``TokenCredential`` or ``AsyncTokenCredential``.
Returns:
A callable that returns a bearer token string (sync or async).
Raises:
ServiceInvalidAuthError: If the token endpoint is empty when needed for credential wrapping.
"""
# Already a token provider callable (not a credential object) — use directly
if callable(credential) and not isinstance(credential, (TokenCredential, AsyncTokenCredential)):
return credential
if not token_endpoint:
raise ChatClientInvalidAuthException(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
if isinstance(credential, AsyncTokenCredential):
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
return get_async_bearer_token_provider(credential, token_endpoint)
from azure.identity import get_bearer_token_provider
return get_bearer_token_provider(credential, token_endpoint) # type: ignore[arg-type]
python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
-261
View File
@@ -1,261 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Foundry Memory Context Provider using BaseContextProvider.
This module provides ``FoundryMemoryProvider``, built on
:class:`BaseContextProvider`.
"""
from __future__ import annotations
import logging
import sys
from contextlib import AbstractAsyncContextManager
from typing import TYPE_CHECKING, Any, ClassVar
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import load_settings
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from azure.ai.projects.aio import AIProjectClient
from openai.types.responses import ResponseInputItemParam
from ._shared import AzureAISettings
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
logger = logging.getLogger(__name__)
class FoundryMemoryProvider(BaseContextProvider):
"""Foundry Memory context provider using the new BaseContextProvider hooks pattern.
Integrates Azure AI Foundry Memory Store for persistent semantic memory,
searching and storing memories via the Azure AI Projects SDK.
Args:
source_id: Unique identifier for this provider instance.
project_client: Azure AI Project client for memory operations.
memory_store_name: The name of the memory store to use.
scope: The namespace that logically groups and isolates memories (e.g., user ID).
context_prompt: The prompt to prepend to retrieved memories.
update_delay: Timeout period before processing memory update in seconds.
Defaults to 300 (5 minutes). Set to 0 to immediately trigger updates.
"""
DEFAULT_SOURCE_ID: ClassVar[str] = "foundry_memory"
DEFAULT_CONTEXT_PROMPT = "## Memories\nConsider the following memories when answering user questions:"
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
*,
project_client: AIProjectClient | None = None,
project_endpoint: str | None = None,
credential: AzureCredentialTypes | None = None,
allow_preview: bool | None = None,
memory_store_name: str,
scope: str | None = None,
context_prompt: str | None = None,
update_delay: int = 300,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize the Foundry Memory context provider.
Args:
source_id: Unique identifier for this provider instance.
project_client: Azure AI Project client for memory operations.
project_endpoint: Azure AI project endpoint URL. Used when project_client is not provided.
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
Required when project_client is not provided.
allow_preview: Enables preview opt-in on internally-created ``AIProjectClient``.
memory_store_name: The name of the memory store to use.
scope: The namespace that logically groups and isolates memories (e.g., user ID).
If None, `session_id` will be used.
context_prompt: The prompt to prepend to retrieved memories.
update_delay: Timeout period before processing memory update in seconds.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
"""
super().__init__(source_id)
azure_ai_settings = load_settings(
AzureAISettings,
env_prefix="AZURE_AI_",
project_endpoint=project_endpoint,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if project_client is None:
resolved_endpoint = azure_ai_settings.get("project_endpoint")
if not resolved_endpoint:
raise ValueError(
"Azure AI project endpoint is required. Set via 'project_endpoint' parameter "
"or 'AZURE_AI_PROJECT_ENDPOINT' environment variable."
)
if not credential:
raise ValueError("Azure credential is required when project_client is not provided.")
project_client_kwargs: dict[str, Any] = {
"endpoint": resolved_endpoint,
"credential": credential, # type: ignore[arg-type]
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
project_client = AIProjectClient(**project_client_kwargs)
if not memory_store_name:
raise ValueError("memory_store_name is required")
if not scope:
raise ValueError("scope is required")
self.project_client = project_client
self.memory_store_name = memory_store_name
self.scope = scope
self.context_prompt = context_prompt or self.DEFAULT_CONTEXT_PROMPT
self.update_delay = update_delay
async def __aenter__(self) -> Self:
"""Async context manager entry."""
if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager):
await self.project_client.__aenter__()
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Async context manager exit."""
if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager):
await self.project_client.__aexit__(exc_type, exc_val, exc_tb)
# -- Hooks pattern ---------------------------------------------------------
async def before_run(
self,
*,
agent: SupportsAgentRun,
session: AgentSession,
context: SessionContext,
state: dict[str, Any],
) -> None:
"""Search Foundry Memory for relevant memories and add to the session context.
This method:
1. Retrieves static memories (user profile) on first call per session
2. Searches for contextual memories based on input messages
3. Combines and injects memories into the context
"""
# On first run, retrieve static memories (user profile memories)
if not state.get("initialized"):
try:
static_search_result = await self.project_client.beta.memory_stores.search_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
)
static_memories = [{"content": memory.memory_item.content} for memory in static_search_result.memories]
state["static_memories"] = static_memories
except Exception as e:
# Log but don't fail - memory retrieval is non-critical
logger.warning(f"Failed to retrieve static memories: {e}")
state["static_memories"] = []
finally:
# Mark as initialized regardless of success to avoid repeated attempts
state["initialized"] = True
# Search for contextual memories based on input messages
# Check if there are any non-empty input messages
has_input = any(msg and msg.text and msg.text.strip() for msg in context.input_messages)
if not has_input:
return
# Convert input messages to memory search item format
items: list[ResponseInputItemParam] = [
{"type": "message", "role": "user", "content": msg.text}
for msg in context.input_messages
if msg and msg.text and msg.text.strip()
]
try:
search_result = await self.project_client.beta.memory_stores.search_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
items=items,
previous_search_id=state.get("previous_search_id"),
)
# Extract search_id for next incremental search
if search_result.memories:
state["previous_search_id"] = search_result.search_id
# Combine static and contextual memories
contextual_memories = [{"content": memory.memory_item.content} for memory in search_result.memories]
all_memories = state.get("static_memories", []) + contextual_memories
# Inject memories into context
if all_memories:
line_separated_memories = "\n".join(
str(memory.get("content", "")) for memory in all_memories if memory.get("content")
)
if line_separated_memories:
context.extend_messages(
self.source_id,
[Message(role="user", text=f"{self.context_prompt}\n{line_separated_memories}")],
)
except Exception as e:
# Log but don't fail - memory retrieval is non-critical
logger.warning(f"Failed to search contextual memories: {e}")
async def after_run(
self,
*,
agent: SupportsAgentRun,
session: AgentSession,
context: SessionContext,
state: dict[str, Any],
) -> None:
"""Store request/response messages to Foundry Memory for future retrieval.
This method updates the memory store with conversation messages.
The update is debounced by the configured update_delay.
"""
messages_to_store: list[Message] = list(context.input_messages)
if context.response and context.response.messages:
messages_to_store.extend(context.response.messages)
# Filter and convert messages to memory update item format
items: list[ResponseInputItemParam] = []
for message in messages_to_store:
if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
if message.role == "user":
items.append({"role": "user", "type": "message", "content": message.text})
elif message.role == "assistant":
items.append({"role": "assistant", "type": "message", "content": message.text})
if not items:
return
try:
# Fire and forget - don't wait for the update to complete
update_poller = await self.project_client.beta.memory_stores.begin_update_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
items=items,
previous_update_id=state.get("previous_update_id"),
update_delay=self.update_delay,
)
# Store the update_id for next incremental update
state["previous_update_id"] = update_poller.update_id
except Exception as e:
# Log but don't fail - memory storage is non-critical
logger.warning(f"Failed to update memories: {e}")
__all__ = ["FoundryMemoryProvider"]
python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+10 -8
View File
@@ -18,7 +18,6 @@ from agent_framework import (
from agent_framework._mcp import MCPTool
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
AgentVersionDetails,
@@ -29,13 +28,15 @@ from azure.ai.projects.models import (
FunctionTool as AzureFunctionTool,
)
from ._client import AzureAIClient, AzureAIProjectAgentOptions
from ._client import AzureAIClient, AzureAIProjectAgentOptions # pyright: ignore[reportDeprecated]
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import AzureAISettings, create_text_format_config, from_azure_ai_tools, to_azure_ai_tools
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
from warnings import deprecated # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
from typing_extensions import TypeVar, deprecated # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self, TypedDict # type: ignore # pragma: no cover
else:
@@ -55,11 +56,12 @@ OptionsCoT = TypeVar(
)
@deprecated("AzureAIProjectAgentProvider is deprecated. Use FoundryAgent instead.")
class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
"""Provider for Azure AI Agent Service (Responses API).
"""Deprecated provider for Azure AI Agent Service (Responses API).
This provider allows you to create, retrieve, and manage Azure AI agents
using the AIProjectClient from the Azure AI Projects SDK.
This provider is deprecated. Use ``FoundryAgent`` instead to connect to
pre-configured agents in Foundry.
Examples:
Using with explicit AIProjectClient:
@@ -200,7 +202,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
)
# Extract options from default_options if present
opts = dict(default_options) if default_options else {}
opts: dict[str, Any] = dict(default_options) if default_options else {}
response_format = opts.get("response_format")
rai_config = opts.get("rai_config")
reasoning = opts.get("reasoning")
@@ -384,7 +386,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
if not isinstance(details.definition, PromptAgentDefinition):
raise ValueError("Agent definition must be PromptAgentDefinition to get a Agent.")
client = AzureAIClient(
client = AzureAIClient( # pyright: ignore[reportDeprecated]
project_client=self._project_client,
agent_name=details.name,
agent_version=details.version,
python/packages/azure-ai/pyproject.toml
+1
View File
@@ -24,6 +24,7 @@ classifiers = [
]
dependencies = [
"agent-framework-core>=1.0.0rc5",
"agent-framework-openai>=1.0.0rc5",
"azure-ai-agents>=1.2.0b5,<1.2.0b6",
"azure-ai-inference>=1.0.0b9,<1.0.0b10",
"aiohttp>=3.7.0,<4",
python/packages/azure-ai/tests/assets/sample_image.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 178 KiB

python/packages/azure-ai/tests/azure_openai/conftest.py
+61
View File
@@ -0,0 +1,61 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from agent_framework import Message
from pytest import fixture
# 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[Message]:
return []
python/packages/azure-ai/tests/azure_openai/test_azure_assistants_client.py
+409
View File
@@ -0,0 +1,409 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import (
SupportsChatGetResponse,
tool,
)
from agent_framework._settings import SecretString
from agent_framework.azure import AzureOpenAIAssistantsClient
from pydantic import Field
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."""
client = AzureOpenAIAssistantsClient(
deployment_name=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",
async_client=mock_async_azure_openai,
)
# Set the _should_delete_assistant flag directly if needed
if should_delete_assistant:
object.__setattr__(client, "_should_delete_assistant", True)
return client
@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."""
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 client.client is mock_async_azure_openai
assert client.model == "test_chat_deployment"
assert client.assistant_id == "existing-assistant-id"
assert client.thread_id == "test-thread-id"
assert not client._should_delete_assistant # type: ignore
assert isinstance(client, SupportsChatGetResponse)
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."""
client = AzureOpenAIAssistantsClient(
deployment_name=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
assistant_name="TestAssistant",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
async_client=mock_async_azure_openai,
)
assert client.client is mock_async_azure_openai
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert client.assistant_id is None
assert client.assistant_name == "TestAssistant"
assert not client._should_delete_assistant # type: ignore
def test_azure_assistants_client_init_validation_fail() -> None:
"""Test AzureOpenAIAssistantsClient initialization with validation failure."""
with pytest.raises(ValueError):
# 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(ValueError):
AzureOpenAIAssistantsClient(api_key=azure_openai_unit_test_env.get("AZURE_OPENAI_API_KEY", "test-key"))
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"}
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 client.model == "test_chat_deployment"
assert isinstance(client, SupportsChatGetResponse)
# 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 client.client.default_headers
assert 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."""
client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "existing-assistant-id"
assert not 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."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
)
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "test-assistant-id"
assert 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."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
)
await client.close() # type: ignore
# Verify assistant deletion was not called
mock_async_azure_openai.beta.assistants.delete.assert_not_called()
assert not 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."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
await client.close()
# Verify assistant deletion was called
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
assert not 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."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
# Test context manager
async with 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
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 = client.to_dict()
assert dumped_settings["model"] == "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 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"]
@tool(approval_mode="never_require")
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."
def test_azure_assistants_client_entra_id_authentication() -> None:
"""Test credential authentication path with sync credential."""
mock_credential = MagicMock()
mock_provider = MagicMock(return_value="token-string")
with (
patch("agent_framework_azure_ai._deprecated_azure_openai.load_settings") as mock_load_settings,
patch(
"agent_framework_azure_ai._deprecated_azure_openai.resolve_credential_to_token_provider",
return_value=mock_provider,
) as mock_resolve,
patch("agent_framework_azure_ai._deprecated_azure_openai.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": "https://cognitiveservices.azure.com/.default",
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_credential,
token_endpoint="https://cognitiveservices.azure.com/.default",
)
# Verify credential was resolved to a token provider
mock_resolve.assert_called_once_with(mock_credential, "https://cognitiveservices.azure.com/.default")
# Verify client was created with the token provider
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_provider
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_ai._deprecated_azure_openai.load_settings") as mock_load_settings:
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
# Test missing authentication raises error
with pytest.raises(ValueError, match="api_key, credential, or a client"):
AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
# No authentication provided at all
)
def test_azure_assistants_client_callable_credential() -> None:
"""Test callable token provider as credential."""
mock_provider = MagicMock(return_value="my-token")
with (
patch("agent_framework_azure_ai._deprecated_azure_openai.load_settings") as mock_load_settings,
patch(
"agent_framework_azure_ai._deprecated_azure_openai.resolve_credential_to_token_provider",
return_value=mock_provider,
),
patch("agent_framework_azure_ai._deprecated_azure_openai.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": "https://cognitiveservices.azure.com/.default",
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_provider,
token_endpoint="https://cognitiveservices.azure.com/.default",
)
# Verify client was created with the token provider
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_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_ai._deprecated_azure_openai.load_settings") as mock_load_settings,
patch("agent_framework_azure_ai._deprecated_azure_openai.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": SecretString("test-api-key"),
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": None,
"base_url": "https://custom-base-url.com",
}
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_ai._deprecated_azure_openai.load_settings") as mock_load_settings,
patch("agent_framework_azure_ai._deprecated_azure_openai.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": SecretString("test-api-key"),
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
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)
python/packages/azure-ai/tests/azure_openai/test_azure_chat_client.py
+1063
View File
File diff suppressed because it is too large Load Diff
python/packages/azure-ai/tests/azure_openai/test_azure_embedding_client.py
+158
View File
@@ -0,0 +1,158 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework.azure import AzureOpenAIEmbeddingClient
from agent_framework_openai import OpenAIEmbeddingOptions
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
def _make_openai_response(
embeddings: list[list[float]],
model: str = "text-embedding-3-small",
prompt_tokens: int = 5,
total_tokens: int = 5,
) -> CreateEmbeddingResponse:
"""Helper to create a mock OpenAI embeddings response."""
data = [OpenAIEmbedding(embedding=emb, index=i, object="embedding") for i, emb in enumerate(embeddings)]
return CreateEmbeddingResponse(
data=data,
model=model,
object="list",
usage=Usage(prompt_tokens=prompt_tokens, total_tokens=total_tokens),
)
@pytest.fixture
def azure_embedding_unit_test_env(monkeypatch: pytest.MonkeyPatch) -> None:
"""Clear ambient Azure OpenAI embedding env vars for deterministic unit tests."""
for key in (
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_TOKEN_ENDPOINT",
):
monkeypatch.delenv(key, raising=False)
def test_azure_construction_with_deployment_name(azure_embedding_unit_test_env: None) -> None:
client = AzureOpenAIEmbeddingClient(
deployment_name="text-embedding-3-small",
api_key="test-key",
endpoint="https://test.openai.azure.com/",
)
assert client.model == "text-embedding-3-small"
def test_azure_construction_with_existing_client(azure_embedding_unit_test_env: None) -> None:
mock_client = MagicMock()
client = AzureOpenAIEmbeddingClient(
deployment_name="my-deployment",
async_client=mock_client,
)
assert client.model == "my-deployment"
assert client.client is mock_client
def test_azure_construction_missing_deployment_name_raises(azure_embedding_unit_test_env: None) -> None:
with pytest.raises(ValueError, match="deployment name is required"):
AzureOpenAIEmbeddingClient(
api_key="test-key",
endpoint="https://test.openai.azure.com/",
)
def test_azure_construction_missing_credentials_raises(azure_embedding_unit_test_env: None) -> None:
with pytest.raises(ValueError, match="api_key, credential, or a client"):
AzureOpenAIEmbeddingClient(
deployment_name="test",
endpoint="https://test.openai.azure.com/",
)
async def test_azure_get_embeddings(azure_embedding_unit_test_env: None) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1, 0.2]],
)
mock_async_client = MagicMock()
mock_async_client.embeddings = MagicMock()
mock_async_client.embeddings.create = AsyncMock(return_value=mock_response)
client = AzureOpenAIEmbeddingClient(
deployment_name="text-embedding-3-small",
async_client=mock_async_client,
)
result = await client.get_embeddings(["hello"])
assert len(result) == 1
assert result[0].vector == [0.1, 0.2]
def test_azure_otel_provider_name(azure_embedding_unit_test_env: None) -> None:
mock_client = MagicMock()
client = AzureOpenAIEmbeddingClient(
deployment_name="test",
async_client=mock_client,
)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
not os.getenv("AZURE_OPENAI_ENDPOINT")
or (not os.getenv("AZURE_OPENAI_API_KEY") and not os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")),
reason="No Azure OpenAI credentials provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings() -> None:
"""End-to-end test of Azure OpenAI embedding generation."""
client = AzureOpenAIEmbeddingClient()
result = await client.get_embeddings(["hello world"])
assert len(result) == 1
assert isinstance(result[0].vector, list)
assert len(result[0].vector) > 0
assert all(isinstance(v, float) for v in result[0].vector)
assert result[0].model_id is not None
assert result.usage is not None
assert result.usage["input_token_count"] > 0
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings_multiple() -> None:
"""Test Azure OpenAI embedding generation for multiple inputs."""
client = AzureOpenAIEmbeddingClient()
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(e.vector) for e in result]
assert all(d == dims[0] for d in dims)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
"""Test Azure OpenAI embedding generation with custom dimensions."""
client = AzureOpenAIEmbeddingClient()
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 256
python/packages/azure-ai/tests/azure_openai/test_azure_responses_client.py
+729
View File
@@ -0,0 +1,729 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
import os
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import MagicMock
import pytest
from agent_framework import (
Agent,
AgentResponse,
ChatResponse,
Content,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework.azure import AzureOpenAIResponsesClient
from azure.identity import AzureCliCredential
from pydantic import BaseModel
from pytest import param
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests.",
)
logger = logging.getLogger(__name__)
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str
@tool(approval_mode="never_require")
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, Content]:
"""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, Content.from_hosted_vector_store(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(credential=AzureCliCredential())
assert azure_responses_client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"}) # type: ignore
def test_init_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
model_id = "test_model_id"
azure_responses_client = AzureOpenAIResponsesClient(deployment_name=model_id)
assert azure_responses_client.model == model_id
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test that model_id kwarg correctly sets the deployment name (issue #4299)."""
azure_responses_client = AzureOpenAIResponsesClient(model_id="gpt-4o")
assert azure_responses_client.model == "gpt-4o"
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg_does_not_override_deployment_name(
azure_openai_unit_test_env: dict[str, str],
) -> None:
"""Test that deployment_name takes precedence over model_id kwarg (issue #4299)."""
azure_responses_client = AzureOpenAIResponsesClient(deployment_name="my-deployment", model_id="gpt-4o")
assert azure_responses_client.model == "my-deployment"
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg_none(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test that model_id=None does not override the env-var deployment name."""
azure_responses_client = AzureOpenAIResponsesClient(model_id=None)
assert azure_responses_client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
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.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, SupportsChatGetResponse)
# 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(ValueError):
AzureOpenAIResponsesClient()
def test_init_with_project_client(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with an existing AIProjectClient."""
from unittest.mock import patch
from openai import AsyncOpenAI
# Create a mock AIProjectClient that returns a mock AsyncOpenAI client
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_client=mock_project_client,
deployment_name="gpt-4o",
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_with_project_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with a project endpoint and credential."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_endpoint="https://test-project.services.ai.azure.com",
deployment_name="gpt-4o",
credential=AzureCliCredential(),
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_create_client_from_project_with_project_client() -> None:
"""Test _create_client_from_project with an existing project client."""
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=mock_project_client,
project_endpoint=None,
credential=None,
)
assert result is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
def test_create_client_from_project_with_endpoint() -> None:
"""Test _create_client_from_project with a project endpoint."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_credential = MagicMock()
with patch("agent_framework_azure_ai._deprecated_azure_openai.AIProjectClient") as MockAIProjectClient:
mock_instance = MockAIProjectClient.return_value
mock_instance.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=mock_credential,
)
assert result is mock_openai_client
MockAIProjectClient.assert_called_once()
mock_instance.get_openai_client.assert_called_once()
def test_create_client_from_project_missing_endpoint() -> None:
"""Test _create_client_from_project raises error when endpoint is missing."""
with pytest.raises(ValueError, match="project endpoint is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint=None,
credential=MagicMock(),
)
def test_create_client_from_project_missing_credential() -> None:
"""Test _create_client_from_project raises error when credential is missing."""
with pytest.raises(ValueError, match="credential is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=None,
)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
settings = {
"deployment_name": 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["deployment_name"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert "api_key" not in dumped_settings
# 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"]
# region Integration Tests
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
# Simple ChatOptions - just verify they don't fail
param("temperature", 0.7, False, id="temperature"),
param("top_p", 0.9, False, id="top_p"),
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("frequency_penalty", 0.5, False, id="frequency_penalty"),
param("presence_penalty", 0.3, False, id="presence_penalty"),
param("stop", ["END"], False, id="stop"),
param("allow_multiple_tool_calls", True, False, id="allow_multiple_tool_calls"),
param("tool_choice", "none", True, id="tool_choice_none"),
# OpenAIResponsesOptions - just verify they don't fail
param("safety_identifier", "user-hash-abc123", False, id="safety_identifier"),
param("truncation", "auto", False, id="truncation"),
param("top_logprobs", 5, False, id="top_logprobs"),
param("prompt_cache_key", "test-cache-key", False, id="prompt_cache_key"),
param("max_tool_calls", 3, False, id="max_tool_calls"),
# Complex options requiring output validation
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param(
"tool_choice",
{"mode": "required", "required_function_name": "get_weather"},
True,
id="tool_choice_required",
),
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": [
"location",
"conditions",
"temperature_c",
"advisory",
],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
"""Parametrized test covering all ChatOptions and OpenAIResponsesOptions.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
# Prepare test message
if option_name == "tools" or option_name == "tool_choice":
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name == "response_format":
# Use prompt that works well with structured output
messages = [
Message(role="user", text="The weather in Seattle is sunny"),
Message(role="user", text="What is the weather in Seattle?"),
]
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name == "tool_choice":
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name == "tools" or option_name == "tool_choice":
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name == "response_format":
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_web_search() -> None:
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
for streaming in [False, True]:
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [AzureOpenAIResponsesClient.get_web_search_tool()],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [
AzureOpenAIResponsesClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})
],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
file_id, vector_store = await create_vector_store(azure_responses_client)
try:
# Test that the client will use the file search tool
response = await azure_responses_client.get_response(
messages=[
Message(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
options={
"tools": [
AzureOpenAIResponsesClient.get_file_search_tool(vector_store_ids=[vector_store.vector_store_id])
],
"tool_choice": "auto",
},
)
assert "sunny" in response.text.lower()
assert "75" in response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the file search tool
try:
response_stream = azure_responses_client.get_response(
messages=[
Message(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
stream=True,
options={
"tools": [
AzureOpenAIResponsesClient.get_file_search_tool(vector_store_ids=[vector_store.vector_store_id])
],
"tool_choice": "auto",
},
)
full_response = await response_stream.get_final_response()
assert "sunny" in full_response.text.lower()
assert "75" in full_response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_mcp_tool() -> None:
"""Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
messages=[Message(role="user", text="How to create an Azure storage account using az cli?")],
options={
# this needs to be high enough to handle the full MCP tool response.
"max_tokens": 5000,
"tools": AzureOpenAIResponsesClient.get_mcp_tool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
),
},
)
assert isinstance(response, ChatResponse)
# MCP server may return empty response intermittently - skip test rather than fail
if not response.text:
pytest.skip("MCP server returned empty response - service-side issue")
# 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"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with code interpreter tool."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
messages=[
Message(
role="user",
text="Calculate the sum of numbers from 1 to 10 using Python code.",
)
],
options={
"tools": [AzureOpenAIResponsesClient.get_code_interpreter_tool()],
},
)
# 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
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My hobby is photography. Remember this.", session=session, store=True)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved session
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_responses_client_tool_rich_content_image() -> None:
"""Test that Azure OpenAI Responses client can handle tool results containing images."""
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@tool(approval_mode="never_require")
def get_test_image() -> Content:
"""Return a test image for analysis."""
return Content.from_data(data=image_bytes, media_type="image/jpeg")
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
messages = [
Message(
role="user",
text="Call the get_test_image tool and describe what you see.",
)
]
options: dict[str, Any] = {"tools": [get_test_image], "tool_choice": "auto"}
if streaming:
response = await client.get_response(messages=messages, stream=True, options=options).get_final_response()
else:
response = await client.get_response(messages=messages, options=options)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
# sample_image.jpg contains a photo of a house; the model should mention it.
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
# region Integration with Foundry V2
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/")
or os.getenv("AZURE_AI_MODEL", "") == "",
reason="No real AZURE_AI_PROJECT_ENDPOINT or AZURE_AI_MODEL provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_function_call_roundtrip_preserves_fidelity():
"""Test that function calls roundtrip correctly with full fidelity preserved.
This verifies the changes where:
1. raw_representation is preserved when parsing function calls
2. fc_id and status are included in additional_properties
3. When re-sending messages, the full object fidelity is preserved
"""
call_count = 0
@tool(name="get_weather", approval_mode="never_require")
async def get_weather_tool(location: str) -> str:
"""Get weather for a location."""
nonlocal call_count
call_count += 1
return f"Weather in {location} is sunny, 72F"
client = AzureOpenAIResponsesClient(
project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
deployment_name=os.environ["AZURE_AI_MODEL"],
credential=AzureCliCredential(),
)
async with Agent(
client=client,
name="WeatherAgent",
instructions="You help check weather. Use get_weather when asked about weather.",
tools=[get_weather_tool],
default_options={"store": False}, # Store messages locally to test fidelity across messages
) as agent:
session = agent.create_session()
# First request - should invoke the tool
response1 = await agent.run("What is the weather in Seattle?", session=session)
assert response1 is not None
assert response1.text is not None
assert call_count >= 1
# Verify the response contains expected content
response_text = response1.text.lower()
assert "seattle" in response_text or "sunny" in response_text or "72" in response_text
# Second request - should work correctly with the preserved conversation
response2 = await agent.run("And how about in Portland?", session=session)
assert response2 is not None
assert response2.text is not None
assert call_count >= 2
# endregion
python/packages/azure-ai/tests/azure_openai/test_entra_id_authentication.py
+61
View File
@@ -0,0 +1,61 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import MagicMock, patch
import pytest
from agent_framework.exceptions import ChatClientInvalidAuthException
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from agent_framework_azure_ai._entra_id_authentication import (
resolve_credential_to_token_provider,
)
TOKEN_ENDPOINT = "https://cognitiveservices.azure.com/.default"
def test_resolve_sync_credential_returns_provider() -> None:
"""Test that a sync TokenCredential is resolved via azure.identity.get_bearer_token_provider."""
mock_credential = MagicMock(spec=TokenCredential)
mock_provider = MagicMock(return_value="token-string")
with patch("azure.identity.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
assert result is mock_provider
def test_resolve_async_credential_returns_provider() -> None:
"""Test that an AsyncTokenCredential is resolved via azure.identity.aio.get_bearer_token_provider."""
mock_credential = MagicMock(spec=AsyncTokenCredential)
mock_provider = MagicMock(return_value="token-string")
with patch("azure.identity.aio.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
assert result is mock_provider
def test_resolve_callable_provider_passthrough() -> None:
"""Test that a callable token provider is returned as-is, without needing token_endpoint."""
my_provider = lambda: "my-token" # noqa: E731
# Works with token_endpoint
assert resolve_credential_to_token_provider(my_provider, TOKEN_ENDPOINT) is my_provider
# Also works without token_endpoint
assert resolve_credential_to_token_provider(my_provider, None) is my_provider
assert resolve_credential_to_token_provider(my_provider, "") is my_provider
def test_resolve_missing_endpoint_raises() -> None:
"""Test that missing token endpoint raises ChatClientInvalidAuthException."""
mock_credential = MagicMock(spec=TokenCredential)
with pytest.raises(ChatClientInvalidAuthException, match="A token endpoint must be provided"):
resolve_credential_to_token_provider(mock_credential, "")
with pytest.raises(ChatClientInvalidAuthException, match="A token endpoint must be provided"):
resolve_credential_to_token_provider(mock_credential, None) # type: ignore[arg-type]
python/packages/azure-ai/tests/test_agent_provider.py
-80
View File
@@ -15,7 +15,6 @@ from azure.ai.agents.models import (
from azure.ai.agents.models import (
CodeInterpreterToolDefinition,
)
from azure.identity.aio import AzureCliCredential
from pydantic import BaseModel
from agent_framework_azure_ai import (
@@ -772,82 +771,3 @@ def test_from_azure_ai_agent_tools_unknown_dict() -> None:
# endregion
# region Integration Tests
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_create_agent() -> None:
"""Integration test: Create an agent using the provider."""
async with (
AzureCliCredential() as credential,
AzureAIAgentsProvider(credential=credential) as provider,
):
agent = await provider.create_agent(
name="IntegrationTestAgent",
instructions="You are a helpful assistant for testing.",
)
try:
assert isinstance(agent, Agent)
assert agent.name == "IntegrationTestAgent"
assert agent.id is not None
finally:
# Cleanup: delete the agent
if agent.id:
await provider._agents_client.delete_agent(agent.id) # type: ignore
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_get_agent() -> None:
"""Integration test: Get an existing agent using the provider."""
async with (
AzureCliCredential() as credential,
AzureAIAgentsProvider(credential=credential) as provider,
):
# First create an agent
created = await provider._agents_client.create_agent( # type: ignore
model=os.getenv("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o"),
name="GetAgentTest",
instructions="Test agent",
)
try:
# Then get it using the provider
agent = await provider.get_agent(created.id)
assert isinstance(agent, Agent)
assert agent.id == created.id
finally:
await provider._agents_client.delete_agent(created.id) # type: ignore
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_create_and_run() -> None:
"""Integration test: Create an agent and run a conversation."""
async with (
AzureCliCredential() as credential,
AzureAIAgentsProvider(credential=credential) as provider,
):
agent = await provider.create_agent(
name="RunTestAgent",
instructions="You are a helpful assistant. Always respond with 'Hello!' to any greeting.",
)
try:
result = await agent.run("Hi there!")
assert result is not None
assert len(result.messages) > 0
finally:
if agent.id:
await provider._agents_client.delete_agent(agent.id) # type: ignore
# endregion
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+9 -409
View File
@@ -1,17 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentSession,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
@@ -28,7 +22,6 @@ from azure.ai.agents.models import (
AgentsNamedToolChoiceType,
AgentsToolChoiceOptionMode,
CodeInterpreterToolDefinition,
FileInfo,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextFileCitationAnnotation,
@@ -41,19 +34,12 @@ from azure.ai.agents.models import (
SubmitToolApprovalAction,
SubmitToolOutputsAction,
ThreadRun,
VectorStore,
)
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from pydantic import BaseModel, Field
from agent_framework_azure_ai import AzureAIAgentClient, AzureAISettings
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/"),
reason="No real AZURE_AI_PROJECT_ENDPOINT provided; skipping integration tests.",
)
def create_test_azure_ai_chat_client(
mock_agents_client: MagicMock,
@@ -102,6 +88,15 @@ def create_test_azure_ai_chat_client(
return client
def test_init_emits_updated_deprecation_warning(mock_agents_client: MagicMock) -> None:
"""Test that construction emits the updated class deprecation warning."""
with pytest.deprecated_call(match="V1 Agents Service API and has no direct replacement"):
AzureAIAgentClient(
agents_client=mock_agents_client,
agent_id="test-agent",
)
def test_azure_ai_settings_init(azure_ai_unit_test_env: dict[str, str]) -> None:
"""Test AzureAISettings initialization."""
settings = load_settings(AzureAISettings, env_prefix="AZURE_AI_")
@@ -1527,401 +1522,6 @@ def get_weather(
return f"The weather in {location} is sunny with a high of 25°C."
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_get_response() -> None:
"""Test Azure AI Chat Client response."""
async with AzureAIAgentClient(credential=AzureCliCredential()) as azure_ai_chat_client:
assert isinstance(azure_ai_chat_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(
Message(
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(Message(role="user", text="What's the weather like today?"))
# Test that the agents_client can be used to get a response
response = await azure_ai_chat_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"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_get_response_tools() -> None:
"""Test Azure AI Chat Client response with tools."""
async with AzureAIAgentClient(credential=AzureCliCredential()) as azure_ai_chat_client:
assert isinstance(azure_ai_chat_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the agents_client can be used to get a response
response = await azure_ai_chat_client.get_response(
messages=messages,
options={"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"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_streaming() -> None:
"""Test Azure AI Chat Client streaming response."""
async with AzureAIAgentClient(credential=AzureCliCredential()) as azure_ai_chat_client:
assert isinstance(azure_ai_chat_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(
Message(
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(Message(role="user", text="What's the weather like today?"))
# Test that the agents_client can be used to get a response
response = azure_ai_chat_client.get_response(messages=messages, stream=True)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if content.type == "text" and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_streaming_tools() -> None:
"""Test Azure AI Chat Client streaming response with tools."""
async with AzureAIAgentClient(credential=AzureCliCredential()) as azure_ai_chat_client:
assert isinstance(azure_ai_chat_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the agents_client can be used to get a response
response = azure_ai_chat_client.get_response(
messages=messages,
stream=True,
options={"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 content.type == "text" and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_basic_run() -> None:
"""Test Agent basic run functionality with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
assert "Hello World" in response.text
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_basic_run_streaming() -> None:
"""Test Agent basic streaming functionality with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run("Please respond with exactly: 'This is a streaming response test.'", stream=True):
assert chunk is not None
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
full_message += chunk.text
# Validate streaming response
assert len(full_message) > 0
assert "streaming response test" in full_message.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_thread_persistence() -> None:
"""Test Agent session persistence across runs with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new session that will be reused
session = agent.create_session()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", session=session
)
assert isinstance(first_response, AgentResponse)
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?", session=session
)
assert isinstance(second_response, AgentResponse)
assert "42" in second_response.text
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_existing_thread_id() -> None:
"""Test Agent existing thread ID functionality with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and get the session ID
session = first_agent.create_session()
first_response = await first_agent.run("My name is Alice. Remember this.", session=session)
# Validate first response
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# The thread ID is set after the first response
existing_thread_id = session.service_session_id
assert existing_thread_id is not None
# Now continue with the same thread ID in a new agent instance
async with Agent(
client=AzureAIAgentClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_thread_id)
# Ask about the previous conversation
response2 = await second_agent.run("What is my name?", session=session)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentResponse)
assert response2.text is not None
# Should reference Alice from the previous conversation
assert "alice" in response2.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_code_interpreter():
"""Test Agent with code interpreter through AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[AzureAIAgentClient.get_code_interpreter_tool()],
) 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, AgentResponse)
assert response.text is not None
# Factorial of 5 is 120
assert "120" in response.text or "factorial" in response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_file_search():
"""Test Agent with file search through AzureAIAgentClient."""
client = AzureAIAgentClient(credential=AzureCliCredential())
file: FileInfo | None = None
vector_store: VectorStore | None = None
try:
# 1. Read and upload the test file to the Azure AI agent service
test_file_path = Path(__file__).parent / "resources" / "employees.pdf"
file = await client.agents_client.files.upload_and_poll(file_path=str(test_file_path), purpose="assistants")
vector_store = await client.agents_client.vector_stores.create_and_poll(
file_ids=[file.id], name="test_employees_vectorstore"
)
# 2. Create file search tool with uploaded resources
file_search_tool = AzureAIAgentClient.get_file_search_tool(vector_store_ids=[vector_store.id])
async with Agent(
client=client,
instructions="You are a helpful assistant that can search through uploaded employee files.",
tools=[file_search_tool],
) as agent:
# 3. Test file search functionality
response = await agent.run("Who is the youngest employee in the files?")
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
# Should find information about Alice Johnson (age 24) being the youngest
assert any(term in response.text.lower() for term in ["alice", "johnson", "24"])
finally:
# 4. Cleanup: Delete the vector store and file
try:
if vector_store:
await client.agents_client.vector_stores.delete(vector_store.id)
if file:
await client.agents_client.files.delete(file.id)
except Exception:
# Ignore cleanup errors to avoid masking the actual test failure
pass
finally:
await client.close()
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_hosted_mcp_tool() -> None:
"""Integration test for MCP tool with Azure AI Agent using Microsoft Learn MCP."""
mcp_tool = AzureAIAgentClient.get_mcp_tool(
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 Agent(
client=AzureAIAgentClient(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?",
options={"max_tokens": 200},
)
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
# With never_require approval mode, there should be no approval requests
assert len(response.user_input_requests) == 0, (
f"Expected no approval requests with never_require mode, but got {len(response.user_input_requests)}"
)
# 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"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather],
) 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, AgentResponse)
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", "25"])
# 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, AgentResponse)
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", "25"])
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_chat_options_run_level() -> None:
"""Test ChatOptions parameter coverage at run level."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
tools=[get_weather],
options={
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9,
"tool_choice": "auto",
"metadata": {"test": "value"},
},
)
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_chat_options_agent_level() -> None:
"""Test ChatOptions parameter coverage agent level."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
tools=[get_weather],
default_options={
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9,
"tool_choice": "auto",
"metadata": {"test": "value"},
},
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
)
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
async def test_azure_ai_chat_client_cleanup_agent_when_enabled_and_created(
mock_agents_client: MagicMock,
) -> None:
python/packages/azure-ai/tests/test_azure_ai_client.py
+18 -385
View File
@@ -11,8 +11,6 @@ from uuid import uuid4
import pytest
from agent_framework import (
Agent,
AgentResponse,
Annotation,
ChatOptions,
ChatResponse,
@@ -24,7 +22,7 @@ from agent_framework import (
tool,
)
from agent_framework._settings import load_settings
from agent_framework.openai._responses_client import RawOpenAIResponsesClient
from agent_framework_openai._chat_client import RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
ApproximateLocation,
@@ -41,17 +39,11 @@ from azure.identity.aio import AzureCliCredential
from openai.types.responses.parsed_response import ParsedResponse
from openai.types.responses.response import Response as OpenAIResponse
from pydantic import BaseModel, ConfigDict, Field
from pytest import fixture, param
from pytest import fixture
from agent_framework_azure_ai import AzureAIClient, AzureAISettings
from agent_framework_azure_ai._shared import from_azure_ai_tools
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/")
or os.getenv("AZURE_AI_MODEL_DEPLOYMENT_NAME", "") == "",
reason="No real AZURE_AI_PROJECT_ENDPOINT or AZURE_AI_MODEL_DEPLOYMENT_NAME provided; skipping integration tests.",
)
@pytest.fixture
def mock_project_client() -> MagicMock:
@@ -415,7 +407,7 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
@@ -452,7 +444,7 @@ async def test_prepare_options_with_application_endpoint(
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
@@ -494,7 +486,7 @@ async def test_prepare_options_with_application_project_client(
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model"},
),
patch.object(
@@ -512,19 +504,6 @@ async def test_prepare_options_with_application_project_client(
assert "extra_body" not in run_options
async def test_initialize_client(mock_project_client: MagicMock) -> None:
"""Test _initialize_client method."""
client = create_test_azure_ai_client(mock_project_client)
mock_openai_client = MagicMock()
mock_project_client.get_openai_client = MagicMock(return_value=mock_openai_client)
await client._initialize_client()
assert client.client is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
def test_update_agent_name_and_description(mock_project_client: MagicMock) -> None:
"""Test _update_agent_name_and_description method."""
client = create_test_azure_ai_client(mock_project_client)
@@ -827,14 +806,14 @@ async def test_runtime_tools_override_logs_warning(
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
with patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]},
):
await client._prepare_options(messages, {})
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_two"}]},
),
patch("agent_framework_azure_ai._client.logger.warning") as mock_warning,
@@ -853,7 +832,7 @@ async def test_prepare_options_logs_warning_for_tools_with_existing_agent_versio
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]},
),
patch("agent_framework_azure_ai._client.logger.warning") as mock_warning,
@@ -875,7 +854,7 @@ async def test_prepare_options_logs_warning_for_tools_on_application_endpoint(
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]},
),
patch.object(client, "_get_agent_reference_or_create", new_callable=AsyncMock) as mock_get_agent_reference,
@@ -1101,14 +1080,14 @@ async def test_runtime_structured_output_override_logs_warning(
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
with patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model"},
):
await client._prepare_options(messages, {"response_format": ResponseFormatModel})
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={"model": "test-model"},
),
patch("agent_framework_azure_ai._client.logger.warning") as mock_warning,
@@ -1129,7 +1108,7 @@ async def test_prepare_options_excludes_response_format(
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={
"model": "test-model",
"response_format": ResponseFormatModel,
@@ -1164,7 +1143,7 @@ async def test_prepare_options_keeps_values_for_unsupported_option_keys(
with (
patch(
"agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options",
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
return_value={
"model": "test-model",
"tools": [{"type": "function", "name": "weather"}],
@@ -1365,352 +1344,6 @@ async def client() -> AsyncGenerator[AzureAIClient, None]:
await project_client.agents.delete(agent_name=agent_name)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
# Simple ChatOptions - just verify they don't fail
param("top_p", 0.9, False, id="top_p"),
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("frequency_penalty", 0.5, False, id="frequency_penalty"),
param("presence_penalty", 0.3, False, id="presence_penalty"),
param("stop", ["END"], False, id="stop"),
param("allow_multiple_tool_calls", True, False, id="allow_multiple_tool_calls"),
param("tool_choice", "none", True, id="tool_choice_none"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param("tool_choice", "required", True, id="tool_choice_required_any"),
param(
"tool_choice",
{"mode": "required", "required_function_name": "get_weather"},
True,
id="tool_choice_required",
),
# OpenAIResponsesOptions - just verify they don't fail
param("safety_identifier", "user-hash-abc123", False, id="safety_identifier"),
param("truncation", "auto", False, id="truncation"),
param("top_logprobs", 5, False, id="top_logprobs"),
param("prompt_cache_key", "test-cache-key", False, id="prompt_cache_key"),
param("max_tool_calls", 3, False, id="max_tool_calls"),
],
)
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
client: AzureAIClient,
) -> None:
"""Parametrized test covering options that can be set at runtime for a Foundry Agent.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
This test reuses a single agent.
"""
# Prepare test message
if option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value, "tools": [get_weather]}
for streaming in [False, True]:
if streaming:
# Test streaming mode
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
# For tool_choice="required", we return after tool execution without a model text response
is_required_tool_choice = option_name == "tool_choice" and (
option_value == "required" or (isinstance(option_value, dict) and option_value.get("mode") == "required")
)
if is_required_tool_choice:
# Response should have function call and function result, but no text from model
assert len(response.messages) >= 2, f"Expected function call + result for {option_name}"
has_function_call = any(c.type == "function_call" for msg in response.messages for c in msg.contents)
has_function_result = any(c.type == "function_result" for msg in response.messages for c in msg.contents)
assert has_function_call, f"No function call in response for {option_name}"
assert has_function_result, f"No function result in response for {option_name}"
else:
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name.startswith("tool_choice") and not is_required_tool_choice:
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name == "response_format":
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
param("temperature", 0.7, False, id="temperature"),
# Complex options requiring output validation
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": ["location", "conditions", "temperature_c", "advisory"],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
async def test_integration_agent_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
"""Test Foundry agent level options in both streaming and non-streaming modes.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
This test create a new client and uses it for both streaming and non-streaming tests.
"""
async with temporary_chat_client(agent_name=f"test-agent-{option_name.replace('_', '-')}-{uuid4()}") as client:
for streaming in [False, True]:
# Prepare test message
if option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options = {option_name: option_value}
if streaming:
# Test streaming mode
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation and option_name.startswith("response_format"):
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_web_search() -> None:
async with temporary_chat_client(agent_name="af-int-test-web-search") as client:
for streaming in [False, True]:
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [client.get_web_search_tool()],
},
}
if streaming:
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
content = {
"messages": [
Message(role="user", text="What is the current weather? Do not ask for my current location.")
],
"options": {
"tool_choice": "auto",
"tools": [client.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})],
},
}
if streaming:
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_agent_hosted_mcp_tool() -> None:
"""Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP."""
async with temporary_chat_client(agent_name="af-int-test-mcp") as client:
response = await client.get_response(
messages=[Message(role="user", text="How to create an Azure storage account using az cli?")],
options={
# this needs to be high enough to handle the full MCP tool response.
"max_tokens": 5000,
"tools": client.get_mcp_tool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
),
},
)
assert isinstance(response, ChatResponse)
assert response.text
# 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"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with code interpreter tool through AzureAIClient."""
async with temporary_chat_client(agent_name="af-int-test-code-interpreter") as client:
response = await client.get_response(
messages=[Message(role="user", text="Calculate the sum of numbers from 1 to 10 using Python code.")],
options={
"tools": [client.get_code_interpreter_tool()],
},
)
# 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
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with (
temporary_chat_client(agent_name="af-int-test-existing-session") as client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as first_agent,
):
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My hobby is photography. Remember this.", session=session, store=True)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with (
temporary_chat_client(agent_name="af-int-test-existing-session-2") as client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as second_agent,
):
# Reuse the preserved session
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
# region Factory Method Tests
@@ -2031,7 +1664,7 @@ async def test_inner_get_response_enriches_non_streaming(mock_project_client: Ma
async def _fake_awaitable() -> ChatResponse:
return base_response
with patch.object(RawOpenAIResponsesClient, "_inner_get_response", return_value=_fake_awaitable()):
with patch.object(RawOpenAIChatClient, "_inner_get_response", return_value=_fake_awaitable()):
result_awaitable = client._inner_get_response(messages=[], options={}, stream=False)
result = await result_awaitable # type: ignore[misc]
@@ -2054,7 +1687,7 @@ async def test_inner_get_response_no_search_output_non_streaming(mock_project_cl
async def _fake_awaitable() -> ChatResponse:
return base_response
with patch.object(RawOpenAIResponsesClient, "_inner_get_response", return_value=_fake_awaitable()):
with patch.object(RawOpenAIChatClient, "_inner_get_response", return_value=_fake_awaitable()):
result_awaitable = client._inner_get_response(messages=[], options={}, stream=False)
result = await result_awaitable # type: ignore[misc]
@@ -2075,7 +1708,7 @@ def test_inner_get_response_streaming_registers_hook(mock_project_client: MagicM
mock_stream = _create_mock_stream()
with patch.object(RawOpenAIResponsesClient, "_inner_get_response", return_value=mock_stream):
with patch.object(RawOpenAIChatClient, "_inner_get_response", return_value=mock_stream):
result = client._inner_get_response(messages=[], options={}, stream=True)
assert result is mock_stream
@@ -2088,7 +1721,7 @@ def test_streaming_hook_captures_search_urls(mock_project_client: MagicMock) ->
mock_stream = _create_mock_stream()
with patch.object(RawOpenAIResponsesClient, "_inner_get_response", return_value=mock_stream):
with patch.object(RawOpenAIChatClient, "_inner_get_response", return_value=mock_stream):
client._inner_get_response(messages=[], options={}, stream=True)
hook = mock_stream._transform_hooks[0]
@@ -2116,7 +1749,7 @@ def test_streaming_hook_enriches_url_citation(mock_project_client: MagicMock) ->
mock_stream = _create_mock_stream()
with patch.object(RawOpenAIResponsesClient, "_inner_get_response", return_value=mock_stream):
with patch.object(RawOpenAIChatClient, "_inner_get_response", return_value=mock_stream):
client._inner_get_response(messages=[], options={}, stream=True)
hook = mock_stream._transform_hooks[0]
python/packages/azure-ai/tests/test_foundry_memory_provider.py
-507
View File
@@ -1,507 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework_azure_ai._foundry_memory_provider import FoundryMemoryProvider
@pytest.fixture
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
mock_client.beta = AsyncMock()
mock_client.beta.memory_stores = AsyncMock()
mock_client.beta.memory_stores.search_memories = AsyncMock()
mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@pytest.fixture
def mock_credential() -> Mock:
"""Create a mock Azure credential."""
return Mock()
# -- Initialization tests ------------------------------------------------------
class TestInit:
"""Test FoundryMemoryProvider initialization."""
def test_init_with_all_params(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
source_id="custom_source",
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
context_prompt="Custom prompt",
update_delay=60,
)
assert provider.source_id == "custom_source"
assert provider.project_client is mock_project_client
assert provider.memory_store_name == "test_store"
assert provider.scope == "user_123"
assert provider.context_prompt == "Custom prompt"
assert provider.update_delay == 60
def test_init_default_source_id(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.source_id == FoundryMemoryProvider.DEFAULT_SOURCE_ID
def test_init_default_context_prompt(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.context_prompt == FoundryMemoryProvider.DEFAULT_CONTEXT_PROMPT
def test_init_default_update_delay(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.update_delay == 300
def test_init_with_project_endpoint_and_credential(
self, mock_project_client: AsyncMock, mock_credential: Mock
) -> None:
with patch("agent_framework_azure_ai._foundry_memory_provider.AIProjectClient") as mock_ai_project_client:
mock_ai_project_client.return_value = mock_project_client
provider = FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
credential=mock_credential, # type: ignore[arg-type]
allow_preview=True,
memory_store_name="test_store",
scope="user_123",
)
assert provider.project_client is mock_project_client
mock_ai_project_client.assert_called_once_with(
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
def test_init_requires_project_endpoint_without_project_client(self) -> None:
with (
patch("agent_framework_azure_ai._foundry_memory_provider.load_settings") as mock_load_settings,
patch.dict(os.environ, {}, clear=True),
pytest.raises(ValueError, match="project endpoint is required"),
):
mock_load_settings.return_value = {"project_endpoint": None}
FoundryMemoryProvider(
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_credential_without_project_client(self) -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_memory_store_name(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="memory_store_name is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="",
scope="user_123",
)
def test_init_requires_scope(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="scope is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="",
)
# -- before_run tests ----------------------------------------------------------
class TestBeforeRun:
"""Test before_run hook."""
async def test_retrieves_static_memories_on_first_run(self, mock_project_client: AsyncMock) -> None:
"""First call retrieves static (user profile) memories."""
mem1 = Mock()
mem1.memory_item.content = "User prefers Python"
mem2 = Mock()
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should call search_memories twice: once for static, once for contextual
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
async def test_contextual_memories_added_to_context(self, mock_project_client: AsyncMock) -> None:
"""Contextual search returns memories → messages added to context with prompt."""
# Mock static search (first call)
static_mem = Mock()
static_mem.memory_item.content = "User prefers Python"
static_result = Mock()
static_result.memories = [static_mem]
# Mock contextual search (second call)
contextual_mem = Mock()
contextual_mem.memory_item.content = "Last discussed async patterns"
contextual_result = Mock()
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Check that memories were added to context
assert provider.source_id in ctx.context_messages
added = ctx.context_messages[provider.source_id]
assert len(added) == 1
assert "User prefers Python" in added[0].text # type: ignore[operator]
assert "Last discussed async patterns" in added[0].text # type: ignore[operator]
assert provider.context_prompt in added[0].text # type: ignore[operator]
assert session.state[provider.source_id]["previous_search_id"] == "search-123"
async def test_empty_input_skips_contextual_search(self, mock_project_client: AsyncMock) -> None:
"""Empty input messages → only static search performed, no contextual search."""
static_result = Mock()
static_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should only call search_memories once for static memories
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
"""Empty search results → no messages added."""
mock_search_result = Mock()
mock_search_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="test")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert provider.source_id not in ctx.context_messages
async def test_static_memories_only_retrieved_once(self, mock_project_client: AsyncMock) -> None:
"""Static memories are only retrieved on the first call."""
static_mem = Mock()
static_mem.memory_item.content = "Static memory"
static_result = Mock()
static_result.memories = [static_mem]
contextual_result = Mock()
contextual_result.memories = []
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# First call
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Search exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# Should not raise exception
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# No memories added
assert provider.source_id not in ctx.context_messages
# -- after_run tests -----------------------------------------------------------
class TestAfterRun:
"""Test after_run hook."""
async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -> None:
"""Stores input+response messages via begin_update_memories."""
mock_poller = Mock()
mock_poller.update_id = "update-456"
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="question")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="answer")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
assert call_kwargs["items"][0]["content"] == "question"
assert call_kwargs["items"][1]["content"] == "answer"
assert session.state[provider.source_id]["previous_update_id"] == "update-456"
async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
"""Only stores user/assistant/system messages with text."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text="hello"),
Message(role="tool", text="tool output"),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[Message(role="assistant", text="reply")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
assert items[1]["content"] == "reply"
async def test_skips_empty_messages(self, mock_project_client: AsyncMock) -> None:
"""Skips messages with empty text."""
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text=""),
Message(role="user", text=" "),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
"""Uses the configured update_delay parameter."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
update_delay=60,
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
"""Uses previous_update_id for incremental updates."""
mock_poller1 = Mock()
mock_poller1.update_id = "update-1"
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx1 = SessionContext(input_messages=[Message(role="user", text="first")], session_id="s1")
ctx1._response = AgentResponse(messages=[Message(role="assistant", text="response1")])
# First update
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx1, state=session.state.setdefault(provider.source_id, {})
)
assert session.state[provider.source_id]["previous_update_id"] == "update-1"
# Second update should use previous_update_id
ctx2 = SessionContext(input_messages=[Message(role="user", text="second")], session_id="s1")
ctx2._response = AgentResponse(messages=[Message(role="assistant", text="response2")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Update exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
# Should not raise exception
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# -- Context manager tests -----------------------------------------------------
class TestContextManager:
"""Test __aenter__/__aexit__ delegation."""
async def test_aenter_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
result = await provider.__aenter__()
assert result is provider
mock_project_client.__aenter__.assert_awaited_once()
async def test_aexit_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
await provider.__aexit__(None, None, None)
mock_project_client.__aexit__.assert_awaited_once()
async def test_async_with_syntax(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
async with provider as p:
assert p is provider
python/packages/azure-ai/tests/test_provider.py
-48
View File
@@ -1,12 +1,10 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import Agent, FunctionTool
from agent_framework._mcp import MCPTool
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
AgentVersionDetails,
PromptAgentDefinition,
@@ -14,16 +12,9 @@ from azure.ai.projects.models import (
from azure.ai.projects.models import (
FunctionTool as AzureFunctionTool,
)
from azure.identity.aio import AzureCliCredential
from agent_framework_azure_ai import AzureAIProjectAgentProvider
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/")
or os.getenv("AZURE_AI_MODEL_DEPLOYMENT_NAME", "") == "",
reason="No real AZURE_AI_PROJECT_ENDPOINT or AZURE_AI_MODEL_DEPLOYMENT_NAME provided; skipping integration tests.",
)
@pytest.fixture
def mock_project_client() -> MagicMock:
@@ -689,42 +680,3 @@ async def test_provider_create_agent_with_mcp_and_regular_tools(
assert "regular_function" in tool_names
assert "mcp_function_1" in tool_names
assert "mcp_function_2" in tool_names
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_provider_create_and_get_agent_integration() -> None:
"""Integration test for provider create_agent and get_agent."""
endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"]
model = os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"]
async with (
AzureCliCredential() as credential,
AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
):
provider = AzureAIProjectAgentProvider(project_client=project_client)
try:
# Create agent
agent = await provider.create_agent(
name="ProviderTestAgent",
model=model,
instructions="You are a helpful assistant. Always respond with 'Hello from provider!'",
)
assert isinstance(agent, Agent)
assert agent.name == "ProviderTestAgent"
# Run the agent
response = await agent.run("Hi!")
assert response.text is not None
assert len(response.text) > 0
# Get the same agent
retrieved_agent = await provider.get_agent(name="ProviderTestAgent")
assert retrieved_agent.name == "ProviderTestAgent"
finally:
# Cleanup
await project_client.agents.delete(agent_name="ProviderTestAgent")