mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Added changes (#1909)
This commit is contained in:
Dmytro Struk
committed by
GitHub
Unverified
GitHub
parent
8b4aa1ebb5
commit
39d3111734
@@ -3,6 +3,7 @@
|
||||
import importlib.metadata
|
||||
|
||||
from ._chat_client import AzureAIAgentClient, AzureAISettings
|
||||
from ._chat_client_v2 import AzureAIAgentClientV2
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
@@ -11,6 +12,7 @@ except importlib.metadata.PackageNotFoundError:
|
||||
|
||||
__all__ = [
|
||||
"AzureAIAgentClient",
|
||||
"AzureAIAgentClientV2",
|
||||
"AzureAISettings",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
@@ -40,9 +40,9 @@ from agent_framework import (
|
||||
use_chat_middleware,
|
||||
use_function_invocation,
|
||||
)
|
||||
from agent_framework._pydantic import AFBaseSettings
|
||||
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
|
||||
from agent_framework.observability import use_observability
|
||||
from azure.ai.agents.aio import AgentsClient
|
||||
from azure.ai.agents.models import (
|
||||
Agent,
|
||||
AgentsNamedToolChoice,
|
||||
@@ -85,11 +85,11 @@ from azure.ai.agents.models import (
|
||||
ToolDefinition,
|
||||
ToolOutput,
|
||||
)
|
||||
from azure.ai.projects.aio import AIProjectClient
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
|
||||
from pydantic import ValidationError
|
||||
|
||||
from ._shared import AzureAISettings
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
@@ -99,47 +99,6 @@ else:
|
||||
logger = get_logger("agent_framework.azure")
|
||||
|
||||
|
||||
class AzureAISettings(AFBaseSettings):
|
||||
"""Azure AI Project settings.
|
||||
|
||||
The settings are first loaded from environment variables with the prefix 'AZURE_AI_'.
|
||||
If the environment variables are not found, the settings can be loaded from a .env file
|
||||
with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
|
||||
are ignored; however, validation will fail alerting that the settings are missing.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Azure AI Project endpoint URL.
|
||||
Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
|
||||
model_deployment_name: The name of the model deployment to use.
|
||||
Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
|
||||
env_file_path: If provided, the .env settings are read from this file path location.
|
||||
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework_azure_ai import AzureAISettings
|
||||
|
||||
# Using environment variables
|
||||
# Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
|
||||
# Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
|
||||
settings = AzureAISettings()
|
||||
|
||||
# Or passing parameters directly
|
||||
settings = AzureAISettings(
|
||||
project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4"
|
||||
)
|
||||
|
||||
# Or loading from a .env file
|
||||
settings = AzureAISettings(env_file_path="path/to/.env")
|
||||
"""
|
||||
|
||||
env_prefix: ClassVar[str] = "AZURE_AI_"
|
||||
|
||||
project_endpoint: str | None = None
|
||||
model_deployment_name: str | None = None
|
||||
|
||||
|
||||
TAzureAIAgentClient = TypeVar("TAzureAIAgentClient", bound="AzureAIAgentClient")
|
||||
|
||||
|
||||
@@ -154,7 +113,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project_client: AIProjectClient | None = None,
|
||||
agents_client: AgentsClient | None = None,
|
||||
agent_id: str | None = None,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
@@ -168,16 +127,16 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
"""Initialize an Azure AI Agent client.
|
||||
|
||||
Keyword Args:
|
||||
project_client: An existing AIProjectClient to use. If not provided, one will be created.
|
||||
agent_id: The ID of an existing agent to use. If not provided and project_client is provided,
|
||||
a new agent will be created (and deleted after the request). If neither project_client
|
||||
agents_client: An existing AgentsClient to use. If not provided, one will be created.
|
||||
agent_id: The ID of an existing agent to use. If not provided and agents_client is provided,
|
||||
a new agent will be created (and deleted after the request). If neither agents_client
|
||||
nor agent_id is provided, both will be created and managed automatically.
|
||||
agent_name: The name to use when creating new agents.
|
||||
thread_id: Default thread ID to use for conversations. Can be overridden by
|
||||
conversation_id property when making a request.
|
||||
project_endpoint: The Azure AI Project endpoint URL.
|
||||
Can also be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
|
||||
Ignored when a project_client is passed.
|
||||
Ignored when a agents_client is passed.
|
||||
model_deployment_name: The model deployment name to use for agent creation.
|
||||
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
|
||||
async_credential: Azure async credential to use for authentication.
|
||||
@@ -217,9 +176,9 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to create Azure AI settings.", ex) from ex
|
||||
|
||||
# If no project_client is provided, create one
|
||||
# If no agents_client is provided, create one
|
||||
should_close_client = False
|
||||
if project_client is None:
|
||||
if agents_client is None:
|
||||
if not azure_ai_settings.project_endpoint:
|
||||
raise ServiceInitializationError(
|
||||
"Azure AI project endpoint is required. Set via 'project_endpoint' parameter "
|
||||
@@ -234,10 +193,11 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
|
||||
# Use provided credential
|
||||
if not async_credential:
|
||||
raise ServiceInitializationError("Azure credential is required when project_client is not provided.")
|
||||
project_client = AIProjectClient(
|
||||
raise ServiceInitializationError("Azure credential is required when agents_client is not provided.")
|
||||
agents_client = AgentsClient(
|
||||
endpoint=azure_ai_settings.project_endpoint,
|
||||
credential=async_credential,
|
||||
# TODO (dmytrostruk): Verify if user_agent works with AgentsClient
|
||||
user_agent=AGENT_FRAMEWORK_USER_AGENT,
|
||||
)
|
||||
should_close_client = True
|
||||
@@ -246,7 +206,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Initialize instance variables
|
||||
self.project_client = project_client
|
||||
self.agents_client = agents_client
|
||||
self.credential = async_credential
|
||||
self.agent_id = agent_id
|
||||
self.agent_name = agent_name
|
||||
@@ -256,27 +216,6 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
self._should_close_client = should_close_client # Track whether we should close client connection
|
||||
self._agent_definition: Agent | None = None # Cached definition for existing agent
|
||||
|
||||
async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None:
|
||||
"""Use this method to setup tracing in your Azure AI Project.
|
||||
|
||||
This will take the connection string from the project project_client.
|
||||
It will override any connection string that is set in the environment variables.
|
||||
It will disable any OTLP endpoint that might have been set.
|
||||
"""
|
||||
try:
|
||||
conn_string = await self.project_client.telemetry.get_application_insights_connection_string()
|
||||
except ResourceNotFoundError:
|
||||
logger.warning(
|
||||
"No Application Insights connection string found for the Azure AI Project, "
|
||||
"please call setup_observability() manually."
|
||||
)
|
||||
return
|
||||
from agent_framework.observability import setup_observability
|
||||
|
||||
setup_observability(
|
||||
applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "Self":
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
@@ -286,7 +225,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the project_client and clean up any agents we created."""
|
||||
"""Close the agents_client and clean up any agents we created."""
|
||||
await self._cleanup_agent_if_needed()
|
||||
await self._close_client_if_needed()
|
||||
|
||||
@@ -298,7 +237,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
settings: A dictionary of settings for the service.
|
||||
"""
|
||||
return cls(
|
||||
project_client=settings.get("project_client"),
|
||||
agents_client=settings.get("agents_client"),
|
||||
agent_id=settings.get("agent_id"),
|
||||
thread_id=settings.get("thread_id"),
|
||||
project_endpoint=settings.get("project_endpoint"),
|
||||
@@ -374,11 +313,14 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
args["instructions"] = run_options["instructions"]
|
||||
if "response_format" in run_options:
|
||||
args["response_format"] = run_options["response_format"]
|
||||
|
||||
if "temperature" in run_options:
|
||||
args["temperature"] = run_options["temperature"]
|
||||
if "top_p" in run_options:
|
||||
args["top_p"] = run_options["top_p"]
|
||||
created_agent = await self.project_client.agents.create_agent(**args)
|
||||
|
||||
created_agent = await self.agents_client.create_agent(**args)
|
||||
|
||||
self.agent_id = str(created_agent.id)
|
||||
self._agent_definition = created_agent
|
||||
self._should_delete_agent = True
|
||||
@@ -422,7 +364,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
args["tool_outputs"] = tool_outputs
|
||||
if tool_approvals:
|
||||
args["tool_approvals"] = tool_approvals
|
||||
await self.project_client.agents.runs.submit_tool_outputs_stream(**args) # type: ignore[reportUnknownMemberType]
|
||||
await self.agents_client.runs.submit_tool_outputs_stream(**args) # type: ignore[reportUnknownMemberType]
|
||||
# Pass the handler to the stream to continue processing
|
||||
stream = handler # type: ignore
|
||||
final_thread_id = thread_run.thread_id
|
||||
@@ -432,7 +374,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
|
||||
# Now create a new run and stream the results.
|
||||
run_options.pop("conversation_id", None)
|
||||
stream = await self.project_client.agents.runs.stream( # type: ignore[reportUnknownMemberType]
|
||||
stream = await self.agents_client.runs.stream( # type: ignore[reportUnknownMemberType]
|
||||
final_thread_id, agent_id=agent_id, **run_options
|
||||
)
|
||||
|
||||
@@ -443,9 +385,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
if thread_id is None:
|
||||
return None
|
||||
|
||||
async for run in self.project_client.agents.runs.list(
|
||||
thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING
|
||||
): # type: ignore[reportUnknownMemberType]
|
||||
async for run in self.agents_client.runs.list(thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING): # type: ignore[reportUnknownMemberType]
|
||||
if run.status not in [
|
||||
RunStatus.COMPLETED,
|
||||
RunStatus.CANCELLED,
|
||||
@@ -462,12 +402,12 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
if thread_id is not None:
|
||||
if thread_run is not None:
|
||||
# There was an active run; we need to cancel it before starting a new run.
|
||||
await self.project_client.agents.runs.cancel(thread_id, thread_run.id)
|
||||
await self.agents_client.runs.cancel(thread_id, thread_run.id)
|
||||
|
||||
return thread_id
|
||||
|
||||
# No thread ID was provided, so create a new thread.
|
||||
thread = await self.project_client.agents.threads.create(
|
||||
thread = await self.agents_client.threads.create(
|
||||
tool_resources=run_options.get("tool_resources"), metadata=run_options.get("metadata")
|
||||
)
|
||||
thread_id = thread.id
|
||||
@@ -476,7 +416,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
# once fixed, in the function above, readd:
|
||||
# `messages=run_options.pop("additional_messages")`
|
||||
for msg in run_options.pop("additional_messages", []):
|
||||
await self.project_client.agents.messages.create(
|
||||
await self.agents_client.messages.create(
|
||||
thread_id=thread_id, role=msg.role, content=msg.content, metadata=msg.metadata
|
||||
)
|
||||
# and remove until here.
|
||||
@@ -709,21 +649,21 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
return []
|
||||
|
||||
async def _close_client_if_needed(self) -> None:
|
||||
"""Close project_client session if we created it."""
|
||||
"""Close agents_client session if we created it."""
|
||||
if self._should_close_client:
|
||||
await self.project_client.close()
|
||||
await self.agents_client.close()
|
||||
|
||||
async def _cleanup_agent_if_needed(self) -> None:
|
||||
"""Clean up the agent if we created it."""
|
||||
if self._should_delete_agent and self.agent_id is not None:
|
||||
await self.project_client.agents.delete_agent(self.agent_id)
|
||||
await self.agents_client.delete_agent(self.agent_id)
|
||||
self.agent_id = None
|
||||
self._should_delete_agent = False
|
||||
|
||||
async def _load_agent_definition_if_needed(self) -> Agent | None:
|
||||
"""Load and cache agent details if not already loaded."""
|
||||
if self._agent_definition is None and self.agent_id is not None:
|
||||
self._agent_definition = await self.project_client.agents.get_agent(self.agent_id)
|
||||
self._agent_definition = await self.agents_client.get_agent(self.agent_id)
|
||||
return self._agent_definition
|
||||
|
||||
def _prepare_tool_choice(self, chat_options: ChatOptions) -> None:
|
||||
@@ -915,57 +855,32 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
config_args["set_lang"] = set_lang
|
||||
# Bing Grounding (support both connection_id and connection_name)
|
||||
connection_id = additional_props.get("connection_id") or os.getenv("BING_CONNECTION_ID")
|
||||
connection_name = additional_props.get("connection_name") or os.getenv("BING_CONNECTION_NAME")
|
||||
# Custom Bing Search
|
||||
custom_connection_name = additional_props.get("custom_connection_name") or os.getenv(
|
||||
"BING_CUSTOM_CONNECTION_NAME"
|
||||
custom_connection_id = additional_props.get("custom_connection_id") or os.getenv(
|
||||
"BING_CUSTOM_CONNECTION_ID"
|
||||
)
|
||||
custom_configuration_name = additional_props.get("custom_instance_name") or os.getenv(
|
||||
custom_instance_name = additional_props.get("custom_instance_name") or os.getenv(
|
||||
"BING_CUSTOM_INSTANCE_NAME"
|
||||
)
|
||||
bing_search: BingGroundingTool | BingCustomSearchTool | None = None
|
||||
if (
|
||||
(connection_id or connection_name)
|
||||
and not custom_connection_name
|
||||
and not custom_configuration_name
|
||||
):
|
||||
if (connection_id) and not custom_connection_id and not custom_instance_name:
|
||||
if connection_id:
|
||||
conn_id = connection_id
|
||||
elif connection_name:
|
||||
try:
|
||||
bing_connection = await self.project_client.connections.get(name=connection_name)
|
||||
except HttpResponseError as err:
|
||||
raise ServiceInitializationError(
|
||||
f"Bing connection '{connection_name}' not found in the Azure AI Project.",
|
||||
err,
|
||||
) from err
|
||||
else:
|
||||
conn_id = bing_connection.id
|
||||
else:
|
||||
raise ServiceInitializationError("Neither connection_id nor connection_name provided.")
|
||||
bing_search = BingGroundingTool(connection_id=conn_id, **config_args)
|
||||
if custom_connection_name and custom_configuration_name:
|
||||
try:
|
||||
bing_custom_connection = await self.project_client.connections.get(
|
||||
name=custom_connection_name
|
||||
)
|
||||
except HttpResponseError as err:
|
||||
raise ServiceInitializationError(
|
||||
f"Bing custom connection '{custom_connection_name}' not found in the Azure AI Project.",
|
||||
err,
|
||||
) from err
|
||||
else:
|
||||
bing_search = BingCustomSearchTool(
|
||||
connection_id=bing_custom_connection.id,
|
||||
instance_name=custom_configuration_name,
|
||||
**config_args,
|
||||
)
|
||||
if custom_connection_id and custom_instance_name:
|
||||
bing_search = BingCustomSearchTool(
|
||||
connection_id=custom_connection_id,
|
||||
instance_name=custom_instance_name,
|
||||
**config_args,
|
||||
)
|
||||
if not bing_search:
|
||||
raise ServiceInitializationError(
|
||||
"Bing search tool requires either 'connection_id' or 'connection_name' for Bing Grounding "
|
||||
"or both 'custom_connection_name' and 'custom_instance_name' for Custom Bing Search. "
|
||||
"Bing search tool requires either 'connection_id' for Bing Grounding "
|
||||
"or both 'custom_connection_id' and 'custom_instance_name' for Custom Bing Search. "
|
||||
"These can be provided via additional_properties or environment variables: "
|
||||
"'BING_CONNECTION_ID', 'BING_CONNECTION_NAME', 'BING_CUSTOM_CONNECTION_NAME', "
|
||||
"'BING_CONNECTION_ID', 'BING_CUSTOM_CONNECTION_ID', "
|
||||
"'BING_CUSTOM_INSTANCE_NAME'"
|
||||
)
|
||||
tool_definitions.extend(bing_search.definitions)
|
||||
@@ -1056,4 +971,4 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
Returns:
|
||||
The service URL for the chat client, or None if not set.
|
||||
"""
|
||||
return self.project_client._config.endpoint
|
||||
return self.agents_client._config.endpoint # type: ignore
|
||||
|
||||
@@ -0,0 +1,310 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import sys
|
||||
from collections.abc import MutableSequence
|
||||
from typing import Any, ClassVar, TypeVar
|
||||
|
||||
from agent_framework import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
TextContent,
|
||||
get_logger,
|
||||
use_chat_middleware,
|
||||
use_function_invocation,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.observability import use_observability
|
||||
from agent_framework.openai._responses_client import OpenAIBaseResponsesClient
|
||||
from azure.ai.projects.aio import AIProjectClient
|
||||
from azure.ai.projects.models import PromptAgentDefinition
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
from azure.core.exceptions import ResourceNotFoundError
|
||||
from openai.types.responses.parsed_response import (
|
||||
ParsedResponse,
|
||||
)
|
||||
from openai.types.responses.response import Response as OpenAIResponse
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
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
|
||||
|
||||
|
||||
logger = get_logger("agent_framework.azure")
|
||||
|
||||
|
||||
TAzureAIAgentClient = TypeVar("TAzureAIAgentClient", bound="AzureAIAgentClientV2")
|
||||
|
||||
|
||||
@use_function_invocation
|
||||
@use_observability
|
||||
@use_chat_middleware
|
||||
class AzureAIAgentClientV2(OpenAIBaseResponsesClient):
|
||||
"""Azure AI Agent Chat client."""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project_client: AIProjectClient | None = None,
|
||||
agent_name: str | None = None,
|
||||
agent_version: str | None = None,
|
||||
conversation_id: str | None = None,
|
||||
project_endpoint: str | None = None,
|
||||
model_deployment_name: str | None = None,
|
||||
async_credential: AsyncTokenCredential | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize an Azure AI Agent client.
|
||||
|
||||
Keyword Args:
|
||||
project_client: An existing AIProjectClient to use. If not provided, one will be created.
|
||||
agent_name: The name to use when creating new agents.
|
||||
agent_version: The version of the agent to use.
|
||||
conversation_id: Default conversation ID to use for conversations. Can be overridden by
|
||||
conversation_id property when making a request.
|
||||
project_endpoint: The Azure AI Project endpoint URL.
|
||||
Can also be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
|
||||
Ignored when a project_client is passed.
|
||||
model_deployment_name: The model deployment name to use for agent creation.
|
||||
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
|
||||
async_credential: Azure async credential to use for authentication.
|
||||
env_file_path: Path to environment file for loading settings.
|
||||
env_file_encoding: Encoding of the environment file.
|
||||
kwargs: Additional keyword arguments passed to the parent class.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import DefaultAzureCredential
|
||||
|
||||
# Using environment variables
|
||||
# Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
|
||||
# Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
|
||||
credential = DefaultAzureCredential()
|
||||
client = AzureAIAgentClient(async_credential=credential)
|
||||
|
||||
# Or passing parameters directly
|
||||
client = AzureAIAgentClient(
|
||||
project_endpoint="https://your-project.cognitiveservices.azure.com",
|
||||
model_deployment_name="gpt-4",
|
||||
async_credential=credential,
|
||||
)
|
||||
|
||||
# Or loading from a .env file
|
||||
client = AzureAIAgentClient(async_credential=credential, env_file_path="path/to/.env")
|
||||
"""
|
||||
try:
|
||||
azure_ai_settings = AzureAISettings(
|
||||
project_endpoint=project_endpoint,
|
||||
model_deployment_name=model_deployment_name,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to create Azure AI settings.", ex) from ex
|
||||
|
||||
# If no project_client is provided, create one
|
||||
should_close_client = False
|
||||
if project_client is None:
|
||||
if not azure_ai_settings.project_endpoint:
|
||||
raise ServiceInitializationError(
|
||||
"Azure AI project endpoint is required. Set via 'project_endpoint' parameter "
|
||||
"or 'AZURE_AI_PROJECT_ENDPOINT' environment variable."
|
||||
)
|
||||
|
||||
if not azure_ai_settings.model_deployment_name:
|
||||
raise ServiceInitializationError(
|
||||
"Azure AI model deployment name is required. Set via 'model_deployment_name' parameter "
|
||||
"or 'AZURE_AI_MODEL_DEPLOYMENT_NAME' environment variable."
|
||||
)
|
||||
|
||||
# Use provided credential
|
||||
if not async_credential:
|
||||
raise ServiceInitializationError("Azure credential is required when project_client is not provided.")
|
||||
project_client = AIProjectClient(
|
||||
endpoint=azure_ai_settings.project_endpoint,
|
||||
credential=async_credential,
|
||||
user_agent=AGENT_FRAMEWORK_USER_AGENT,
|
||||
)
|
||||
should_close_client = True
|
||||
|
||||
# Initialize parent
|
||||
super().__init__(
|
||||
model_id=azure_ai_settings.model_deployment_name, # type: ignore
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Initialize instance variables
|
||||
self.agent_name = agent_name
|
||||
self.agent_version = agent_version
|
||||
self.project_client = project_client
|
||||
self.credential = async_credential
|
||||
self.model_id = azure_ai_settings.model_deployment_name
|
||||
self.conversation_id = conversation_id
|
||||
self._should_close_client = should_close_client # Track whether we should close client connection
|
||||
|
||||
async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None:
|
||||
"""Use this method to setup tracing in your Azure AI Project.
|
||||
|
||||
This will take the connection string from the project project_client.
|
||||
It will override any connection string that is set in the environment variables.
|
||||
It will disable any OTLP endpoint that might have been set.
|
||||
"""
|
||||
try:
|
||||
conn_string = await self.project_client.telemetry.get_application_insights_connection_string()
|
||||
except ResourceNotFoundError:
|
||||
logger.warning(
|
||||
"No Application Insights connection string found for the Azure AI Project, "
|
||||
"please call setup_observability() manually."
|
||||
)
|
||||
return
|
||||
from agent_framework.observability import setup_observability
|
||||
|
||||
setup_observability(
|
||||
applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "Self":
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the project_client."""
|
||||
await self._close_client_if_needed()
|
||||
|
||||
async def _get_agent_reference_or_create(
|
||||
self, run_options: dict[str, Any], messages_instructions: str | None
|
||||
) -> dict[str, str]:
|
||||
"""Determine which agent to use and create if needed.
|
||||
|
||||
Returns:
|
||||
str: The agent_name to use
|
||||
"""
|
||||
agent_name = self.agent_name or "UnnamedAgent"
|
||||
|
||||
# If no agent_version is provided, create a new agent
|
||||
if self.agent_version is None:
|
||||
if "model" not in run_options or not run_options["model"]:
|
||||
raise ServiceInitializationError(
|
||||
"Model deployment name is required for agent creation, "
|
||||
"can also be passed to the get_response methods."
|
||||
)
|
||||
|
||||
args: dict[str, Any] = {
|
||||
"model": run_options["model"],
|
||||
}
|
||||
if "tools" in run_options:
|
||||
args["tools"] = run_options["tools"]
|
||||
|
||||
# Combine instructions from messages and options
|
||||
combined_instructions = [
|
||||
instructions
|
||||
for instructions in [messages_instructions, run_options.get("instructions")]
|
||||
if instructions
|
||||
]
|
||||
if combined_instructions:
|
||||
args["instructions"] = "".join(combined_instructions)
|
||||
|
||||
# TODO (dmytrostruk): Add response format
|
||||
|
||||
created_agent = await self.project_client.agents.create_version(
|
||||
agent_name=agent_name, definition=PromptAgentDefinition(**args)
|
||||
)
|
||||
|
||||
self.agent_name = created_agent.name
|
||||
self.agent_version = created_agent.version
|
||||
|
||||
return {"name": agent_name, "version": self.agent_version, "type": "agent_reference"}
|
||||
|
||||
async def _get_conversation_id_or_create(self, run_options: dict[str, Any]) -> str:
|
||||
# Since "conversation" property is used, remove "previous_response_id" from options
|
||||
# Use global conversation_id as fallback
|
||||
conversation_id = run_options.pop("previous_response_id", self.conversation_id)
|
||||
|
||||
if conversation_id:
|
||||
return conversation_id
|
||||
|
||||
# Create a new conversation with messages
|
||||
created_conversation = await self.client.conversations.create()
|
||||
return created_conversation.id
|
||||
|
||||
async def _close_client_if_needed(self) -> None:
|
||||
"""Close project_client session if we created it."""
|
||||
if self._should_close_client:
|
||||
await self.project_client.close()
|
||||
|
||||
def _prepare_input(self, messages: MutableSequence[ChatMessage]) -> tuple[list[ChatMessage], str | None]:
|
||||
"""Prepare input from messages and convert system/developer messages to instructions."""
|
||||
result: list[ChatMessage] = []
|
||||
instructions_list: list[str] = []
|
||||
instructions: str | None = None
|
||||
|
||||
# System/developer messages are turned into instructions, since there is no such message roles in Azure AI.
|
||||
for message in messages:
|
||||
if message.role.value in ["system", "developer"]:
|
||||
for text_content in [content for content in message.contents if isinstance(content, TextContent)]:
|
||||
instructions_list.append(text_content.text)
|
||||
else:
|
||||
result.append(message)
|
||||
|
||||
if len(instructions_list) > 0:
|
||||
instructions = "".join(instructions_list)
|
||||
|
||||
return result, instructions
|
||||
|
||||
async def prepare_options(
|
||||
self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
|
||||
) -> dict[str, Any]:
|
||||
prepared_messages, instructions = self._prepare_input(messages)
|
||||
run_options = await super().prepare_options(prepared_messages, chat_options)
|
||||
agent_reference = await self._get_agent_reference_or_create(run_options, instructions)
|
||||
|
||||
store = run_options.get("store", False)
|
||||
|
||||
if store:
|
||||
conversation_id = await self._get_conversation_id_or_create(run_options)
|
||||
run_options["conversation"] = conversation_id
|
||||
|
||||
run_options["extra_body"] = {"agent": agent_reference}
|
||||
|
||||
# Remove properties that are not supported
|
||||
# Model and tools captured in the agent setup
|
||||
if "model" in run_options:
|
||||
run_options.pop("model", None)
|
||||
|
||||
if "tools" in run_options:
|
||||
run_options.pop("tools", None)
|
||||
|
||||
return run_options
|
||||
|
||||
async def initialize_client(self):
|
||||
"""Initialize OpenAI client asynchronously."""
|
||||
self.client = await self.project_client.get_openai_client() # type: ignore
|
||||
|
||||
def get_conversation_id(self, response: OpenAIResponse | ParsedResponse[BaseModel], store: bool) -> str | None:
|
||||
"""Get the conversation ID from the response if store is True."""
|
||||
return response.conversation.id if response.conversation and store else None
|
||||
|
||||
def _update_agent_name(self, agent_name: str | None) -> None:
|
||||
"""Update the agent name in the chat client.
|
||||
|
||||
Args:
|
||||
agent_name: The new name for the agent.
|
||||
"""
|
||||
# This is a no-op in the base class, but can be overridden by subclasses
|
||||
# to update the agent name in the client.
|
||||
if agent_name and not self.agent_name:
|
||||
self.agent_name = agent_name
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import ClassVar
|
||||
|
||||
from agent_framework._pydantic import AFBaseSettings
|
||||
|
||||
|
||||
class AzureAISettings(AFBaseSettings):
|
||||
"""Azure AI Project settings.
|
||||
|
||||
The settings are first loaded from environment variables with the prefix 'AZURE_AI_'.
|
||||
If the environment variables are not found, the settings can be loaded from a .env file
|
||||
with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
|
||||
are ignored; however, validation will fail alerting that the settings are missing.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Azure AI Project endpoint URL.
|
||||
Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
|
||||
model_deployment_name: The name of the model deployment to use.
|
||||
Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
|
||||
env_file_path: If provided, the .env settings are read from this file path location.
|
||||
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.azure import AzureAISettings
|
||||
|
||||
# Using environment variables
|
||||
# Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
|
||||
# Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
|
||||
settings = AzureAISettings()
|
||||
|
||||
# Or passing parameters directly
|
||||
settings = AzureAISettings(
|
||||
project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4"
|
||||
)
|
||||
|
||||
# Or loading from a .env file
|
||||
settings = AzureAISettings(env_file_path="path/to/.env")
|
||||
"""
|
||||
|
||||
env_prefix: ClassVar[str] = "AZURE_AI_"
|
||||
|
||||
project_endpoint: str | None = None
|
||||
model_deployment_name: str | None = None
|
||||
@@ -23,7 +23,7 @@ classifiers = [
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core",
|
||||
"azure-ai-projects >= 1.0.0b11",
|
||||
"azure-ai-projects >= 2.0.0a20251103001",
|
||||
"azure-ai-agents == 1.2.0b5",
|
||||
]
|
||||
|
||||
@@ -84,3 +84,6 @@ test = "pytest --cov=agent_framework_azure_ai --cov-report=term-missing:skip-cov
|
||||
[build-system]
|
||||
requires = ["flit-core >= 3.11,<4.0"]
|
||||
build-backend = "flit_core.buildapi"
|
||||
|
||||
[tool.uv.sources]
|
||||
azure-ai-projects = { index = "azure-sdk-for-python" }
|
||||
@@ -44,31 +44,30 @@ def azure_ai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_ai_project_client() -> MagicMock:
|
||||
"""Fixture that provides a mock AIProjectClient."""
|
||||
def mock_agents_client() -> MagicMock:
|
||||
"""Fixture that provides a mock AgentsClient."""
|
||||
mock_client = MagicMock()
|
||||
|
||||
# Mock agents property
|
||||
mock_client.agents = MagicMock()
|
||||
mock_client.agents.create_agent = AsyncMock()
|
||||
mock_client.agents.delete_agent = AsyncMock()
|
||||
mock_client.create_agent = AsyncMock()
|
||||
mock_client.delete_agent = AsyncMock()
|
||||
|
||||
# Mock agent creation response
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.id = "test-agent-id"
|
||||
mock_client.agents.create_agent.return_value = mock_agent
|
||||
mock_client.create_agent.return_value = mock_agent
|
||||
|
||||
# Mock threads property
|
||||
mock_client.agents.threads = MagicMock()
|
||||
mock_client.agents.threads.create = AsyncMock()
|
||||
mock_client.agents.messages.create = AsyncMock()
|
||||
mock_client.threads = MagicMock()
|
||||
mock_client.threads.create = AsyncMock()
|
||||
mock_client.messages.create = AsyncMock()
|
||||
|
||||
# Mock runs property
|
||||
mock_client.agents.runs = MagicMock()
|
||||
mock_client.agents.runs.list = AsyncMock()
|
||||
mock_client.agents.runs.cancel = AsyncMock()
|
||||
mock_client.agents.runs.stream = AsyncMock()
|
||||
mock_client.agents.runs.submit_tool_outputs_stream = AsyncMock()
|
||||
mock_client.runs = MagicMock()
|
||||
mock_client.runs.list = AsyncMock()
|
||||
mock_client.runs.cancel = AsyncMock()
|
||||
mock_client.runs.stream = AsyncMock()
|
||||
mock_client.runs.submit_tool_outputs_stream = AsyncMock()
|
||||
|
||||
return mock_client
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,473 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
ChatClientProtocol,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from pydantic import ValidationError
|
||||
|
||||
from agent_framework_azure_ai import AzureAIAgentClientV2, AzureAISettings
|
||||
|
||||
|
||||
def create_test_azure_ai_client_v2(
|
||||
mock_project_client: MagicMock,
|
||||
agent_name: str | None = None,
|
||||
agent_version: str | None = None,
|
||||
conversation_id: str | None = None,
|
||||
azure_ai_settings: AzureAISettings | None = None,
|
||||
should_close_client: bool = False,
|
||||
) -> AzureAIAgentClientV2:
|
||||
"""Helper function to create AzureAIAgentClientV2 instances for testing, bypassing normal validation."""
|
||||
if azure_ai_settings is None:
|
||||
azure_ai_settings = AzureAISettings(env_file_path="test.env")
|
||||
|
||||
# Create client instance directly
|
||||
client = object.__new__(AzureAIAgentClientV2)
|
||||
|
||||
# Set attributes directly
|
||||
client.project_client = mock_project_client
|
||||
client.credential = None
|
||||
client.agent_name = agent_name
|
||||
client.agent_version = agent_version
|
||||
client.model_id = azure_ai_settings.model_deployment_name
|
||||
client.conversation_id = conversation_id
|
||||
client._should_close_client = should_close_client # type: ignore
|
||||
client.additional_properties = {}
|
||||
client.middleware = None
|
||||
|
||||
# Mock the OpenAI client attribute
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.conversations = MagicMock()
|
||||
mock_openai_client.conversations.create = AsyncMock()
|
||||
client.client = mock_openai_client
|
||||
|
||||
return client
|
||||
|
||||
|
||||
def test_azure_ai_settings_init(azure_ai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test AzureAISettings initialization."""
|
||||
settings = AzureAISettings()
|
||||
|
||||
assert settings.project_endpoint == azure_ai_unit_test_env["AZURE_AI_PROJECT_ENDPOINT"]
|
||||
assert settings.model_deployment_name == azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]
|
||||
|
||||
|
||||
def test_azure_ai_settings_init_with_explicit_values() -> None:
|
||||
"""Test AzureAISettings initialization with explicit values."""
|
||||
settings = AzureAISettings(
|
||||
project_endpoint="https://custom-endpoint.com/",
|
||||
model_deployment_name="custom-model",
|
||||
)
|
||||
|
||||
assert settings.project_endpoint == "https://custom-endpoint.com/"
|
||||
assert settings.model_deployment_name == "custom-model"
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_with_project_client(mock_project_client: MagicMock) -> None:
|
||||
"""Test AzureAIAgentClientV2 initialization with existing project_client."""
|
||||
with patch("agent_framework_azure_ai._chat_client_v2.AzureAISettings") as mock_settings:
|
||||
mock_settings.return_value.project_endpoint = None
|
||||
mock_settings.return_value.model_deployment_name = "test-model"
|
||||
|
||||
client = AzureAIAgentClientV2(
|
||||
project_client=mock_project_client,
|
||||
agent_name="test-agent",
|
||||
agent_version="1.0",
|
||||
)
|
||||
|
||||
assert client.project_client is mock_project_client
|
||||
assert client.agent_name == "test-agent"
|
||||
assert client.agent_version == "1.0"
|
||||
assert not client._should_close_client # type: ignore
|
||||
assert isinstance(client, ChatClientProtocol)
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_auto_create_client(
|
||||
azure_ai_unit_test_env: dict[str, str],
|
||||
mock_azure_credential: MagicMock,
|
||||
) -> None:
|
||||
"""Test AzureAIAgentClientV2 initialization with auto-created project_client."""
|
||||
with patch("agent_framework_azure_ai._chat_client_v2.AIProjectClient") as mock_ai_project_client:
|
||||
mock_project_client = MagicMock()
|
||||
mock_ai_project_client.return_value = mock_project_client
|
||||
|
||||
client = AzureAIAgentClientV2(
|
||||
project_endpoint=azure_ai_unit_test_env["AZURE_AI_PROJECT_ENDPOINT"],
|
||||
model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
|
||||
async_credential=mock_azure_credential,
|
||||
agent_name="test-agent",
|
||||
)
|
||||
|
||||
assert client.project_client is mock_project_client
|
||||
assert client.agent_name == "test-agent"
|
||||
assert client._should_close_client # type: ignore
|
||||
|
||||
# Verify AIProjectClient was called with correct parameters
|
||||
mock_ai_project_client.assert_called_once()
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_missing_project_endpoint() -> None:
|
||||
"""Test AzureAIAgentClientV2 initialization when project_endpoint is missing and no project_client provided."""
|
||||
with patch("agent_framework_azure_ai._chat_client_v2.AzureAISettings") as mock_settings:
|
||||
mock_settings.return_value.project_endpoint = None
|
||||
mock_settings.return_value.model_deployment_name = "test-model"
|
||||
|
||||
with pytest.raises(ServiceInitializationError, match="Azure AI project endpoint is required"):
|
||||
AzureAIAgentClientV2(async_credential=MagicMock())
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_missing_model_deployment() -> None:
|
||||
"""Test AzureAIAgentClientV2 initialization when model deployment is missing for agent creation."""
|
||||
with patch("agent_framework_azure_ai._chat_client_v2.AzureAISettings") as mock_settings:
|
||||
mock_settings.return_value.project_endpoint = "https://test.com"
|
||||
mock_settings.return_value.model_deployment_name = None
|
||||
|
||||
with pytest.raises(ServiceInitializationError, match="Azure AI model deployment name is required"):
|
||||
AzureAIAgentClientV2(async_credential=MagicMock())
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_missing_credential(azure_ai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test AzureAIAgentClientV2.__init__ when async_credential is missing and no project_client provided."""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Azure credential is required when project_client is not provided"
|
||||
):
|
||||
AzureAIAgentClientV2(
|
||||
project_endpoint=azure_ai_unit_test_env["AZURE_AI_PROJECT_ENDPOINT"],
|
||||
model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
|
||||
)
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_init_validation_error(mock_azure_credential: MagicMock) -> None:
|
||||
"""Test that ValidationError in AzureAISettings is properly handled."""
|
||||
with patch("agent_framework_azure_ai._chat_client_v2.AzureAISettings") as mock_settings:
|
||||
mock_settings.side_effect = ValidationError.from_exception_data("test", [])
|
||||
|
||||
with pytest.raises(ServiceInitializationError, match="Failed to create Azure AI settings"):
|
||||
AzureAIAgentClientV2(async_credential=mock_azure_credential)
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_get_agent_reference_or_create_existing_version(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_agent_reference_or_create when agent_version is already provided."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="existing-agent", agent_version="1.0")
|
||||
|
||||
agent_ref = await client._get_agent_reference_or_create({}, None) # type: ignore
|
||||
|
||||
assert agent_ref == {"name": "existing-agent", "version": "1.0", "type": "agent_reference"}
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_get_agent_reference_or_create_new_agent(
|
||||
mock_project_client: MagicMock,
|
||||
azure_ai_unit_test_env: dict[str, str],
|
||||
) -> None:
|
||||
"""Test _get_agent_reference_or_create when creating a new agent."""
|
||||
azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"])
|
||||
client = create_test_azure_ai_client_v2(
|
||||
mock_project_client, agent_name="new-agent", azure_ai_settings=azure_ai_settings
|
||||
)
|
||||
|
||||
# Mock agent creation response
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.name = "new-agent"
|
||||
mock_agent.version = "1.0"
|
||||
mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent)
|
||||
|
||||
run_options = {"model": azure_ai_settings.model_deployment_name}
|
||||
agent_ref = await client._get_agent_reference_or_create(run_options, None) # type: ignore
|
||||
|
||||
assert agent_ref == {"name": "new-agent", "version": "1.0", "type": "agent_reference"}
|
||||
assert client.agent_name == "new-agent"
|
||||
assert client.agent_version == "1.0"
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_get_agent_reference_missing_model(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_agent_reference_or_create when model is missing for agent creation."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="test-agent")
|
||||
|
||||
with pytest.raises(ServiceInitializationError, match="Model deployment name is required for agent creation"):
|
||||
await client._get_agent_reference_or_create({}, None) # type: ignore
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_get_conversation_id_or_create_existing(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_conversation_id_or_create when conversation_id is already provided."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, conversation_id="existing-conversation")
|
||||
|
||||
conversation_id = await client._get_conversation_id_or_create({}) # type: ignore
|
||||
|
||||
assert conversation_id == "existing-conversation"
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_get_conversation_id_or_create_new(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_conversation_id_or_create when creating a new conversation."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
# Mock conversation creation response
|
||||
mock_conversation = MagicMock()
|
||||
mock_conversation.id = "new-conversation-123"
|
||||
client.client.conversations.create = AsyncMock(return_value=mock_conversation)
|
||||
|
||||
conversation_id = await client._get_conversation_id_or_create({}) # type: ignore
|
||||
|
||||
assert conversation_id == "new-conversation-123"
|
||||
client.client.conversations.create.assert_called_once()
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_prepare_input_with_system_messages(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _prepare_input converts system/developer messages to instructions."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
messages = [
|
||||
ChatMessage(role=Role.SYSTEM, contents=[TextContent(text="You are a helpful assistant.")]),
|
||||
ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")]),
|
||||
ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="System response")]),
|
||||
]
|
||||
|
||||
result_messages, instructions = client._prepare_input(messages) # type: ignore
|
||||
|
||||
assert len(result_messages) == 2
|
||||
assert result_messages[0].role == Role.USER
|
||||
assert result_messages[1].role == Role.ASSISTANT
|
||||
assert instructions == "You are a helpful assistant."
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_prepare_input_no_system_messages(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _prepare_input with no system/developer messages."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
messages = [
|
||||
ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")]),
|
||||
ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="Hi there!")]),
|
||||
]
|
||||
|
||||
result_messages, instructions = client._prepare_input(messages) # type: ignore
|
||||
|
||||
assert len(result_messages) == 2
|
||||
assert instructions is None
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_prepare_options_basic(mock_project_client: MagicMock) -> None:
|
||||
"""Test prepare_options basic functionality."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="test-agent", agent_version="1.0")
|
||||
|
||||
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
|
||||
chat_options = ChatOptions()
|
||||
|
||||
with (
|
||||
patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}),
|
||||
patch.object(
|
||||
client,
|
||||
"_get_agent_reference_or_create",
|
||||
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
|
||||
),
|
||||
):
|
||||
run_options = await client.prepare_options(messages, chat_options)
|
||||
|
||||
assert "extra_body" in run_options
|
||||
assert run_options["extra_body"]["agent"]["name"] == "test-agent"
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_prepare_options_with_store(mock_project_client: MagicMock) -> None:
|
||||
"""Test prepare_options with store=True creates conversation."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="test-agent", agent_version="1.0")
|
||||
|
||||
# Mock conversation creation
|
||||
mock_conversation = MagicMock()
|
||||
mock_conversation.id = "new-conversation-456"
|
||||
client.client.conversations.create = AsyncMock(return_value=mock_conversation)
|
||||
|
||||
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
|
||||
chat_options = ChatOptions(store=True)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model", "store": True}
|
||||
),
|
||||
patch.object(
|
||||
client,
|
||||
"_get_agent_reference_or_create",
|
||||
return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"},
|
||||
),
|
||||
):
|
||||
run_options = await client.prepare_options(messages, chat_options)
|
||||
|
||||
assert "conversation" in run_options
|
||||
assert run_options["conversation"] == "new-conversation-456"
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_initialize_client(mock_project_client: MagicMock) -> None:
|
||||
"""Test initialize_client method."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_project_client.get_openai_client = AsyncMock(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_azure_ai_client_v2_get_conversation_id_from_response(mock_project_client: MagicMock) -> None:
|
||||
"""Test get_conversation_id method."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
# Test with conversation and store=True
|
||||
mock_response = MagicMock()
|
||||
mock_response.conversation.id = "test-conversation-123"
|
||||
|
||||
conversation_id = client.get_conversation_id(mock_response, store=True)
|
||||
assert conversation_id == "test-conversation-123"
|
||||
|
||||
# Test with store=False
|
||||
conversation_id = client.get_conversation_id(mock_response, store=False)
|
||||
assert conversation_id is None
|
||||
|
||||
# Test with no conversation
|
||||
mock_response.conversation = None
|
||||
conversation_id = client.get_conversation_id(mock_response, store=True)
|
||||
assert conversation_id is None
|
||||
|
||||
|
||||
def test_azure_ai_client_v2_update_agent_name(mock_project_client: MagicMock) -> None:
|
||||
"""Test _update_agent_name method."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client)
|
||||
|
||||
# Test updating agent name when current is None
|
||||
with patch.object(client, "_update_agent_name") as mock_update:
|
||||
mock_update.return_value = None
|
||||
client._update_agent_name("new-agent") # type: ignore
|
||||
mock_update.assert_called_once_with("new-agent")
|
||||
|
||||
# Test behavior when agent name is updated
|
||||
assert client.agent_name is None # Should remain None since we didn't actually update
|
||||
client.agent_name = "test-agent" # Manually set for the test
|
||||
|
||||
# Test with None input
|
||||
with patch.object(client, "_update_agent_name") as mock_update:
|
||||
mock_update.return_value = None
|
||||
client._update_agent_name(None) # type: ignore
|
||||
mock_update.assert_called_once_with(None)
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_async_context_manager(mock_project_client: MagicMock) -> None:
|
||||
"""Test async context manager functionality."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, should_close_client=True)
|
||||
|
||||
mock_project_client.close = AsyncMock()
|
||||
|
||||
async with client as ctx_client:
|
||||
assert ctx_client is client
|
||||
|
||||
# Should call close after exiting context
|
||||
mock_project_client.close.assert_called_once()
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_close_method(mock_project_client: MagicMock) -> None:
|
||||
"""Test close method."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, should_close_client=True)
|
||||
|
||||
mock_project_client.close = AsyncMock()
|
||||
|
||||
await client.close()
|
||||
|
||||
mock_project_client.close.assert_called_once()
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_close_client_when_should_close_false(mock_project_client: MagicMock) -> None:
|
||||
"""Test _close_client_if_needed when should_close_client is False."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, should_close_client=False)
|
||||
|
||||
mock_project_client.close = AsyncMock()
|
||||
|
||||
await client._close_client_if_needed() # type: ignore
|
||||
|
||||
# Should not call close when should_close_client is False
|
||||
mock_project_client.close.assert_not_called()
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_agent_creation_with_instructions(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test agent creation with combined instructions."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="test-agent")
|
||||
|
||||
# Mock agent creation response
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.name = "test-agent"
|
||||
mock_agent.version = "1.0"
|
||||
mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent)
|
||||
|
||||
run_options = {"model": "test-model", "instructions": "Option instructions. "}
|
||||
messages_instructions = "Message instructions. "
|
||||
|
||||
await client._get_agent_reference_or_create(run_options, messages_instructions) # type: ignore
|
||||
|
||||
# Verify agent was created with combined instructions
|
||||
call_args = mock_project_client.agents.create_version.call_args
|
||||
assert call_args[1]["definition"].instructions == "Message instructions. Option instructions. "
|
||||
|
||||
|
||||
async def test_azure_ai_client_v2_agent_creation_with_tools(
|
||||
mock_project_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test agent creation with tools."""
|
||||
client = create_test_azure_ai_client_v2(mock_project_client, agent_name="test-agent")
|
||||
|
||||
# Mock agent creation response
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.name = "test-agent"
|
||||
mock_agent.version = "1.0"
|
||||
mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent)
|
||||
|
||||
test_tools = [{"type": "function", "function": {"name": "test_tool"}}]
|
||||
run_options = {"model": "test-model", "tools": test_tools}
|
||||
|
||||
await client._get_agent_reference_or_create(run_options, None) # type: ignore
|
||||
|
||||
# Verify agent was created with tools
|
||||
call_args = mock_project_client.agents.create_version.call_args
|
||||
assert call_args[1]["definition"].tools == test_tools
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_project_client() -> MagicMock:
|
||||
"""Fixture that provides a mock AIProjectClient."""
|
||||
mock_client = MagicMock()
|
||||
|
||||
# Mock agents property
|
||||
mock_client.agents = MagicMock()
|
||||
mock_client.agents.create_version = AsyncMock()
|
||||
|
||||
# Mock conversations property
|
||||
mock_client.conversations = MagicMock()
|
||||
mock_client.conversations.create = AsyncMock()
|
||||
|
||||
# Mock telemetry property
|
||||
mock_client.telemetry = MagicMock()
|
||||
mock_client.telemetry.get_application_insights_connection_string = AsyncMock()
|
||||
|
||||
# Mock get_openai_client method
|
||||
mock_client.get_openai_client = AsyncMock()
|
||||
|
||||
# Mock close method
|
||||
mock_client.close = AsyncMock()
|
||||
|
||||
return mock_client
|
||||
@@ -6,6 +6,7 @@ from typing import Any
|
||||
|
||||
_IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"),
|
||||
"AzureAIAgentClientV2": ("agent_framework_azure_ai", "azure-ai"),
|
||||
"AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "core"),
|
||||
"AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "core"),
|
||||
"AzureAISettings": ("agent_framework_azure_ai", "azure-ai"),
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework_azure_ai import AzureAIAgentClient, AzureAISettings
|
||||
from agent_framework_azure_ai import AzureAIAgentClient, AzureAIAgentClientV2, AzureAISettings
|
||||
|
||||
from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
|
||||
from agent_framework.azure._chat_client import AzureOpenAIChatClient
|
||||
@@ -10,6 +10,7 @@ from agent_framework.azure._shared import AzureOpenAISettings
|
||||
|
||||
__all__ = [
|
||||
"AzureAIAgentClient",
|
||||
"AzureAIAgentClientV2",
|
||||
"AzureAISettings",
|
||||
"AzureOpenAIAssistantsClient",
|
||||
"AzureOpenAIChatClient",
|
||||
|
||||
@@ -89,23 +89,24 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
options_dict = self._prepare_options(messages, chat_options)
|
||||
client = await self.ensure_client()
|
||||
run_options = await self.prepare_options(messages, chat_options)
|
||||
try:
|
||||
if not chat_options.response_format:
|
||||
response = await self.client.responses.create(
|
||||
response = await client.responses.create(
|
||||
stream=False,
|
||||
**options_dict,
|
||||
**run_options,
|
||||
)
|
||||
chat_options.conversation_id = response.id if chat_options.store is True else None
|
||||
chat_options.conversation_id = self.get_conversation_id(response, chat_options.store)
|
||||
return self._create_response_content(response, chat_options=chat_options)
|
||||
# create call does not support response_format, so we need to handle it via parse call
|
||||
resp_format = chat_options.response_format
|
||||
parsed_response: ParsedResponse[BaseModel] = await self.client.responses.parse(
|
||||
parsed_response: ParsedResponse[BaseModel] = await client.responses.parse(
|
||||
text_format=resp_format,
|
||||
stream=False,
|
||||
**options_dict,
|
||||
**run_options,
|
||||
)
|
||||
chat_options.conversation_id = parsed_response.id if chat_options.store is True else None
|
||||
chat_options.conversation_id = self.get_conversation_id(parsed_response, chat_options.store)
|
||||
return self._create_response_content(parsed_response, chat_options=chat_options)
|
||||
except BadRequestError as ex:
|
||||
if ex.code == "content_filter":
|
||||
@@ -130,13 +131,14 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
options_dict = self._prepare_options(messages, chat_options)
|
||||
client = await self.ensure_client()
|
||||
run_options = await self.prepare_options(messages, chat_options)
|
||||
function_call_ids: dict[int, tuple[str, str]] = {} # output_index: (call_id, name)
|
||||
try:
|
||||
if not chat_options.response_format:
|
||||
response = await self.client.responses.create(
|
||||
response = await client.responses.create(
|
||||
stream=True,
|
||||
**options_dict,
|
||||
**run_options,
|
||||
)
|
||||
async for chunk in response:
|
||||
update = self._create_streaming_response_content(
|
||||
@@ -145,9 +147,9 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
yield update
|
||||
return
|
||||
# create call does not support response_format, so we need to handle it via stream call
|
||||
async with self.client.responses.stream(
|
||||
async with client.responses.stream(
|
||||
text_format=chat_options.response_format,
|
||||
**options_dict,
|
||||
**run_options,
|
||||
) as response:
|
||||
async for chunk in response:
|
||||
update = self._create_streaming_response_content(
|
||||
@@ -170,6 +172,10 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
inner_exception=ex,
|
||||
) from ex
|
||||
|
||||
def get_conversation_id(self, response: OpenAIResponse | ParsedResponse[BaseModel], store: bool) -> str | None:
|
||||
"""Get the conversation ID from the response if store is True."""
|
||||
return response.id if store else None
|
||||
|
||||
# region Prep methods
|
||||
|
||||
def _tools_to_response_tools(
|
||||
@@ -298,9 +304,11 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
response_tools.append(tool_dict)
|
||||
return response_tools
|
||||
|
||||
def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
|
||||
async def prepare_options(
|
||||
self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
|
||||
) -> dict[str, Any]:
|
||||
"""Take ChatOptions and create the specific options for Responses API."""
|
||||
options_dict: dict[str, Any] = chat_options.to_dict(
|
||||
run_options: dict[str, Any] = chat_options.to_dict(
|
||||
exclude={
|
||||
"type",
|
||||
"response_format", # handled in inner get methods
|
||||
@@ -319,35 +327,35 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
"max_tokens": "max_output_tokens",
|
||||
}
|
||||
for old_key, new_key in translations.items():
|
||||
if old_key in options_dict and old_key != new_key:
|
||||
options_dict[new_key] = options_dict.pop(old_key)
|
||||
if old_key in run_options and old_key != new_key:
|
||||
run_options[new_key] = run_options.pop(old_key)
|
||||
|
||||
# tools
|
||||
if chat_options.tools is None:
|
||||
options_dict.pop("parallel_tool_calls", None)
|
||||
run_options.pop("parallel_tool_calls", None)
|
||||
else:
|
||||
options_dict["tools"] = self._tools_to_response_tools(chat_options.tools)
|
||||
run_options["tools"] = self._tools_to_response_tools(chat_options.tools)
|
||||
|
||||
# model id
|
||||
if not options_dict.get("model"):
|
||||
options_dict["model"] = self.model_id
|
||||
if not run_options.get("model"):
|
||||
run_options["model"] = self.model_id
|
||||
|
||||
# messages
|
||||
request_input = self._prepare_chat_messages_for_request(messages)
|
||||
if not request_input:
|
||||
raise ServiceInvalidRequestError("Messages are required for chat completions")
|
||||
options_dict["input"] = request_input
|
||||
run_options["input"] = request_input
|
||||
|
||||
# additional provider specific settings
|
||||
if additional_properties := options_dict.pop("additional_properties", None):
|
||||
if additional_properties := run_options.pop("additional_properties", None):
|
||||
for key, value in additional_properties.items():
|
||||
if value is not None:
|
||||
options_dict[key] = value
|
||||
if "store" not in options_dict:
|
||||
options_dict["store"] = False
|
||||
if (tool_choice := options_dict.get("tool_choice")) and len(tool_choice.keys()) == 1:
|
||||
options_dict["tool_choice"] = tool_choice["mode"]
|
||||
return options_dict
|
||||
run_options[key] = value
|
||||
if "store" not in run_options:
|
||||
run_options["store"] = False
|
||||
if (tool_choice := run_options.get("tool_choice")) and len(tool_choice.keys()) == 1:
|
||||
run_options["tool_choice"] = tool_choice["mode"]
|
||||
return run_options
|
||||
|
||||
def _prepare_chat_messages_for_request(self, chat_messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
|
||||
"""Prepare the chat messages for a request.
|
||||
@@ -749,7 +757,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
"raw_representation": response,
|
||||
}
|
||||
if chat_options.store:
|
||||
args["conversation_id"] = response.id
|
||||
args["conversation_id"] = self.get_conversation_id(response, chat_options.store)
|
||||
if response.usage and (usage_details := self._usage_details_from_openai(response.usage)):
|
||||
args["usage_details"] = usage_details
|
||||
if structured_response:
|
||||
@@ -849,7 +857,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
|
||||
contents.append(TextReasoningContent(text=event.text, raw_representation=event))
|
||||
metadata.update(self._get_metadata_from_response(event))
|
||||
case "response.completed":
|
||||
conversation_id = event.response.id if chat_options.store is True else None
|
||||
conversation_id = self.get_conversation_id(event.response, chat_options.store)
|
||||
model = event.response.model
|
||||
if event.response.usage:
|
||||
usage = self._usage_details_from_openai(event.response.usage)
|
||||
|
||||
@@ -127,7 +127,7 @@ class OpenAIBase(SerializationMixin):
|
||||
|
||||
INJECTABLE: ClassVar[set[str]] = {"client"}
|
||||
|
||||
def __init__(self, *, client: AsyncOpenAI, model_id: str, **kwargs: Any) -> None:
|
||||
def __init__(self, *, model_id: str, client: AsyncOpenAI | None = None, **kwargs: Any) -> None:
|
||||
"""Initialize OpenAIBase.
|
||||
|
||||
Keyword Args:
|
||||
@@ -162,6 +162,21 @@ class OpenAIBase(SerializationMixin):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
async def initialize_client(self):
|
||||
"""Initialize OpenAI client asynchronously.
|
||||
|
||||
Override in subclasses to initialize the OpenAI client asynchronously.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def ensure_client(self) -> AsyncOpenAI:
|
||||
"""Ensure OpenAI client is initialized."""
|
||||
await self.initialize_client()
|
||||
if self.client is None:
|
||||
raise ServiceInitializationError("OpenAI client is not initialized")
|
||||
|
||||
return self.client
|
||||
|
||||
def _get_api_key(
|
||||
self, api_key: str | SecretStr | Callable[[], str | Awaitable[str]] | None
|
||||
) -> str | Callable[[], str | Awaitable[str]] | None:
|
||||
|
||||
Reference in New Issue
Block a user