agent/agent-framework
1
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: openai updates (#388)

Browse Source 
* openai updates

* rebuild of openai structure

* updated responses structure

* renamed sample

* added file id support to code interpreter

* added hosted file ids to code interpretor

* mypy fixes

* removed default az cred from codebase

* updated agent name setup

* added kwargs to entra methods

* and further kwargs

* extra comment

* updated all samples

* readded custom get methods for responses

* updated int tests with ad credential

* missed one
This commit is contained in:
Eduard van Valkenburg
2025-08-12 08:14:22 +02:00
committed by GitHub
Unverified
parent 19676978e9
commit df9d85d1f0
53 changed files with 1668 additions and 1470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure/agent_framework_azure/_assistants_client.py
+9 -9
View File
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Mapping
from typing import Any, ClassVar
from typing import TYPE_CHECKING, Any, ClassVar
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai import OpenAIAssistantsClient
@@ -9,9 +9,10 @@ from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
from ._shared import (
AzureOpenAISettings,
)
from ._shared import AzureOpenAISettings
if TYPE_CHECKING:
from azure.identity import ChainedTokenCredential
__all__ = ["AzureAssistantsClient"]
@@ -20,7 +21,6 @@ class AzureAssistantsClient(OpenAIAssistantsClient):
"""Azure OpenAI Assistants client."""
DEFAULT_AZURE_API_VERSION: ClassVar[str] = "2024-05-01-preview"
MODEL_PROVIDER_NAME: ClassVar[str] = "azure_openai" # type: ignore[reportIncompatibleVariableOverride, misc]
def __init__(
self,
@@ -35,6 +35,7 @@ class AzureAssistantsClient(OpenAIAssistantsClient):
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
ad_credential: "ChainedTokenCredential | None" = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
@@ -61,6 +62,7 @@ class AzureAssistantsClient(OpenAIAssistantsClient):
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
ad_credential: The Azure AD credential to use for authentication. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
@@ -93,11 +95,9 @@ class AzureAssistantsClient(OpenAIAssistantsClient):
and not ad_token
and not ad_token_provider
and azure_openai_settings.token_endpoint
and ad_credential
):
# Try to get token using Entra ID if no other auth method is provided
from ._entra_id_authentication import get_entra_auth_token
ad_token = get_entra_auth_token(azure_openai_settings.token_endpoint)
ad_token = azure_openai_settings.get_azure_auth_token(ad_credential)
if not async_client and not azure_openai_settings.api_key and not ad_token and not ad_token_provider:
raise ServiceInitializationError("The Azure OpenAI API key, ad_token, or ad_token_provider is required.")
python/packages/azure/agent_framework_azure/_chat_client.py
+58 -104
View File
@@ -2,25 +2,21 @@
import json
import logging
import sys
from collections.abc import Mapping
from copy import deepcopy
from typing import Any, TypeVar
from uuid import uuid4
from agent_framework import (
ChatFinishReason,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
FunctionResultContent,
CitationAnnotation,
TextContent,
)
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai._chat_client import OpenAIChatClientBase
from agent_framework.openai._shared import OpenAIModelTypes
from azure.identity import ChainedTokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
@@ -30,9 +26,15 @@ from ._shared import (
AzureOpenAISettings,
)
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
logger: logging.Logger = logging.getLogger(__name__)
TChatResponse = TypeVar("TChatResponse", ChatResponse, ChatResponseUpdate)
TAzureChatClient = TypeVar("TAzureChatClient", bound="AzureChatClient")
class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
@@ -48,6 +50,7 @@ class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
ad_credential: ChainedTokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
@@ -70,6 +73,7 @@ class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
ad_credential: The Azure Active Directory credential. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
@@ -107,14 +111,14 @@ class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
ad_token=ad_token,
ad_token_provider=ad_token_provider,
token_endpoint=azure_openai_settings.token_endpoint,
ad_credential=ad_credential,
default_headers=default_headers,
ai_model_type=OpenAIModelTypes.CHAT,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls, settings: dict[str, Any]) -> "AzureChatClient":
def from_dict(cls: type[TAzureChatClient], settings: dict[str, Any]) -> TAzureChatClient:
"""Initialize an Azure OpenAI service from a dictionary of settings.
Args:
@@ -122,7 +126,7 @@ class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
should contain keys: service_id, and optionally:
ad_auth, ad_token_provider, default_headers
"""
return AzureChatClient(
return cls(
api_key=settings.get("api_key"),
deployment_name=settings.get("deployment_name"),
endpoint=settings.get("endpoint"),
@@ -134,99 +138,49 @@ class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
env_file_path=settings.get("env_file_path"),
)
def _create_chat_message_content(
self, response: ChatCompletion, choice: Choice, response_metadata: dict[str, Any]
) -> ChatResponse:
"""Create an Azure chat message content object from a choice."""
content = super()._create_chat_message_content(response, choice, response_metadata)
return self._add_tool_message_to_chat_message_content(content, choice)
@override
def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
"""Parse the choice into a TextContent object.
def _create_streaming_chat_message_content(
self,
chunk: ChatCompletionChunk,
choice: ChunkChoice,
chunk_metadata: dict[str, Any],
) -> ChatResponseUpdate:
"""Create an Azure streaming chat message content object from a choice."""
content = super()._create_streaming_chat_message_content(chunk, choice, chunk_metadata)
assert isinstance(content, ChatResponseUpdate) and isinstance(choice, ChunkChoice) # nosec # noqa: S101
return self._add_tool_message_to_chat_message_content(content, choice)
def _add_tool_message_to_chat_message_content(
self,
content: TChatResponse,
choice: Choice | ChunkChoice,
) -> TChatResponse:
if tool_message := self._get_tool_message_from_chat_choice(choice=choice):
if not isinstance(tool_message, dict):
# try to json, to ensure it is a dictionary
try:
tool_message = json.loads(tool_message)
except json.JSONDecodeError:
logger.warning("Tool message is not a dictionary, ignore context.")
return content
function_call = FunctionCallContent(
call_id=str(uuid4()),
name="Azure-OnYourData",
arguments={"query": tool_message.get("intent", [])},
)
result = FunctionResultContent(
call_id=function_call.call_id,
result=tool_message["citations"],
exception=function_call.exception,
additional_properties=function_call.additional_properties,
)
inner_content = content.messages[0].contents if isinstance(content, ChatResponse) else content.contents
inner_content.insert(0, function_call)
inner_content.insert(1, result)
return content
def _get_tool_message_from_chat_choice(self, choice: Choice | ChunkChoice) -> dict[str, Any] | None:
"""Get the tool message from a choice."""
content = choice.message if isinstance(choice, Choice) else choice.delta
# When you enable asynchronous content filtering in Azure OpenAI, you may receive empty deltas
if content and content.model_extra is not None:
return content.model_extra.get("context", None)
# openai allows extra content, so model_extra will be a dict, but we need to check anyway, but no way to test.
return None # pragma: no cover
@staticmethod
def _split_message(message: "ChatResponse") -> ChatResponse:
"""Split an Azure On Your Data response into separate ChatMessages within the ChatResponse.
If the message does not have three contents, and those three are one each of:
FunctionCallContent, FunctionResultContent, and TextContent,
it will not return three messages, potentially only one or two.
The order of the returned messages is as expected by OpenAI.
Overwritten from OpenAIChatClientBase to deal with Azure On Your Data function.
For docs see:
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
"""
if len(message.messages) == 0:
return message
if len(message.messages[0].contents) != 3:
return message
messages = {
"tool_call": deepcopy(message.messages[0]),
"tool_result": deepcopy(message.messages[0]),
"assistant": deepcopy(message.messages[0]),
}
for key, msg in messages.items():
if key == "tool_call":
msg.contents = [item for item in msg.contents if isinstance(item, FunctionCallContent)]
message.finish_reason = ChatFinishReason.TOOL_CALLS
if key == "tool_result":
msg.contents = [item for item in msg.contents if isinstance(item, FunctionResultContent)]
if key == "assistant":
msg.contents = [item for item in msg.contents if isinstance(item, TextContent)]
message = choice.message if isinstance(choice, Choice) else choice.delta
if hasattr(message, "refusal") and message.refusal:
return TextContent(text=message.refusal, raw_representation=choice)
if not message.content:
return None
text_content = TextContent(text=message.content, raw_representation=choice)
if not message.model_extra or "context" not in message.model_extra:
return text_content
return ChatResponse(
response_id=message.response_id,
conversation_id=message.conversation_id,
messages=[messages["tool_call"], messages["tool_result"], messages["assistant"]],
created_at=message.created_at,
model_id=message.ai_model_id,
usage_details=message.usage_details,
finish_reason=message.finish_reason,
additional_properties=message.additional_properties,
)
context: dict[str, Any] | str = message.context # type: ignore[assignment, union-attr]
if isinstance(context, str):
try:
context = json.loads(context)
except json.JSONDecodeError:
logger.warning("Context is not a valid JSON string, ignoring context.")
return text_content
if not isinstance(context, dict):
logger.warning("Context is not a valid dictionary, ignoring context.")
return text_content
# `all_retrieved_documents` is currently not used, but can be retrieved
# through the raw_representation in the text content.
if intent := context.get("intent"):
text_content.additional_properties = {"intent": intent}
if citations := context.get("citations"):
text_content.annotations = []
for citation in citations:
text_content.annotations.append(
CitationAnnotation(
title=citation.get("title", ""),
url=citation.get("url", ""),
snippet=citation.get("content", ""),
file_id=citation.get("filepath", ""),
tool_name="Azure-on-your-Data",
additional_properties={"chunk_id": citation.get("chunk_id", "")},
raw_representation=citation,
)
)
return text_content
python/packages/azure/agent_framework_azure/_entra_id_authentication.py
+47 -7
View File
@@ -1,22 +1,33 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from typing import TYPE_CHECKING, Any
from agent_framework.exceptions import ServiceInvalidAuthError
from azure.core.exceptions import ClientAuthenticationError
from azure.identity import DefaultAzureCredential
if TYPE_CHECKING:
from azure.identity import ChainedTokenCredential
from azure.identity.aio import ChainedTokenCredential as AsyncChainedTokenCredential
logger: logging.Logger = logging.getLogger(__name__)
def get_entra_auth_token(token_endpoint: str) -> str | None:
def get_entra_auth_token(
credential: "ChainedTokenCredential",
token_endpoint: str,
**kwargs: Any,
) -> str | None:
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint.
The token endpoint may be specified as an environment variable, via the .env
file or as an argument. If the token endpoint is not provided, the default is None.
Args:
credential: The Azure credential to use for authentication.
for dev, you can use `DefaultAzureCredential`, but this is not recommended for production.
token_endpoint: The token endpoint to use to retrieve the authentication token.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
@@ -26,12 +37,41 @@ def get_entra_auth_token(token_endpoint: str) -> str | None:
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
credential = DefaultAzureCredential()
try:
auth_token = credential.get_token(token_endpoint)
except ClientAuthenticationError:
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`.")
auth_token = credential.get_token(token_endpoint, **kwargs)
except ClientAuthenticationError as ex:
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
return None
return auth_token.token if auth_token else None
async def get_entra_auth_token_async(
credential: "AsyncChainedTokenCredential", token_endpoint: str, **kwargs: Any
) -> str | None:
"""Retrieve a async Microsoft Entra Auth Token for a given token endpoint.
The token endpoint may be specified as an environment variable, via the .env
file or as an argument. If the token endpoint is not provided, the default is None.
Args:
credential: The async Azure credential to use for authentication.
for dev, you can use `DefaultAzureCredential`, but this is not recommended for production.
token_endpoint: The token endpoint to use to retrieve the authentication token.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
"""
if not token_endpoint:
raise ServiceInvalidAuthError(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
try:
auth_token = await credential.get_token(token_endpoint, **kwargs)
except ClientAuthenticationError as ex:
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
return None
return auth_token.token if auth_token else None
python/packages/azure/agent_framework_azure/_responses_client.py
+9 -7
View File
@@ -1,14 +1,14 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Mapping
from typing import Any, ClassVar
from typing import Any, TypeVar
from urllib.parse import urljoin
from agent_framework import use_tool_calling
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai._responses_client import OpenAIResponsesClientBase
from agent_framework.openai._shared import OpenAIModelTypes
from agent_framework.telemetry import use_telemetry
from azure.identity import ChainedTokenCredential
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
from pydantic import SecretStr, ValidationError
from pydantic.networks import AnyUrl
@@ -18,14 +18,14 @@ from ._shared import (
AzureOpenAISettings,
)
TAzureResponsesClient = TypeVar("TAzureResponsesClient", bound="AzureResponsesClient")
@use_telemetry
@use_tool_calling
class AzureResponsesClient(AzureOpenAIConfigBase, OpenAIResponsesClientBase):
"""Azure Responses completion class."""
MODEL_PROVIDER_NAME: ClassVar[str] = "azure_openai" # type: ignore[reportIncompatibleVariableOverride, misc]
def __init__(
self,
api_key: str | None = None,
@@ -36,6 +36,7 @@ class AzureResponsesClient(AzureOpenAIConfigBase, OpenAIResponsesClientBase):
ad_token: str | None = None,
ad_token_provider: AsyncAzureADTokenProvider | None = None,
token_endpoint: str | None = None,
ad_credential: ChainedTokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
@@ -58,6 +59,7 @@ class AzureResponsesClient(AzureOpenAIConfigBase, OpenAIResponsesClientBase):
ad_token: The Azure Active Directory token. (Optional)
ad_token_provider: The Azure Active Directory token provider. (Optional)
token_endpoint: The token endpoint to request an Azure token. (Optional)
ad_credential: The Azure Active Directory credential. (Optional)
default_headers: The default headers mapping of string keys to
string values for HTTP requests. (Optional)
async_client: An existing client to use. (Optional)
@@ -105,20 +107,20 @@ class AzureResponsesClient(AzureOpenAIConfigBase, OpenAIResponsesClientBase):
ad_token=ad_token,
ad_token_provider=ad_token_provider,
token_endpoint=azure_openai_settings.token_endpoint,
ad_credential=ad_credential,
default_headers=default_headers,
ai_model_type=OpenAIModelTypes.RESPONSE,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls, settings: dict[str, Any]) -> "AzureResponsesClient":
def from_dict(cls: type[TAzureResponsesClient], settings: dict[str, Any]) -> TAzureResponsesClient:
"""Initialize an Open AI service from a dictionary of settings.
Args:
settings: A dictionary of settings for the service.
"""
return AzureResponsesClient(
return cls(
api_key=settings.get("api_key"),
deployment_name=settings.get("deployment_name"),
endpoint=settings.get("endpoint"),
python/packages/azure/agent_framework_azure/_shared.py
+30 -28
View File
@@ -8,8 +8,9 @@ from typing import Any, ClassVar, Final
from agent_framework._pydantic import AFBaseSettings, HttpsUrl
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai._shared import OpenAIHandler, OpenAIModelTypes
from agent_framework.openai._shared import OpenAIHandler
from agent_framework.telemetry import USER_AGENT_KEY
from azure.identity import ChainedTokenCredential
from openai.lib.azure import AsyncAzureOpenAI
from pydantic import ConfigDict, SecretStr, model_validator, validate_call
@@ -20,6 +21,7 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import Self # pragma: no cover
logger: logging.Logger = logging.getLogger(__name__)
@@ -132,7 +134,9 @@ class AzureOpenAISettings(AFBaseSettings):
default_api_version: str = DEFAULT_AZURE_API_VERSION
default_token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT
def get_azure_openai_auth_token(self, token_endpoint: str | None = None) -> str | None:
def get_azure_auth_token(
self, credential: "ChainedTokenCredential", token_endpoint: str | None = None, **kwargs: Any
) -> str | None:
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint for the use with Azure OpenAI.
The required role for the token is `Cognitive Services OpenAI Contributor`.
@@ -141,7 +145,9 @@ class AzureOpenAISettings(AFBaseSettings):
The `token_endpoint` argument takes precedence over the `token_endpoint` attribute.
Args:
credential: The Azure AD credential to use.
token_endpoint: The token endpoint to use. Defaults to `https://cognitiveservices.azure.com/.default`.
**kwargs: Additional keyword arguments to pass to the token retrieval method.
Returns:
The Azure token or None if the token could not be retrieved.
@@ -149,10 +155,8 @@ class AzureOpenAISettings(AFBaseSettings):
Raises:
ServiceInitializationError: If the token endpoint is not provided.
"""
endpoint_to_use = token_endpoint or self.token_endpoint
if endpoint_to_use is None: # type: ignore
raise ServiceInitializationError("Please provide a token endpoint to retrieve the authentication token.")
return get_entra_auth_token(endpoint_to_use)
endpoint_to_use = token_endpoint or self.token_endpoint or self.default_token_endpoint
return get_entra_auth_token(credential, endpoint_to_use, **kwargs)
@model_validator(mode="after")
def _validate_fields(self) -> Self:
@@ -164,11 +168,12 @@ class AzureOpenAISettings(AFBaseSettings):
class AzureOpenAIConfigBase(OpenAIHandler):
"""Internal class for configuring a connection to an Azure OpenAI service."""
MODEL_PROVIDER_NAME: ClassVar[str] = "azure_openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def __init__(
self,
deployment_name: str,
ai_model_type: OpenAIModelTypes,
endpoint: HttpsUrl | None = None,
base_url: HttpsUrl | None = None,
api_version: str = DEFAULT_AZURE_API_VERSION,
@@ -176,6 +181,7 @@ class AzureOpenAIConfigBase(OpenAIHandler):
ad_token: str | None = None,
ad_token_provider: Callable[[], str | Awaitable[str]] | None = None,
token_endpoint: str | None = None,
ad_credential: ChainedTokenCredential | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncAzureOpenAI | None = None,
instruction_role: str | None = None,
@@ -187,20 +193,20 @@ class AzureOpenAIConfigBase(OpenAIHandler):
This is necessary for types like `HttpsUrl` and `OpenAIModelTypes`.
Args:
deployment_name (str): Name of the deployment.
ai_model_type (OpenAIModelTypes): The type of OpenAI model to deploy.
endpoint (HttpsUrl): The specific endpoint URL for the deployment. (Optional)
base_url (Url): The base URL for Azure services. (Optional)
api_version (str): Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
api_key (str): API key for Azure services. (Optional)
ad_token (str): Azure AD token for authentication. (Optional)
ad_token_provider (Callable[[], Union[str, Awaitable[str]]]): A callable
or coroutine function providing Azure AD tokens. (Optional)
token_endpoint (str): Azure AD token endpoint use to get the token. (Optional)
default_headers (Union[Mapping[str, str], None]): Default headers for HTTP requests. (Optional)
client (AsyncAzureOpenAI): An existing client to use. (Optional)
instruction_role (str | None): The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`. (Optional)
deployment_name: Name of the deployment.
ai_model_type: The type of OpenAI model to deploy.
endpoint: The specific endpoint URL for the deployment.
base_url: The base URL for Azure services.
api_version: Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
api_key: API key for Azure services.
ad_token: Azure AD token for authentication.
ad_token_provider: A callable or coroutine function providing Azure AD tokens.
token_endpoint: Azure AD token endpoint use to get the token.
ad_credential: Azure AD credential for authentication.
default_headers: Default headers for HTTP requests.
client: An existing client to use.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
kwargs: Additional keyword arguments.
"""
@@ -211,8 +217,8 @@ class AzureOpenAIConfigBase(OpenAIHandler):
# If the client is None, the api_key is none, the ad_token is none, and the ad_token_provider is none,
# then we will attempt to get the ad_token using the default endpoint specified in the Azure OpenAI
# settings.
if not api_key and not ad_token_provider and not ad_token and token_endpoint:
ad_token = get_entra_auth_token(token_endpoint)
if not api_key and not ad_token_provider and not ad_token and token_endpoint and ad_credential:
ad_token = get_entra_auth_token(ad_credential, token_endpoint)
if not api_key and not ad_token and not ad_token_provider:
raise ServiceInitializationError(
@@ -237,10 +243,8 @@ class AzureOpenAIConfigBase(OpenAIHandler):
args["base_url"] = str(base_url)
if endpoint and not base_url:
args["azure_endpoint"] = str(endpoint)
# TODO (eavanvalkenburg): Remove the check on model type when the package fixes: https://github.com/openai/openai-python/issues/2120
if deployment_name and ai_model_type != OpenAIModelTypes.REALTIME:
if deployment_name:
args["azure_deployment"] = deployment_name
if "websocket_base_url" in kwargs:
args["websocket_base_url"] = kwargs.pop("websocket_base_url")
@@ -248,7 +252,6 @@ class AzureOpenAIConfigBase(OpenAIHandler):
args = {
"ai_model_id": deployment_name,
"client": client,
"ai_model_type": ai_model_type,
}
if instruction_role:
args["instruction_role"] = instruction_role
@@ -271,7 +274,6 @@ class AzureOpenAIConfigBase(OpenAIHandler):
"total_tokens",
"api_type",
"org_id",
"ai_model_type",
"service_id",
"client",
},
python/packages/azure/tests/test_azure_assistants_client.py
+9 -6
View File
@@ -13,6 +13,7 @@ from agent_framework import (
TextContent,
)
from agent_framework.exceptions import ServiceInitializationError
from azure.identity import DefaultAzureCredential
from pydantic import Field
from agent_framework_azure import AzureAssistantsClient
@@ -259,7 +260,7 @@ def get_weather(
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response() -> None:
"""Test Azure Assistants Client response."""
async with AzureAssistantsClient() as azure_assistants_client:
async with AzureAssistantsClient(ad_credential=DefaultAzureCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClient)
messages: list[ChatMessage] = []
@@ -283,7 +284,7 @@ async def test_azure_assistants_client_get_response() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response_tools() -> None:
"""Test Azure Assistants Client response with tools."""
async with AzureAssistantsClient() as azure_assistants_client:
async with AzureAssistantsClient(ad_credential=DefaultAzureCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClient)
messages: list[ChatMessage] = []
@@ -304,7 +305,7 @@ async def test_azure_assistants_client_get_response_tools() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming() -> None:
"""Test Azure Assistants Client streaming response."""
async with AzureAssistantsClient() as azure_assistants_client:
async with AzureAssistantsClient(ad_credential=DefaultAzureCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClient)
messages: list[ChatMessage] = []
@@ -334,7 +335,7 @@ async def test_azure_assistants_client_streaming() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming_tools() -> None:
"""Test Azure Assistants Client streaming response with tools."""
async with AzureAssistantsClient() as azure_assistants_client:
async with AzureAssistantsClient(ad_credential=DefaultAzureCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClient)
messages: list[ChatMessage] = []
@@ -361,14 +362,16 @@ async def test_azure_assistants_client_streaming_tools() -> None:
async def test_azure_assistants_client_with_existing_assistant() -> None:
"""Test Azure Assistants Client with existing assistant ID."""
# First create an assistant to use in the test
async with AzureAssistantsClient() as temp_client:
async with AzureAssistantsClient(ad_credential=DefaultAzureCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
# Now test using the existing assistant
async with AzureAssistantsClient(assistant_id=assistant_id) as azure_assistants_client:
async with AzureAssistantsClient(
assistant_id=assistant_id, ad_credential=DefaultAzureCredential()
) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClient)
assert azure_assistants_client.assistant_id == assistant_id
python/packages/azure/tests/test_azure_chat_client.py
+35 -80
View File
@@ -12,8 +12,6 @@ from agent_framework import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
FunctionResultContent,
TextContent,
ai_function,
)
@@ -23,6 +21,7 @@ from agent_framework.openai import (
OpenAIContentFilterException,
)
from agent_framework.telemetry import USER_AGENT_KEY
from azure.identity import DefaultAzureCredential
from httpx import Request, Response
from openai import AsyncAzureOpenAI, AsyncStream
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
@@ -123,6 +122,7 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
"default_headers": default_headers,
"env_file_path": "test.env",
}
azure_chat_client = AzureChatClient.from_dict(settings)
@@ -258,13 +258,15 @@ async def test_azure_on_your_data(
content="test",
role="assistant",
context={ # type: ignore
"citations": {
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
},
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
},
),
@@ -298,11 +300,11 @@ async def test_azure_on_your_data(
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 3
assert isinstance(content.messages[0].contents[0], FunctionCallContent)
assert isinstance(content.messages[0].contents[1], FunctionResultContent)
assert isinstance(content.messages[0].contents[2], TextContent)
assert content.messages[0].contents[2].text == "test"
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
@@ -326,13 +328,15 @@ async def test_azure_on_your_data_string(
content="test",
role="assistant",
context=json.dumps({ # type: ignore
"citations": {
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
},
"citations": [
{
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
}
],
"intent": "query used",
}),
),
@@ -366,11 +370,11 @@ async def test_azure_on_your_data_string(
additional_properties={"extra_body": expected_data_settings},
)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 3
assert isinstance(content.messages[0].contents[0], FunctionCallContent)
assert isinstance(content.messages[0].contents[1], FunctionResultContent)
assert isinstance(content.messages[0].contents[2], TextContent)
assert content.messages[0].contents[2].text == "test"
assert len(content.messages[0].contents) == 1
assert isinstance(content.messages[0].contents[0], TextContent)
assert len(content.messages[0].contents[0].annotations) == 1
assert content.messages[0].contents[0].annotations[0].title == "test title"
assert content.messages[0].contents[0].text == "test"
mock_create.assert_awaited_once_with(
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
@@ -437,55 +441,6 @@ async def test_azure_on_your_data_fail(
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_azure_on_your_data_split_messages(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
Choice(
index=0,
message=ChatCompletionMessage(
content="test",
role="assistant",
context={ # type: ignore
"citations": {
"content": "test content",
"title": "test title",
"url": "test url",
"filepath": "test filepath",
"chunk_id": "test chunk_id",
},
"intent": "query used",
},
),
finish_reason="stop",
)
]
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
azure_chat_client = AzureChatClient()
content = await azure_chat_client.get_response(
messages=messages_in,
)
message = azure_chat_client._split_message(content)
assert len(content.messages) == 1
assert len(content.messages[0].contents) == 3
assert isinstance(content.messages[0].contents[0], FunctionCallContent)
assert isinstance(content.messages[0].contents[1], FunctionResultContent)
assert isinstance(content.messages[0].contents[2], TextContent)
assert content.messages[0].contents[2].text == "test"
assert message.messages[0].contents == [content.messages[0].contents[0]]
CONTENT_FILTERED_ERROR_MESSAGE = (
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy. Please "
"modify your prompt and retry. To learn more about our content filtering policies please read our "
@@ -607,7 +562,7 @@ async def test_bad_request_non_content_filter(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_streaming(
async def test_get_streaming(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
@@ -646,7 +601,7 @@ def get_story_text() -> str:
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureChatClient()
azure_chat_client = AzureChatClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -672,7 +627,7 @@ async def test_azure_openai_chat_client_response() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_response_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureChatClient()
azure_chat_client = AzureChatClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -693,7 +648,7 @@ async def test_azure_openai_chat_client_response_tools() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureChatClient()
azure_chat_client = AzureChatClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -725,7 +680,7 @@ async def test_azure_openai_chat_client_streaming() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureChatClient()
azure_chat_client = AzureChatClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_chat_client, ChatClient)
messages: list[ChatMessage] = []
python/packages/azure/tests/test_azure_responses_client.py
+36 -39
View File
@@ -6,7 +6,8 @@ from typing import Annotated
import pytest
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
from agent_framework.azure import AzureResponsesClient
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.exceptions import ServiceInitializationError
from azure.identity import DefaultAzureCredential
from pydantic import BaseModel
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
@@ -104,7 +105,7 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response() -> None:
"""Test azure responses client responses."""
azure_responses_client = AzureResponsesClient()
azure_responses_client = AzureResponsesClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_responses_client, ChatClient)
@@ -147,7 +148,7 @@ async def test_azure_responses_client_response() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_response_tools() -> None:
"""Test azure responses client tools."""
azure_responses_client = AzureResponsesClient()
azure_responses_client = AzureResponsesClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_responses_client, ChatClient)
@@ -186,7 +187,7 @@ async def test_azure_responses_client_response_tools() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming() -> None:
"""Test Azure azure responses client streaming responses."""
azure_responses_client = AzureResponsesClient()
azure_responses_client = AzureResponsesClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_responses_client, ChatClient)
@@ -219,29 +220,27 @@ async def test_azure_responses_client_streaming() -> None:
messages.append(ChatMessage(role="user", text="The weather in Seattle is sunny"))
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
# This is currently broken. See https://github.com/azure/azure-python/issues/2305
with pytest.raises(ServiceResponseException):
response = azure_responses_client.get_streaming_response(
messages=messages,
response_format=OutputStruct,
)
full_message = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
response = azure_responses_client.get_streaming_response(
messages=messages,
response_format=OutputStruct,
)
full_message = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()
@skip_if_azure_integration_tests_disabled
async def test_azure_responses_client_streaming_tools() -> None:
"""Test azure responses client streaming tools."""
azure_responses_client = AzureResponsesClient()
azure_responses_client = AzureResponsesClient(ad_credential=DefaultAzureCredential())
assert isinstance(azure_responses_client, ChatClient)
@@ -266,22 +265,20 @@ async def test_azure_responses_client_streaming_tools() -> None:
messages.clear()
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
# This is currently broken. See https://github.com/azure/azure-python/issues/2305
with pytest.raises(ServiceResponseException):
response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
full_message = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
response = azure_responses_client.get_streaming_response(
messages=messages,
tools=[get_weather],
tool_choice="auto",
response_format=OutputStruct,
)
full_message = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if isinstance(content, TextContent) and content.text:
full_message += content.text
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather.lower()