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()
python/packages/foundry/agent_framework_foundry/_chat_client.py
+31 -32
View File
@@ -4,7 +4,7 @@ import contextlib
import json
import sys
from collections.abc import AsyncIterable, MutableMapping, MutableSequence
from typing import Any, ClassVar
from typing import Any, ClassVar, TypeVar
from agent_framework import (
AIContents,
@@ -26,7 +26,6 @@ from agent_framework import (
UsageDetails,
use_tool_calling,
)
from agent_framework._clients import ai_function_to_json_schema_spec
from agent_framework._pydantic import AFBaseSettings
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.telemetry import use_telemetry
@@ -93,6 +92,9 @@ class FoundrySettings(AFBaseSettings):
agent_name: str | None = "UnnamedAgent"
TFoundryChatClient = TypeVar("TFoundryChatClient", bound="FoundryChatClient")
@use_telemetry
@use_tool_calling
class FoundryChatClient(ChatClientBase):
@@ -102,21 +104,22 @@ class FoundryChatClient(ChatClientBase):
client: AIProjectClient = Field(...)
credential: AsyncTokenCredential | None = Field(...)
agent_id: str | None = Field(default=None)
agent_name: str | None = Field(default=None)
ai_model_deployment_name: str | None = Field(default=None)
thread_id: str | None = Field(default=None)
_should_delete_agent: bool = PrivateAttr(default=False) # Track whether we should delete the agent
_should_close_client: bool = PrivateAttr(default=False) # Track whether we should close client connection
_should_close_credential: bool = PrivateAttr(default=False) # Track whether we should close credential
_foundry_settings: FoundrySettings = PrivateAttr()
def __init__(
self,
*,
client: AIProjectClient | None = None,
agent_id: str | None = None,
agent_name: str | None = None,
thread_id: str | None = None,
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
async_ad_credential: AsyncTokenCredential | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
@@ -133,8 +136,7 @@ class FoundryChatClient(ChatClientBase):
conversation_id property, when making a request.
project_endpoint: The Azure AI Foundry project endpoint URL. Used if client is not provided.
model_deployment_name: The model deployment name to use for agent creation.
credential: Azure async credential to use for authentication. If not provided,
DefaultAzureCredential will be used.
async_ad_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.
@@ -152,7 +154,6 @@ class FoundryChatClient(ChatClientBase):
# If no client is provided, create one
should_close_client = False
should_close_credential = False
if client is None:
if not foundry_settings.project_endpoint:
raise ServiceInitializationError("Project endpoint is required when client is not provided.")
@@ -161,27 +162,20 @@ class FoundryChatClient(ChatClientBase):
raise ServiceInitializationError("Model deployment name is required for agent creation.")
# Use provided credential or fallback to DefaultAzureCredential
if credential is None:
from azure.identity.aio import DefaultAzureCredential
credential = DefaultAzureCredential()
should_close_credential = True
client = AIProjectClient(endpoint=foundry_settings.project_endpoint, credential=credential)
should_close_client = True
if not async_ad_credential:
raise ServiceInitializationError("Azure AD credential is required when client is not provided.")
client = AIProjectClient(endpoint=foundry_settings.project_endpoint, credential=async_ad_credential)
super().__init__(
client=client, # type: ignore[reportCallIssue]
credential=credential, # type: ignore[reportCallIssue]
credential=async_ad_credential, # type: ignore[reportCallIssue]
agent_id=agent_id, # type: ignore[reportCallIssue]
thread_id=thread_id, # type: ignore[reportCallIssue]
agent_name=foundry_settings.agent_name, # type: ignore[reportCallIssue]
ai_model_deployment_name=foundry_settings.model_deployment_name, # type: ignore[reportCallIssue]
**kwargs,
)
self._should_delete_agent = False
self._should_close_client = should_close_client
self._should_close_credential = should_close_credential
self._foundry_settings = foundry_settings
async def __aenter__(self) -> "Self":
"""Async context manager entry."""
@@ -195,10 +189,9 @@ class FoundryChatClient(ChatClientBase):
"""Close the client and clean up any agents we created."""
await self._cleanup_agent_if_needed()
await self._close_client_if_needed()
await self._close_credential_if_needed()
@classmethod
def from_dict(cls, settings: dict[str, Any]) -> "FoundryChatClient":
def from_dict(cls: type[TFoundryChatClient], settings: dict[str, Any]) -> TFoundryChatClient:
"""Initialize a FoundryChatClient from a dictionary of settings.
Args:
@@ -262,11 +255,11 @@ class FoundryChatClient(ChatClientBase):
"""
# If no agent_id is provided, create a temporary agent
if self.agent_id is None:
if not self._foundry_settings.model_deployment_name:
if not self.ai_model_deployment_name:
raise ServiceInitializationError("Model deployment name is required for agent creation.")
agent_name = self._foundry_settings.agent_name
args = {"model": self._foundry_settings.model_deployment_name, "name": agent_name}
agent_name = self.agent_name
args = {"model": self.ai_model_deployment_name, "name": agent_name}
if run_options:
if "tools" in run_options:
args["tools"] = run_options["tools"]
@@ -459,11 +452,6 @@ class FoundryChatClient(ChatClientBase):
return contents
async def _close_credential_if_needed(self) -> None:
"""Close credential if we created it."""
if self._should_close_credential and self.credential is not None:
await self.credential.close()
async def _close_client_if_needed(self) -> None:
"""Close client session if we created it."""
if self._should_close_client:
@@ -496,7 +484,7 @@ class FoundryChatClient(ChatClientBase):
if chat_options.tool_choice != "none" and chat_options.tools is not None:
for tool in chat_options.tools:
if isinstance(tool, AIFunction):
tool_definitions.append(ai_function_to_json_schema_spec(tool)) # type: ignore[reportUnknownArgumentType]
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
elif isinstance(tool, HostedCodeInterpreterTool):
tool_definitions.append(CodeInterpreterToolDefinition())
elif isinstance(tool, MutableMapping):
@@ -605,3 +593,14 @@ class FoundryChatClient(ChatClientBase):
tool_outputs.append(ToolOutput(tool_call_id=call_id, output=str(function_result_content.result)))
return run_id, tool_outputs
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
python/packages/foundry/tests/test_foundry_chat_client.py
+7 -5
View File
@@ -15,6 +15,7 @@ from agent_framework import (
TextContent,
)
from agent_framework.exceptions import ServiceInitializationError
from azure.identity.aio import DefaultAzureCredential
from pydantic import Field
from agent_framework_foundry import FoundryChatClient, FoundrySettings
@@ -44,7 +45,8 @@ def create_test_foundry_chat_client(
agent_id=agent_id,
thread_id=thread_id,
_should_delete_agent=should_delete_agent,
_foundry_settings=foundry_settings,
agent_name=foundry_settings.agent_name, # type: ignore[reportCallIssue]
ai_model_deployment_name=foundry_settings.model_deployment_name, # type:
credential=None,
)
@@ -304,7 +306,7 @@ def get_weather(
@skip_if_foundry_integration_tests_disabled
async def test_foundry_chat_client_get_response() -> None:
"""Test Foundry Chat Client response."""
async with FoundryChatClient() as foundry_chat_client:
async with FoundryChatClient(async_ad_credential=DefaultAzureCredential()) as foundry_chat_client:
assert isinstance(foundry_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -328,7 +330,7 @@ async def test_foundry_chat_client_get_response() -> None:
@skip_if_foundry_integration_tests_disabled
async def test_foundry_chat_client_get_response_tools() -> None:
"""Test Foundry Chat Client response with tools."""
async with FoundryChatClient() as foundry_chat_client:
async with FoundryChatClient(async_ad_credential=DefaultAzureCredential()) as foundry_chat_client:
assert isinstance(foundry_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -349,7 +351,7 @@ async def test_foundry_chat_client_get_response_tools() -> None:
@skip_if_foundry_integration_tests_disabled
async def test_foundry_chat_client_streaming() -> None:
"""Test Foundry Chat Client streaming response."""
async with FoundryChatClient() as foundry_chat_client:
async with FoundryChatClient(async_ad_credential=DefaultAzureCredential()) as foundry_chat_client:
assert isinstance(foundry_chat_client, ChatClient)
messages: list[ChatMessage] = []
@@ -379,7 +381,7 @@ async def test_foundry_chat_client_streaming() -> None:
@skip_if_foundry_integration_tests_disabled
async def test_foundry_chat_client_streaming_tools() -> None:
"""Test Foundry Chat Client streaming response with tools."""
async with FoundryChatClient() as foundry_chat_client:
async with FoundryChatClient(async_ad_credential=DefaultAzureCredential()) as foundry_chat_client:
assert isinstance(foundry_chat_client, ChatClient)
messages: list[ChatMessage] = []
python/packages/main/agent_framework/_agents.py
+14 -1
View File
@@ -381,6 +381,8 @@ class ChatClientAgent(AgentBase):
kwargs: any additional keyword arguments.
Unused, can be used by subclasses of this Agent.
"""
kwargs.update(additional_properties or {})
args: dict[str, Any] = {
"chat_client": chat_client,
"chat_options": ChatOptions(
@@ -399,7 +401,7 @@ class ChatClientAgent(AgentBase):
tools=tools, # type: ignore
top_p=top_p,
user=user,
additional_properties=additional_properties or {},
additional_properties=kwargs,
),
}
if instructions is not None:
@@ -412,6 +414,7 @@ class ChatClientAgent(AgentBase):
args["id"] = id
super().__init__(**args)
self._update_agent_name()
async def __aenter__(self) -> "Self":
"""Async context manager entry.
@@ -430,6 +433,16 @@ class ChatClientAgent(AgentBase):
if isinstance(self.chat_client, AbstractAsyncContextManager):
await self.chat_client.__aexit__(exc_type, exc_val, exc_tb) # type: ignore[reportUnknownMemberType]
def _update_agent_name(self) -> None:
"""Update the agent name in a chat client.
Checks if there is a agent name, the implementation
should check if there is already a agent name defined, and if not
set it to this value.
"""
if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name): # type: ignore[reportAttributeAccessIssue]
self.chat_client._update_agent_name(self.name) # type: ignore[reportAttributeAccessIssue]
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
python/packages/main/agent_framework/_clients.py
+16 -27
View File
@@ -73,18 +73,6 @@ async def _auto_invoke_function(
)
def ai_function_to_json_schema_spec(function: AIFunction[BaseModel, Any]) -> dict[str, Any]:
"""Convert a AIFunction to the JSON Schema function specification format."""
return {
"type": "function",
"function": {
"name": function.name,
"description": function.description,
"parameters": function.parameters(),
},
}
def _tool_call_non_streaming(
func: Callable[..., Awaitable["ChatResponse"]],
) -> Callable[..., Awaitable["ChatResponse"]]:
@@ -117,14 +105,15 @@ def _tool_call_non_streaming(
_auto_invoke_function(
function_call,
custom_args=kwargs,
tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
tool_map={t.name: t for t in chat_options.tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
sequence_index=seq_idx,
request_index=attempt_idx,
)
for seq_idx, function_call in enumerate(function_calls)
])
# add a single ChatMessage to the response with the results
response.messages.append(ChatMessage(role="tool", contents=results))
result_message = ChatMessage(role="tool", contents=results)
response.messages.append(result_message)
# response should contain 2 messages after this,
# one with function call contents
# and one with function result contents
@@ -133,7 +122,11 @@ def _tool_call_non_streaming(
# we need to keep track of all function call messages
fcc_messages.extend(response.messages)
# and add them as additional context to the messages
messages.extend(response.messages)
if chat_options.store:
messages.clear()
messages.append(result_message)
else:
messages.extend(response.messages)
continue
# If we reach this point, it means there were no function calls to handle,
# we'll add the previous function call and responses
@@ -146,7 +139,7 @@ def _tool_call_non_streaming(
# Failsafe: give up on tools, ask model for plain answer
chat_options.tool_choice = "none"
self._prepare_tools_and_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
self._prepare_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
response = await func(self, messages=messages, chat_options=chat_options)
if fcc_messages:
for msg in reversed(fcc_messages):
@@ -202,7 +195,7 @@ def _tool_call_streaming(
_auto_invoke_function(
function_call,
custom_args=kwargs,
tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
tool_map={t.name: t for t in chat_options.tools or [] if isinstance(t, AIFunction)}, # type: ignore[reportPrivateUsage]
sequence_index=seq_idx,
request_index=attempt_idx,
)
@@ -216,7 +209,7 @@ def _tool_call_streaming(
# Failsafe: give up on tools, ask model for plain answer
chat_options.tool_choice = "none"
self._prepare_tools_and_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
self._prepare_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
async for update in func(self, messages=messages, chat_options=chat_options, **kwargs):
yield update
@@ -529,7 +522,7 @@ class ChatClientBase(AFBaseModel, ABC):
additional_properties=additional_properties or {},
)
prepped_messages = self._prepare_messages(messages)
self._prepare_tools_and_tool_choice(chat_options=chat_options)
self._prepare_tool_choice(chat_options=chat_options)
return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **kwargs)
async def get_streaming_response(
@@ -610,13 +603,13 @@ class ChatClientBase(AFBaseModel, ABC):
**kwargs,
)
prepped_messages = self._prepare_messages(messages)
self._prepare_tools_and_tool_choice(chat_options=chat_options)
self._prepare_tool_choice(chat_options=chat_options)
async for update in self._inner_get_streaming_response(
messages=prepped_messages, chat_options=chat_options, **kwargs
):
yield update
def _prepare_tools_and_tool_choice(self, chat_options: ChatOptions) -> None:
def _prepare_tool_choice(self, chat_options: ChatOptions) -> None:
"""Prepare the tools and tool choice for the chat options.
This function should be overridden by subclasses to customize tool handling.
@@ -627,10 +620,6 @@ class ChatClientBase(AFBaseModel, ABC):
chat_options.tools = None
chat_options.tool_choice = ChatToolMode.NONE.mode
return
chat_options.tools = [
(ai_function_to_json_schema_spec(t) if isinstance(t, AIFunction) else t) # type: ignore[reportUnknownArgumentType]
for t in chat_options._ai_tools or [] # type: ignore[reportPrivateUsage]
]
if not chat_options.tools:
chat_options.tool_choice = ChatToolMode.NONE.mode
else:
@@ -647,8 +636,8 @@ class ChatClientBase(AFBaseModel, ABC):
def create_agent(
self,
*,
name: str,
instructions: str,
name: str | None = None,
instructions: str | None = None,
tools: AITool
| list[AITool]
| Callable[..., Any]
python/packages/main/agent_framework/_tools.py
+140 -112
View File
@@ -7,9 +7,10 @@ from time import perf_counter
from typing import TYPE_CHECKING, Annotated, Any, Generic, Protocol, TypeVar, get_args, get_origin, runtime_checkable
from opentelemetry import metrics, trace
from pydantic import BaseModel, Field, create_model
from pydantic import BaseModel, Field, PrivateAttr, create_model
from ._logging import get_logger
from ._pydantic import AFBaseModel
from .telemetry import GenAIAttributes, start_as_current_span
if TYPE_CHECKING:
@@ -22,6 +23,48 @@ logger = get_logger()
__all__ = ["AIFunction", "AITool", "HostedCodeInterpreterTool", "ai_function"]
def _parse_inputs(
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None",
) -> list["AIContents"]:
"""Parse the inputs for a tool, ensuring they are of type AIContents."""
if inputs is None:
return []
from ._types import AIContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent
parsed_inputs: list["AIContents"] = []
if not isinstance(inputs, list):
inputs = [inputs]
for input_item in inputs:
if isinstance(input_item, str):
# If it's a string, we assume it's a URI or similar identifier.
# Convert it to a UriContent or similar type as needed.
parsed_inputs.append(UriContent(uri=input_item, media_type="text/plain"))
elif isinstance(input_item, dict):
# If it's a dict, we assume it contains properties for a specific content type.
# we check if the required keys are present to determine the type.
# for instance, if it has "uri" and "media_type", we treat it as UriContent.
# if is only has uri, then we treat it as DataContent.
# etc.
if "uri" in input_item:
parsed_inputs.append(
UriContent(**input_item) if "media_type" in input_item else DataContent(**input_item)
)
elif "file_id" in input_item:
parsed_inputs.append(HostedFileContent(**input_item))
elif "vector_store_id" in input_item:
parsed_inputs.append(HostedVectorStoreContent(**input_item))
elif "data" in input_item:
parsed_inputs.append(DataContent(**input_item))
else:
raise ValueError(f"Unsupported input type: {input_item}")
elif isinstance(input_item, AIContent):
parsed_inputs.append(input_item)
else:
raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected AIContents or dict.")
return parsed_inputs
@runtime_checkable
class AITool(Protocol):
"""Represents a generic tool that can be specified to an AI service.
@@ -37,9 +80,9 @@ class AITool(Protocol):
name: str
"""The name of the tool."""
description: str | None = None
description: str
"""A description of the tool, suitable for use in describing the purpose to a model."""
additional_properties: dict[str, Any] | None = None
additional_properties: dict[str, Any] | None
"""Additional properties associated with the tool."""
def __str__(self) -> str:
@@ -51,48 +94,96 @@ ArgsT = TypeVar("ArgsT", bound=BaseModel)
ReturnT = TypeVar("ReturnT")
class AIFunction(AITool, Generic[ArgsT, ReturnT]):
"""A AITool that is callable as code."""
class AIToolBase(AFBaseModel):
"""Base class for AI tools, providing common attributes and methods.
Args:
name: The name of the tool.
description: A description of the tool.
additional_properties: Additional properties associated with the tool.
"""
name: str = Field(..., kw_only=False)
description: str = ""
additional_properties: dict[str, Any] | None = None
def __str__(self) -> str:
"""Return a string representation of the tool."""
if self.description:
return f"{self.__class__.__name__}(name={self.name}, description={self.description})"
return f"{self.__class__.__name__}(name={self.name})"
class HostedCodeInterpreterTool(AIToolBase):
"""Represents a hosted tool that can be specified to an AI service to enable it to execute generated code.
This tool does not implement code interpretation itself. It serves as a marker to inform a service
that it is allowed to execute generated code if the service is capable of doing so.
"""
inputs: list[Any] = Field(default_factory=list)
def __init__(
self,
func: Callable[..., Awaitable[ReturnT] | ReturnT],
name: str,
description: str,
input_model: type[ArgsT],
*,
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None,
description: str | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
):
"""Initialize a FunctionTool.
) -> None:
"""Initialize the HostedCodeInterpreterTool.
Args:
func: The function to wrap.
name: The name of the tool.
inputs: A list of contents that the tool can accept as input. Defaults to None.
This should mostly be HostedFileContent or HostedVectorStoreContent.
Can also be DataContent, depending on the service used.
When supplying a list, it can contain:
- AIContents instances
- dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- strings (which will be converted to UriContent with media_type "text/plain").
If None, defaults to an empty list.
description: A description of the tool.
input_model: A Pydantic model that defines the input parameters for the function.
**kwargs: Additional properties to set on the tool.
stored in additional_properties.
additional_properties: Additional properties associated with the tool.
**kwargs: Additional keyword arguments to pass to the base class.
"""
self.name = name
self.description = description
self.input_model = input_model
self.additional_properties: dict[str, Any] | None = kwargs
self._func = func
self.invocation_duration_histogram = meter.create_histogram(
"agent_framework.function.invocation.duration",
args: dict[str, Any] = {
"name": "code_interpreter",
}
if inputs:
args["inputs"] = _parse_inputs(inputs)
if description is not None:
args["description"] = description
if additional_properties is not None:
args["additional_properties"] = additional_properties
if "name" in kwargs:
raise ValueError("The 'name' argument is reserved for the HostedCodeInterpreterTool and cannot be set.")
super().__init__(**args, **kwargs)
class AIFunction(AIToolBase, Generic[ArgsT, ReturnT]):
"""A AITool that is callable as code.
Args:
name: The name of the function.
description: A description of the function.
additional_properties: Additional properties to set on the function.
func: The function to wrap. If None, returns a decorator.
input_model: The Pydantic model that defines the input parameters for the function.
"""
func: Callable[..., Awaitable[ReturnT] | ReturnT]
input_model: type[ArgsT]
_invocation_duration_histogram: metrics.Histogram = PrivateAttr(
default_factory=lambda: meter.create_histogram(
GenAIAttributes.MEASUREMENT_FUNCTION_INVOCATION_DURATION.value,
unit="s",
description="Measures the duration of a function's execution",
)
def parameters(self) -> dict[str, Any]:
"""Return the parameter json schemas of the input model."""
return self.input_model.model_json_schema()
)
def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
"""Call the wrapped function with the provided arguments."""
return self._func(*args, **kwargs)
def __str__(self) -> str:
return f"AIFunction(name={self.name}, description={self.description})"
return self.func(*args, **kwargs)
async def invoke(
self,
@@ -136,9 +227,24 @@ class AIFunction(AITool, Generic[ArgsT, ReturnT]):
raise
finally:
duration = perf_counter() - starting_time_stamp
self.invocation_duration_histogram.record(duration, attributes=attributes)
self._invocation_duration_histogram.record(duration, attributes=attributes)
logger.info("Function completed. Duration: %fs", duration)
def parameters(self) -> dict[str, Any]:
"""Create the json schema of the parameters."""
return self.input_model.model_json_schema()
def to_json_schema_spec(self) -> dict[str, Any]:
"""Convert a AIFunction to the JSON Schema function specification format."""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters(),
},
}
def _parse_annotation(annotation: Any) -> Any:
"""Parse a type annotation and return the corresponding type.
@@ -207,91 +313,13 @@ def ai_function(
raise TypeError(f"Input model for {tool_name} must be a subclass of BaseModel, got {input_model}")
return AIFunction[Any, ReturnT](
func=f,
name=tool_name,
description=tool_desc,
additional_properties=additional_properties or {},
func=f,
input_model=input_model,
**(additional_properties if additional_properties is not None else {}),
)
return wrapper(func)
return decorator(func) if func else decorator # type: ignore[reportReturnType, return-value]
def _parse_inputs(
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None",
) -> list["AIContents"]:
"""Parse the inputs for a tool, ensuring they are of type AIContents."""
if inputs is None:
return []
from ._types import AIContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent
parsed_inputs: list["AIContents"] = []
if not isinstance(inputs, list):
inputs = [inputs]
for input_item in inputs:
if isinstance(input_item, str):
# If it's a string, we assume it's a URI or similar identifier.
# Convert it to a UriContent or similar type as needed.
parsed_inputs.append(UriContent(uri=input_item, media_type="text/plain"))
elif isinstance(input_item, dict):
# If it's a dict, we assume it contains properties for a specific content type.
# we check if the required keys are present to determine the type.
if "uri" in input_item:
parsed_inputs.append(
UriContent(**input_item) if "media_type" in input_item else DataContent(**input_item)
)
elif "file_id" in input_item:
parsed_inputs.append(HostedFileContent(**input_item))
elif "vector_store_id" in input_item:
parsed_inputs.append(HostedVectorStoreContent(**input_item))
elif "data" in input_item:
parsed_inputs.append(DataContent(**input_item))
else:
raise ValueError(f"Unsupported input type: {input_item}")
elif isinstance(input_item, AIContent):
parsed_inputs.append(input_item)
else:
raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected AIContents or dict.")
return parsed_inputs
class HostedCodeInterpreterTool(AITool):
"""Represents a hosted tool that can be specified to an AI service to enable it to execute generated code.
This tool does not implement code interpretation itself. It serves as a marker to inform a service
that it is allowed to execute generated code if the service is capable of doing so.
"""
def __init__(
self,
name: str = "code_interpreter",
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None,
description: str | None = None,
additional_properties: dict[str, Any] | None = None,
):
"""Initialize a HostedCodeInterpreterTool.
Args:
name: The name of the tool. Defaults to "code_interpreter".
inputs: A list of contents that the tool can accept as input. Defaults to None.
This should mostly be HostedFileContent or HostedVectorStoreContent.
Can also be DataContent, depending on the service used.
When supplying a list, it can contain:
- AIContents instances
- dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
- strings (which will be converted to UriContent with media_type "text/plain").
If None, defaults to an empty list.
description: A description of the tool.
additional_properties: Additional properties associated with the tool, specific to the service used.
"""
self.name = name
self.inputs = _parse_inputs(inputs)
self.description = description
self.additional_properties = additional_properties
def __str__(self) -> str:
"""Return a string representation of the tool."""
return f"HostedCodeInterpreterTool(name={self.name})"
python/packages/main/agent_framework/_types.py
+25 -26
View File
@@ -21,11 +21,9 @@ from pydantic import (
BaseModel,
ConfigDict,
Field,
PrivateAttr,
ValidationError,
field_validator,
model_serializer,
model_validator,
)
from ._pydantic import AFBaseModel
@@ -81,7 +79,6 @@ __all__ = [
"AIAnnotations",
"AIContent",
"AIContents",
"AITool",
"AgentRunResponse",
"AgentRunResponseUpdate",
"AnnotatedRegion",
@@ -159,6 +156,10 @@ class UsageDetails(AFBaseModel):
**kwargs,
)
def __str__(self) -> str:
"""Returns a string representation of the usage details."""
return self.model_dump_json(indent=4, exclude_none=True)
@property
def additional_counts(self) -> dict[str, int]:
"""Represents well-known additional counts for usage. This is not an exhaustive list.
@@ -173,6 +174,14 @@ class UsageDetails(AFBaseModel):
"""
return self.model_extra or {}
def __setitem__(self, key: str, value: int) -> None:
"""Sets an additional count for the usage details."""
if not isinstance(value, int):
raise ValueError("Additional counts must be integers.")
if self.model_extra is None:
self.model_extra = {} # type: ignore[reportAttributeAccessIssue, misc]
self.model_extra[key] = value
def __add__(self, other: "UsageDetails | None") -> "UsageDetails":
"""Combines two `UsageDetails` instances."""
if not other:
@@ -394,7 +403,7 @@ class AIContent(AFBaseModel):
type: Literal["ai"] = "ai"
annotations: list[AIAnnotations] | None = None
additional_properties: dict[str, Any] | None = None
raw_representation: Any | None = Field(default=None, repr=False)
raw_representation: Any | None = Field(default=None, repr=False, exclude=True)
class TextContent(AIContent):
@@ -1130,7 +1139,7 @@ class ChatRole(AFBaseModel):
TOOL: The role that provides additional information and references in response to tool use requests.
"""
value: str
value: str = Field(..., kw_only=False)
SYSTEM: ClassVar[Self] # type: ignore[assignment]
"""The role that instructs or sets the behaviour of the AI system."""
@@ -1211,7 +1220,7 @@ class ChatMessage(AFBaseModel):
"""The ID of the chat message."""
additional_properties: dict[str, Any] | None = None
"""Any additional properties associated with the chat message."""
raw_representation: Any | None = None
raw_representation: Any | None = Field(default=None, exclude=True)
"""The raw representation of the chat message from an underlying implementation."""
@overload
@@ -1760,13 +1769,6 @@ class ChatOptions(AFBaseModel):
tools: list[AITool | MutableMapping[str, Any]] | None = None
top_p: Annotated[float | None, Field(ge=0.0, le=1.0)] = None
user: str | None = None
_ai_tools: list[AITool | MutableMapping[str, Any]] | None = PrivateAttr(default=None)
@model_validator(mode="after")
def _copy_to_ai_tools(self) -> Self:
if self.tools and not self._ai_tools:
self._ai_tools = self.tools
return self
@field_validator("tools", mode="before")
@classmethod
@@ -1782,10 +1784,7 @@ class ChatOptions(AFBaseModel):
| None
),
) -> list[AITool | MutableMapping[str, Any]] | None:
"""Parse the tools field.
All tools are stored in both tools and _ai_tools.
"""
"""Parse the tools field."""
if not tools:
return None
if not isinstance(tools, list):
@@ -1849,22 +1848,22 @@ class ChatOptions(AFBaseModel):
"""
if not isinstance(other, ChatOptions):
return self
ai_tools = other._ai_tools
updated_values = other.model_dump(exclude_none=True)
updated_values.pop("tools", [])
other_tools = other.tools
updated_values = other.model_dump(exclude_none=True, exclude={"tools"})
logit_bias = updated_values.pop("logit_bias", {})
metadata = updated_values.pop("metadata", {})
additional_properties = updated_values.pop("additional_properties", {})
combined = self.model_copy(update=updated_values)
if ai_tools:
if not combined._ai_tools:
combined._ai_tools = []
for tool in ai_tools:
if tool not in combined._ai_tools:
combined._ai_tools.append(tool)
combined.logit_bias = {**(combined.logit_bias or {}), **logit_bias}
combined.metadata = {**(combined.metadata or {}), **metadata}
combined.additional_properties = {**(combined.additional_properties or {}), **additional_properties}
if other_tools:
if not combined.tools:
combined.tools = other_tools
else:
for tool in other_tools:
if tool not in combined.tools:
combined.tools.append(tool)
return combined
python/packages/main/agent_framework/exceptions.py
+28 -2
View File
@@ -1,10 +1,24 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from typing import Any
logger = logging.getLogger("agent_framework")
class AgentFrameworkException(Exception):
"""Base class for exceptions in the Agent Framework."""
"""Base exceptions for the Agent Framework.
pass
Automatically logs the message as debug.
"""
def __init__(self, message: str, inner_exception: Exception | None = None, *args: Any, **kwargs: Any):
"""Create an AgentFrameworkException.
This emits a debug log, with the inner_exception if provided.
"""
logger.debug(message, exc_info=inner_exception)
super().__init__(message, *args, **kwargs) # type: ignore
class AgentException(AgentFrameworkException):
@@ -68,3 +82,15 @@ class ServiceInvalidResponseError(ServiceResponseException):
"""An error occurred while validating the response from the service."""
pass
class ToolException(AgentFrameworkException):
"""An error occurred while executing a tool."""
pass
class ToolExecutionException(ToolException):
"""An error occurred while executing a tool."""
pass
python/packages/main/agent_framework/openai/_assistants_client.py
+14 -4
View File
@@ -3,7 +3,7 @@
import json
import sys
from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence
from typing import Any, ClassVar
from typing import Any
from openai import AsyncOpenAI
from openai.types.beta.threads import (
@@ -20,7 +20,7 @@ from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
from openai.types.beta.threads.runs import RunStep
from pydantic import Field, PrivateAttr, SecretStr, ValidationError
from .._clients import ChatClientBase, ai_function_to_json_schema_spec, use_tool_calling
from .._clients import ChatClientBase, use_tool_calling
from .._tools import AIFunction, HostedCodeInterpreterTool
from .._types import (
AIContents,
@@ -54,7 +54,6 @@ __all__ = ["OpenAIAssistantsClient"]
class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
"""OpenAI Assistants client."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
assistant_id: str | None = Field(default=None)
assistant_name: str | None = Field(default=None)
thread_id: str | None = Field(default=None)
@@ -362,7 +361,7 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
if chat_options.tool_choice != "none" and chat_options.tools is not None:
for tool in chat_options.tools:
if isinstance(tool, AIFunction):
tool_definitions.append(ai_function_to_json_schema_spec(tool)) # type: ignore[reportUnknownArgumentType]
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
elif isinstance(tool, HostedCodeInterpreterTool):
tool_definitions.append({"type": "code_interpreter"})
elif isinstance(tool, MutableMapping):
@@ -467,3 +466,14 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
tool_outputs.append(ToolOutput(tool_call_id=call_id, output=str(function_result_content.result)))
return run_id, tool_outputs
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.assistant_name:
self.assistant_name = agent_name
python/packages/main/agent_framework/openai/_chat_client.py
+157 -92
View File
@@ -1,20 +1,24 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from collections.abc import AsyncIterable, Mapping, MutableSequence, Sequence
from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence, Sequence
from datetime import datetime
from itertools import chain
from typing import Any, ClassVar, cast
from typing import Any, TypeVar
from openai import AsyncOpenAI, AsyncStream
from openai import AsyncOpenAI, BadRequestError
from openai.lib._parsing._completions import type_to_response_format_param
from openai.types import CompletionUsage
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDeltaToolCall
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall
from pydantic import SecretStr, ValidationError
from openai.types.chat.chat_completion_message_custom_tool_call import ChatCompletionMessageCustomToolCall
from pydantic import BaseModel, SecretStr, ValidationError
from agent_framework import AIFunction, AITool, UsageContent
from .._clients import ChatClientBase, use_tool_calling
from .._logging import get_logger
from .._types import (
AIContents,
ChatFinishReason,
@@ -28,12 +32,19 @@ from .._types import (
TextContent,
UsageDetails,
)
from ..exceptions import ServiceInitializationError, ServiceInvalidResponseError
from ..exceptions import (
ServiceInitializationError,
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..telemetry import use_telemetry
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAIModelTypes, OpenAISettings
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAISettings, prepare_function_call_results
__all__ = ["OpenAIChatClient"]
logger = get_logger("agent_framework.openai")
# region Base Client
@use_telemetry
@@ -41,8 +52,6 @@ __all__ = ["OpenAIChatClient"]
class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
"""OpenAI Chat completion class."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
async def _inner_get_response(
self,
*,
@@ -50,16 +59,24 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
chat_options.additional_properties["stream"] = False
if not chat_options.ai_model_id:
chat_options.ai_model_id = self.ai_model_id
response = await self._send_request(chat_options, messages=self._prepare_chat_history_for_request(messages))
assert isinstance(response, ChatCompletion) # nosec # noqa: S101
response_metadata = self._get_metadata_from_chat_response(response)
return next(
self._create_chat_message_content(response, choice, response_metadata) for choice in response.choices
)
options_dict = self._prepare_options(messages, chat_options)
try:
return self._create_chat_response(await self.client.chat.completions.create(stream=False, **options_dict))
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
async def _inner_get_streaming_response(
self,
@@ -68,91 +85,138 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
chat_options.additional_properties["stream"] = True
chat_options.additional_properties["stream_options"] = {"include_usage": True}
chat_options.ai_model_id = chat_options.ai_model_id or self.ai_model_id
response = await self._send_request(chat_options, messages=self._prepare_chat_history_for_request(messages))
if not isinstance(response, AsyncStream):
raise ServiceInvalidResponseError("Expected an AsyncStream[ChatCompletionChunk] response.")
async for chunk in response:
assert isinstance(chunk, ChatCompletionChunk) # nosec # noqa: S101
if len(chunk.choices) == 0 and chunk.usage is None:
continue
assert isinstance(chunk, ChatCompletionChunk) # nosec # noqa: S101
chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk)
if chunk.usage is not None:
# Usage is contained in the last chunk where the choices are empty
# We are duplicating the usage metadata to all the choices in the response
yield ChatResponseUpdate(
role=ChatRole.ASSISTANT,
contents=[],
ai_model_id=chat_options.ai_model_id,
additional_properties=chunk_metadata,
)
else:
yield next(
self._create_streaming_chat_message_content(chunk, choice, chunk_metadata)
for choice in chunk.choices
)
options_dict = self._prepare_options(messages, chat_options)
options_dict["stream_options"] = {"include_usage": True}
try:
async for chunk in await self.client.chat.completions.create(stream=True, **options_dict):
if len(chunk.choices) == 0 and chunk.usage is None:
continue
yield self._create_chat_response_update(chunk)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
# region content creation
def _create_chat_message_content(
self, response: ChatCompletion, choice: Choice, response_metadata: dict[str, Any]
) -> "ChatResponse":
"""Create a chat message content object from a choice."""
metadata = self._get_metadata_from_chat_choice(choice)
metadata.update(response_metadata)
items: MutableSequence[ChatMessage] = []
if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_chat_choice(choice)]:
items.append(ChatMessage(role="assistant", contents=parsed_tool_calls))
if choice.message.content:
items.append(ChatMessage(role="assistant", text=choice.message.content))
elif hasattr(choice.message, "refusal") and choice.message.refusal:
items.append(ChatMessage(role="assistant", text=choice.message.refusal))
def _chat_to_tool_spec(self, tools: list[AITool | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
chat_tools: list[dict[str, Any]] = []
for tool in tools:
if isinstance(tool, AITool):
match tool:
case AIFunction():
chat_tools.append(tool.to_json_schema_spec())
case _:
logger.debug("Unsupported tool passed (type: %s), ignoring", type(tool))
else:
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
return chat_tools
def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
options_dict = chat_options.to_provider_settings()
if messages and "messages" not in options_dict:
options_dict["messages"] = self._prepare_chat_history_for_request(messages)
if "messages" not in options_dict:
raise ServiceInvalidRequestError("Messages are required for chat completions")
if chat_options.tools is None:
options_dict.pop("parallel_tool_calls", None)
else:
options_dict["tools"] = self._chat_to_tool_spec(chat_options.tools)
if "model" not in options_dict:
options_dict["model"] = self.ai_model_id
if (
chat_options.response_format
and isinstance(chat_options.response_format, type)
and issubclass(chat_options.response_format, BaseModel)
):
options_dict["response_format"] = type_to_response_format_param(chat_options.response_format)
return options_dict
def _create_chat_response(self, response: ChatCompletion) -> "ChatResponse":
"""Create a chat message content object from a choice."""
response_metadata = self._get_metadata_from_chat_response(response)
messages: list[ChatMessage] = []
finish_reason: ChatFinishReason | None = None
for choice in response.choices:
response_metadata.update(self._get_metadata_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = ChatFinishReason(value=choice.finish_reason)
contents: list[AIContents] = []
if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_chat_choice(choice)]:
contents.extend(parsed_tool_calls)
if text_content := self._parse_text_from_choice(choice):
contents.append(text_content)
messages.append(ChatMessage(role="assistant", contents=contents))
return ChatResponse(
response_id=response.id,
created_at=datetime.fromtimestamp(response.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
usage_details=self._usage_details_from_openai(response.usage) if response.usage else None,
messages=items,
model_id=self.ai_model_id,
additional_properties=metadata,
finish_reason=(ChatFinishReason(value=choice.finish_reason) if choice.finish_reason else None),
messages=messages,
model_id=response.model,
additional_properties=response_metadata,
finish_reason=finish_reason,
)
def _create_streaming_chat_message_content(
def _create_chat_response_update(
self,
chunk: ChatCompletionChunk,
choice: ChunkChoice,
chunk_metadata: dict[str, Any],
) -> ChatResponseUpdate:
"""Create a streaming chat message content object from a choice."""
metadata = self._get_metadata_from_chat_choice(choice)
metadata.update(chunk_metadata)
chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk)
if chunk.usage:
return ChatResponseUpdate(
role=ChatRole.ASSISTANT,
contents=[UsageContent(details=self._usage_details_from_openai(chunk.usage), raw_representation=chunk)],
ai_model_id=chunk.model,
additional_properties=chunk_metadata,
)
contents: list[AIContents] = []
finish_reason: ChatFinishReason | None = None
for choice in chunk.choices:
chunk_metadata.update(self._get_metadata_from_chat_choice(choice))
contents.extend(self._get_tool_calls_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = ChatFinishReason(value=choice.finish_reason)
items: list[Any] = self._get_tool_calls_from_chat_choice(choice)
if choice.delta and choice.delta.content is not None:
items.append(TextContent(text=choice.delta.content))
if text_content := self._parse_text_from_choice(choice):
contents.append(text_content)
return ChatResponseUpdate(
created_at=datetime.fromtimestamp(chunk.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
contents=items,
contents=contents,
role=ChatRole.ASSISTANT,
ai_model_id=self.ai_model_id,
additional_properties=metadata,
finish_reason=(ChatFinishReason(value=choice.finish_reason) if choice.finish_reason else None),
ai_model_id=chunk.model,
additional_properties=chunk_metadata,
finish_reason=finish_reason,
raw_representation=chunk,
)
def _usage_details_from_openai(self, usage: CompletionUsage) -> UsageDetails | None:
def _usage_details_from_openai(self, usage: CompletionUsage) -> UsageDetails:
return UsageDetails(
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
)
def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
"""Parse the choice into a TextContent object."""
message = choice.message if isinstance(choice, Choice) else choice.delta
if message.content:
return TextContent(text=message.content, raw_representation=choice)
if hasattr(message, "refusal") and message.refusal:
return TextContent(text=message.refusal, raw_representation=choice)
return None
def _get_metadata_from_chat_response(self, response: ChatCompletion) -> dict[str, Any]:
"""Get metadata from a chat response."""
return {
@@ -175,13 +239,15 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
"""Get tool calls from a chat choice."""
resp: list[AIContents] = []
content = choice.message if isinstance(choice, Choice) else choice.delta
if content and (tool_calls := getattr(content, "tool_calls", None)) is not None:
for tool in cast(list[ChatCompletionMessageToolCall] | list[ChoiceDeltaToolCall], tool_calls):
if tool.function: # type: ignore[reportAttributeAccessIssue, union-attr]
if content and content.tool_calls:
for tool in content.tool_calls:
if not isinstance(tool, ChatCompletionMessageCustomToolCall) and tool.function:
# ignoring tool.custom
fcc = FunctionCallContent(
call_id=tool.id if tool.id else "",
name=tool.function.name if tool.function.name else "", # type: ignore
arguments=tool.function.arguments if tool.function.arguments else "", # type: ignore
name=tool.function.name if tool.function.name else "",
arguments=tool.function.arguments if tool.function.arguments else "",
raw_representation=tool.function,
)
resp.append(fcc)
@@ -236,7 +302,8 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
args["tool_calls"] = [self._openai_content_parser(content)] # type: ignore
case FunctionResultContent():
args["tool_call_id"] = content.call_id
args["content"] = content.result
if content.result:
args["content"] = prepare_function_call_results(content.result)
case _:
if "content" not in args:
args["content"] = []
@@ -275,6 +342,8 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
# region Public client
TOpenAIChatClient = TypeVar("TOpenAIChatClient", bound="OpenAIChatClient")
class OpenAIChatClient(OpenAIConfigBase, OpenAIChatClientBase):
"""OpenAI Chat completion class."""
@@ -328,23 +397,19 @@ class OpenAIChatClient(OpenAIConfigBase, OpenAIChatClientBase):
ai_model_id=openai_settings.chat_model_id,
api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None,
org_id=openai_settings.org_id,
ai_model_type=OpenAIModelTypes.CHAT,
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls, settings: dict[str, Any]) -> "OpenAIChatClient":
"""Initialize an Open AI service from a dictionary of settings.
def from_dict(cls: type[TOpenAIChatClient], settings: dict[str, Any]) -> TOpenAIChatClient:
"""Initialize an Open AI Chat Client from a dictionary of settings.
Args:
settings: A dictionary of settings for the service.
"""
return OpenAIChatClient(
ai_model_id=settings["ai_model_id"],
default_headers=settings.get("default_headers"),
)
return cls(**settings)
# endregion
python/packages/main/agent_framework/openai/_responses_client.py
+517 -256
View File
@@ -4,70 +4,87 @@ import sys
from collections.abc import AsyncIterable, Callable, Mapping, MutableMapping, MutableSequence, Sequence
from datetime import datetime
from itertools import chain
from typing import Any, ClassVar, Literal
from typing import TYPE_CHECKING, Any, Literal, TypeVar
from openai import AsyncOpenAI, AsyncStream
from openai import AsyncOpenAI, BadRequestError
from openai.types.responses.function_tool_param import FunctionToolParam
from openai.types.responses.parsed_response import (
ParsedResponse,
)
from openai.types.responses.response import Response as OpenAIResponse
from openai.types.responses.response_code_interpreter_tool_call import ResponseCodeInterpreterToolCall
from openai.types.responses.response_completed_event import ResponseCompletedEvent
from openai.types.responses.response_content_part_added_event import ResponseContentPartAddedEvent
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.response_includable import ResponseIncludable
from openai.types.responses.response_output_item import ResponseOutputItem
from openai.types.responses.response_output_message import ResponseOutputMessage
from openai.types.responses.response_function_call_arguments_delta_event import ResponseFunctionCallArgumentsDeltaEvent
from openai.types.responses.response_output_item_added_event import ResponseOutputItemAddedEvent
from openai.types.responses.response_output_refusal import ResponseOutputRefusal
from openai.types.responses.response_output_text import ResponseOutputText
from openai.types.responses.response_stream_event import ResponseStreamEvent as OpenAIResponseStreamEvent
from openai.types.responses.response_text_delta_event import ResponseTextDeltaEvent
from openai.types.responses.response_usage import ResponseUsage
from openai.types.responses.tool_param import (
CodeInterpreter,
CodeInterpreterContainerCodeInterpreterToolAuto,
ToolParam,
)
from pydantic import BaseModel, SecretStr, ValidationError
from agent_framework import DataContent, TextReasoningContent, UriContent, UsageContent
from .._clients import ChatClientBase, use_tool_calling
from .._tools import HostedCodeInterpreterTool
from .._logging import get_logger
from .._tools import AIFunction, AITool, HostedCodeInterpreterTool
from .._types import (
AIContents,
AITool,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
ChatRole,
ChatToolMode,
CitationAnnotation,
FunctionCallContent,
FunctionResultContent,
HostedFileContent,
StructuredResponse,
TextContent,
TextSpanRegion,
UsageDetails,
)
from ..exceptions import ServiceInitializationError, ServiceInvalidResponseError
from ..exceptions import (
ServiceInitializationError,
ServiceInvalidRequestError,
ServiceResponseException,
)
from ..telemetry import use_telemetry
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAIModelTypes, OpenAISettings
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAISettings, prepare_function_call_results
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
if TYPE_CHECKING:
from openai.types.responses.response_includable import ResponseIncludable
from .._types import ChatToolMode
logger = get_logger("agent_framework.openai")
__all__ = ["OpenAIResponsesClient"]
# region ResponsesClient
class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
def _filter_options(self, **kwargs: Any) -> dict[str, Any]:
"""Filter options for the responses call."""
# The responses call does not support all the options that the chat completion call does.
# We filter out the unsupported options.
return {key: value for key, value in kwargs.items() if value is not None}
"""Base class for all OpenAI Responses based API's."""
# The responses create call takes very different parameters than the chat completion call,
# so we override the get_response method to handle the specific parameters for responses.
@override
async def get_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
*,
# TODO(peterychang): enable this option. background: bool | None = None,
include: list[ResponseIncludable] | None = None,
include: list["ResponseIncludable"] | None = None,
instruction: str | None = None,
max_tokens: int | None = None,
parallel_tool_calls: bool | None = None,
@@ -79,7 +96,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
seed: int | None = None,
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tool_choice: "ChatToolMode" | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
| list[AITool]
| Callable[..., Any]
@@ -123,31 +140,27 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
Returns:
A chat response from the model.
"""
filtered_options = self._filter_options(
background=False,
return await super().get_response(
messages=messages,
include=include,
instruction=instruction,
max_tokens=max_tokens,
parallel_tool_calls=parallel_tool_calls,
model=model,
previous_response_id=previous_response_id,
reasoning=reasoning,
service_tier=service_tier,
truncation=truncation,
timeout=timeout,
)
filtered_options.update(additional_properties or {})
return await super().get_response(
messages=messages,
model=model,
max_tokens=max_tokens,
response_format=response_format,
seed=seed,
store=store,
temperature=temperature,
top_p=top_p,
tool_choice=tool_choice,
tools=tools,
top_p=top_p,
user=user,
additional_properties=filtered_options,
truncation=truncation,
timeout=timeout,
additional_properties=additional_properties,
**kwargs,
)
@@ -157,7 +170,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
messages: str | ChatMessage | list[str] | list[ChatMessage],
*,
# TODO(peterychang): enable this option. background: bool | None = None,
include: list[ResponseIncludable] | None = None,
include: list["ResponseIncludable"] | None = None,
instruction: str | None = None,
max_tokens: int | None = None,
parallel_tool_calls: bool | None = None,
@@ -169,7 +182,7 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
seed: int | None = None,
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tool_choice: "ChatToolMode" | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
tools: AITool
| list[AITool]
| Callable[..., Any]
@@ -213,55 +226,32 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
Returns:
A stream representing the response(s) from the LLM.
"""
filtered_options = self._filter_options(
background=False,
async for update in super().get_streaming_response(
messages=messages,
include=include,
instruction=instruction,
max_tokens=max_tokens,
parallel_tool_calls=parallel_tool_calls,
model=model,
previous_response_id=previous_response_id,
reasoning=reasoning,
service_tier=service_tier,
truncation=truncation,
timeout=timeout,
)
filtered_options.update(additional_properties or {})
async for update in super().get_streaming_response(
messages=messages,
model=model,
max_tokens=max_tokens,
response_format=response_format,
seed=seed,
store=store,
temperature=temperature,
top_p=top_p,
tool_choice=tool_choice,
tools=tools,
top_p=top_p,
user=user,
additional_properties=filtered_options,
truncation=truncation,
timeout=timeout,
additional_properties=additional_properties,
**kwargs,
):
yield update
def _chat_to_response_tool_spec(self, tools: list[AITool | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
response_tools: list[dict[str, Any]] = []
for tool in tools:
if isinstance(tool, AITool):
# TODO(peterychang): Support AITools
if isinstance(tool, HostedCodeInterpreterTool):
response_tools.append({"type": "code_interpreter", "container": {"type": "auto"}})
continue
if "function" not in tool:
response_tools.append(tool if isinstance(tool, dict) else dict(tool))
parameters = {"additionalProperties": False}
parameters.update(tool.get("function", {}).get("parameters", {}))
response_tools.append({
"type": "function",
"name": tool.get("function", {}).get("name", ""),
"strict": True,
"description": tool.get("function", {}).get("description", None),
"parameters": parameters,
})
return response_tools
# region Inner Methods
async def _inner_get_response(
self,
@@ -270,16 +260,39 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
chat_options.additional_properties["stream"] = False
if not chat_options.ai_model_id:
chat_options.ai_model_id = self.ai_model_id
if chat_options.tools:
chat_options.additional_properties.update({
"response_tools": self._chat_to_response_tool_spec(chat_options.tools)
})
response = await self._send_request(chat_options, messages=self._prepare_chat_history_for_request(messages))
assert isinstance(response, OpenAIResponse) # nosec # noqa: S101
return next(self._create_response_content(response, item, store=chat_options.store) for item in response.output)
options_dict = self._prepare_options(messages, chat_options)
try:
if not chat_options.response_format:
response = await self.client.responses.create(
stream=False,
**options_dict,
)
chat_options.conversation_id = response.id if chat_options.store is True else None
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(
text_format=resp_format,
stream=False,
**options_dict,
)
chat_options.conversation_id = parsed_response.id if chat_options.store is True else None
return self._create_response_content(parsed_response, chat_options=chat_options)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt, with exception: {ex}",
inner_exception=ex,
) from ex
async def _inner_get_streaming_response(
self,
@@ -288,169 +301,113 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
chat_options.additional_properties["stream"] = True
chat_options.ai_model_id = chat_options.ai_model_id or self.ai_model_id
options_dict = 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(
stream=True,
**options_dict,
)
async for chunk in response:
update = self._create_streaming_response_content(
chunk, chat_options=chat_options, function_call_ids=function_call_ids
)
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(
text_format=chat_options.response_format,
**options_dict,
) as response:
async for chunk in response:
update = self._create_streaming_response_content(
chunk, chat_options=chat_options, function_call_ids=function_call_ids
)
yield update
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
inner_exception=ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
inner_exception=ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
inner_exception=ex,
) from ex
if chat_options.tools:
chat_options.additional_properties.update({
"response_tools": self._chat_to_response_tool_spec(chat_options.tools)
})
response = await self._send_request(chat_options, messages=self._prepare_chat_history_for_request(messages))
if not isinstance(response, AsyncStream):
raise ServiceInvalidResponseError("Expected an AsyncStream[ResponseStreamEvent] response.")
async for chunk in response:
update = self._create_streaming_response_content(chunk, store=chat_options.store) # type: ignore
if not update:
continue
yield update
# region Prep methods
def _create_response_content(
self, response: OpenAIResponse, item: ResponseOutputItem, store: bool | None
) -> "ChatResponse":
"""Create a chat message content object from a choice."""
items: MutableSequence[ChatMessage] = []
metadata: dict[str, Any] = response.metadata or {}
# TODO(peterychang): Add support for other content types
if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_response(response)]:
items.append(ChatMessage(role="assistant", contents=parsed_tool_calls))
if isinstance(item, ResponseOutputMessage):
for content in item.content:
# TODO(peterychang): Add annotations when available
if isinstance(content, ResponseOutputText):
items.append(ChatMessage(role=item.role, text=content.text))
metadata.update(self._get_metadata_from_response(content))
elif isinstance(content, ResponseOutputRefusal):
items.append(ChatMessage(role=item.role, text=content.refusal))
if isinstance(item, ResponseCodeInterpreterToolCall):
items.append(ChatMessage(role=ChatRole.ASSISTANT, text=response.output_text))
return ChatResponse(
response_id=response.id,
conversation_id=response.id if store is True else None,
created_at=datetime.fromtimestamp(response.created_at).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
usage_details=self._usage_details_from_openai(response.usage) if response.usage else None,
messages=items,
model_id=self.ai_model_id,
additional_properties=metadata,
raw_representation=response,
)
def _chat_to_response_tool_spec(
self, tools: list[AITool | MutableMapping[str, Any]]
) -> list[ToolParam | dict[str, Any]]:
response_tools: list[ToolParam | dict[str, Any]] = []
for tool in tools:
if isinstance(tool, AITool):
match tool:
case HostedCodeInterpreterTool():
tool_args: dict[str, Any] = {"type": "auto"}
if tool.inputs:
tool_args["file_ids"] = []
for tool_input in tool.inputs:
if isinstance(tool_input, HostedFileContent):
tool_args["file_ids"].append(tool_input.file_id)
if not tool_args["file_ids"]:
tool_args.pop("file_ids")
response_tools.append(
CodeInterpreter(
type="code_interpreter",
container=CodeInterpreterContainerCodeInterpreterToolAuto(**tool_args), # type: ignore[typeddict-item]
)
)
case AIFunction():
params = tool.parameters()
params["additionalProperties"] = False
response_tools.append(
FunctionToolParam(
name=tool.name,
parameters=params,
strict=False,
type="function",
description=tool.description,
)
)
case _:
logger.debug("Unsupported tool passed (type: %s)", type(tool))
else:
response_tools.append(tool if isinstance(tool, dict) else dict(tool))
return response_tools
def _create_streaming_response_content(
self, event: OpenAIResponseStreamEvent, store: bool | None
) -> ChatResponseUpdate | None:
"""Create a streaming chat message content object from a choice."""
metadata: dict[str, Any] = {}
items: list[AIContents] = []
conversation_id: str | None = None
# TODO(peterychang): Add support for other content types
if isinstance(event, ResponseContentPartAddedEvent):
if isinstance(event.part, ResponseOutputText):
items.append(TextContent(text=event.part.text))
metadata.update(self._get_metadata_from_response(event.part))
elif isinstance(event.part, ResponseOutputRefusal):
items.append(TextContent(text=event.part.refusal))
elif isinstance(event, ResponseTextDeltaEvent):
items.append(TextContent(text=event.delta))
metadata.update(self._get_metadata_from_response(event))
elif isinstance(event, ResponseCompletedEvent):
conversation_id = event.response.id if store is True else None
# Tool calls are available in the completed event
if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_response(event.response)]:
items.extend(parsed_tool_calls)
def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
"""Take ChatOptions and create the specific options for Responses."""
options_dict = chat_options.to_provider_settings(exclude={"response_format"})
# 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
# tools
if chat_options.tools is None:
options_dict.pop("parallel_tool_calls", None)
else:
return None
return ChatResponseUpdate(
contents=items,
conversation_id=conversation_id,
role=ChatRole.ASSISTANT,
ai_model_id=self.ai_model_id,
additional_properties=metadata,
raw_representation=event,
)
options_dict["tools"] = self._chat_to_response_tool_spec(chat_options.tools)
# other settings
if "store" not in options_dict:
options_dict["store"] = False
if "conversation_id" in options_dict:
options_dict["previous_response_id"] = options_dict["conversation_id"]
options_dict.pop("conversation_id")
if "model" not in options_dict:
options_dict["model"] = self.ai_model_id
return options_dict
def _get_tool_calls_from_response(self, response: OpenAIResponse) -> list[AIContents]:
resp: list[AIContents] = []
# TODO(peterychang): Support the other tool calls
for item in (i for i in response.output if isinstance(i, ResponseFunctionToolCall)):
fcc = FunctionCallContent(
call_id=item.id if item.id else "",
name=item.name,
arguments=item.arguments,
additional_properties={"call_id": item.call_id},
)
resp.append(fcc)
return resp
def _usage_details_from_openai(self, usage: ResponseUsage) -> UsageDetails | None:
return UsageDetails(
prompt_tokens=usage.input_tokens,
completion_tokens=usage.output_tokens,
total_tokens=usage.total_tokens,
)
def _openai_chat_message_parser(
self,
message: ChatMessage,
tool_id_to_call_id: dict[str, str],
) -> list[dict[str, Any]]:
"""Parse a chat message into the openai format."""
all_messages: list[dict[str, Any]] = []
args: dict[str, Any] = {
"role": message.role.value if isinstance(message.role, ChatRole) else message.role,
}
if message.additional_properties:
args["metadata"] = message.additional_properties
for content in message.contents:
match content:
case FunctionResultContent():
new_args: dict[str, Any] = {}
new_args.update(self._openai_content_parser(message.role, content, tool_id_to_call_id))
all_messages.append(new_args)
case FunctionCallContent():
function_call = self._openai_content_parser(message.role, content, tool_id_to_call_id)
all_messages.append(function_call) # type: ignore
case _:
if "content" not in args:
args["content"] = []
args["content"].append(self._openai_content_parser(message.role, content, tool_id_to_call_id)) # type: ignore
if "content" in args or "tool_calls" in args:
all_messages.append(args)
return all_messages
def _openai_content_parser(
self,
role: ChatRole,
content: AIContents,
tool_id_to_call_id: dict[str, str],
) -> dict[str, Any]:
"""Parse contents into the openai format."""
match content:
case FunctionCallContent():
return {
"id": content.call_id,
"call_id": tool_id_to_call_id[content.call_id],
"type": "function_call",
"name": content.name,
"arguments": content.arguments,
}
case FunctionResultContent():
# call_id for the result needs to be the same as the call_id for the function call
return {
"call_id": tool_id_to_call_id[content.call_id],
"type": "function_call_output",
"output": content.result,
}
case TextContent():
return {
"type": "output_text" if role == ChatRole.ASSISTANT else "input_text",
"text": content.text,
}
# TODO(peterychang): We'll probably need to specialize the other content types as well
case _:
return content.model_dump(exclude_none=True)
def _prepare_chat_history_for_request(self, chat_messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
"""Prepare the chat history for a request.
def _prepare_chat_messages_for_request(self, chat_messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
"""Prepare the chat messages for a request.
Allowing customization of the key names for role/author, and optionally overriding the role.
@@ -464,24 +421,336 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
chat_messages: The chat history to prepare.
Returns:
prepared_chat_history (Any): The prepared chat history for a request.
The prepared chat messages for a request.
"""
tool_id_to_call_id: dict[str, str] = {}
call_id_to_id: dict[str, str] = {}
for message in chat_messages:
for content in message.contents:
if isinstance(content, FunctionCallContent):
assert content.additional_properties and "call_id" in content.additional_properties # nosec # noqa: S101
call_id = content.additional_properties["call_id"]
tool_id_to_call_id[content.call_id] = call_id
list_of_list = [self._openai_chat_message_parser(message, tool_id_to_call_id) for message in chat_messages]
if (
isinstance(content, FunctionCallContent)
and content.additional_properties
and "fc_id" in content.additional_properties
):
call_id_to_id[content.call_id] = content.additional_properties["fc_id"]
list_of_list = [self._openai_chat_message_parser(message, call_id_to_id) for message in chat_messages]
# Flatten the list of lists into a single list
return list(chain.from_iterable(list_of_list))
# region Response creation methods
def _create_response_content(
self,
response: OpenAIResponse | ParsedResponse[BaseModel],
chat_options: ChatOptions,
) -> "ChatResponse":
"""Create a chat message content object from a choice."""
structured_response: BaseModel | None = response.output_parsed if isinstance(response, ParsedResponse) else None # type: ignore[reportUnknownMemberType]
metadata: dict[str, Any] = response.metadata or {}
contents: list[AIContents] = []
for item in response.output: # type: ignore[reportUnknownMemberType]
match item.type:
# types:
# ParsedResponseOutputMessage[Unknown] |
# ParsedResponseFunctionToolCall |
# ResponseFileSearchToolCall |
# ResponseFunctionWebSearch |
# ResponseComputerToolCall |
# ResponseReasoningItem |
# McpCall |
# McpApprovalRequest |
# ImageGenerationCall |
# LocalShellCall |
# LocalShellCallAction |
# McpListTools |
# ResponseCodeInterpreterToolCall |
# ResponseCustomToolCall |
# ParsedResponseOutputMessage[BaseModel] |
# ResponseOutputMessage |
# ResponseFunctionToolCall
case "message": # ResponseOutputMessage
for message_content in item.content: # type: ignore[reportMissingTypeArgument]
match message_content.type:
case "output_text":
text_content = TextContent(
text=message_content.text, raw_representation=message_content
)
metadata.update(self._get_metadata_from_response(message_content))
if message_content.annotations:
text_content.annotations = []
for annotation in message_content.annotations:
match annotation.type:
case "file_path":
text_content.annotations.append(
CitationAnnotation(
file_id=annotation.file_id,
additional_properties={
"index": annotation.index,
},
raw_representation=annotation,
)
)
case "file_citation":
text_content.annotations.append(
CitationAnnotation(
url=annotation.filename,
file_id=annotation.file_id,
raw_representation=annotation,
additional_properties={
"index": annotation.index,
},
)
)
case "url_citation":
text_content.annotations.append(
CitationAnnotation(
title=annotation.title,
url=annotation.url,
annotated_regions=[
TextSpanRegion(
start_index=annotation.start_index,
end_index=annotation.end_index,
)
],
raw_representation=annotation,
)
)
case "container_file_citation":
text_content.annotations.append(
CitationAnnotation(
file_id=annotation.file_id,
url=annotation.filename,
additional_properties={
"container_id": annotation.container_id,
},
annotated_regions=[
TextSpanRegion(
start_index=annotation.start_index,
end_index=annotation.end_index,
)
],
raw_representation=annotation,
)
)
case _:
logger.debug("Unparsed annotation type: %s", annotation.type)
contents.append(text_content)
case "refusal":
contents.append(
TextContent(text=message_content.refusal, raw_representation=message_content)
)
case "reasoning": # ResponseOutputReasoning
if item.content:
for index, reasoning_content in enumerate(item.content):
additional_properties = None
if item.summary and index < len(item.summary):
additional_properties = {"summary": item.summary[index]}
contents.append(
TextReasoningContent(
text=reasoning_content.text,
raw_representation=reasoning_content,
additional_properties=additional_properties,
)
)
case "code_interpreter_call": # ResponseOutputCodeInterpreterCall
if item.outputs:
for code_output in item.outputs:
if code_output.type == "logs":
contents.append(TextContent(text=code_output.logs, raw_representation=item))
if code_output.type == "image":
contents.append(
UriContent(
uri=code_output.url,
raw_representation=item,
# no more specific media type then this can be inferred
media_type="image",
)
)
elif item.code:
# fallback if no output was returned is the code:
contents.append(TextContent(text=item.code, raw_representation=item))
case "function_call": # ResponseOutputFunctionCall
contents.append(
FunctionCallContent(
call_id=item.call_id if item.call_id else "",
name=item.name,
arguments=item.arguments,
additional_properties={"fc_id": item.id},
raw_representation=item,
)
)
case "image_generation_call": # ResponseOutputImageGenerationCall
if item.result:
contents.append(
DataContent(
uri=item.result,
raw_representation=item,
)
)
# TODO(peterychang): Add support for other content types
case _:
logger.debug("Unparsed content of type: %s: %s", item.type, item)
response_message = ChatMessage(role="assistant", contents=contents)
args: dict[str, Any] = {
"response_id": response.id,
"created_at": datetime.fromtimestamp(response.created_at).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
"messages": response_message,
"model_id": response.model,
"additional_properties": metadata,
"raw_representation": response,
}
if chat_options.store:
args["conversation_id"] = response.id
if response.usage and (usage_details := self._usage_details_from_openai(response.usage)):
args["usage_details"] = usage_details
if structured_response:
args["value"] = structured_response
return StructuredResponse(**args)
return ChatResponse(**args)
def _create_streaming_response_content(
self,
event: OpenAIResponseStreamEvent,
chat_options: ChatOptions,
function_call_ids: dict[int, tuple[str, str]],
) -> ChatResponseUpdate:
"""Create a streaming chat message content object from a choice."""
metadata: dict[str, Any] = {}
items: list[AIContents] = []
conversation_id: str | None = None
model = self.ai_model_id
# TODO(peterychang): Add support for other content types
match event:
case ResponseContentPartAddedEvent():
match event.part:
case ResponseOutputText():
items.append(TextContent(text=event.part.text, raw_representation=event))
metadata.update(self._get_metadata_from_response(event.part))
case ResponseOutputRefusal():
items.append(TextContent(text=event.part.refusal, raw_representation=event))
case ResponseTextDeltaEvent():
items.append(TextContent(text=event.delta, raw_representation=event))
metadata.update(self._get_metadata_from_response(event))
case ResponseCompletedEvent():
conversation_id = event.response.id if chat_options.store is True else None
model = event.response.model
if event.response.usage:
usage = self._usage_details_from_openai(event.response.usage)
if usage:
items.append(UsageContent(details=usage, raw_representation=event))
case ResponseOutputItemAddedEvent():
if event.item.type == "function_call":
function_call_ids[event.output_index] = (event.item.call_id, event.item.name)
case ResponseFunctionCallArgumentsDeltaEvent():
call_id, name = function_call_ids.get(event.output_index, (None, None))
if call_id and name:
items.append(
FunctionCallContent(
call_id=call_id,
name=name,
arguments=event.delta,
additional_properties={"output_index": event.output_index, "fc_id": event.item_id},
raw_representation=event,
)
)
case _:
logger.debug("Unparsed event: %s", event)
return ChatResponseUpdate(
contents=items,
conversation_id=conversation_id,
role=ChatRole.ASSISTANT,
ai_model_id=model,
additional_properties=metadata,
raw_representation=event,
)
def _usage_details_from_openai(self, usage: ResponseUsage) -> UsageDetails | None:
details = UsageDetails(
input_token_count=usage.input_tokens,
output_token_count=usage.output_tokens,
total_token_count=usage.total_tokens,
)
if usage.input_tokens_details and usage.input_tokens_details.cached_tokens:
details["openai.cached_input_tokens"] = usage.input_tokens_details.cached_tokens
if usage.output_tokens_details and usage.output_tokens_details.reasoning_tokens:
details["openai.reasoning_tokens"] = usage.output_tokens_details.reasoning_tokens
return details
def _openai_chat_message_parser(
self,
message: ChatMessage,
call_id_to_id: dict[str, str],
) -> list[dict[str, Any]]:
"""Parse a chat message into the openai format."""
all_messages: list[dict[str, Any]] = []
args: dict[str, Any] = {
"role": message.role.value if isinstance(message.role, ChatRole) else message.role,
}
if message.additional_properties:
args["metadata"] = message.additional_properties
for content in message.contents:
match content:
case FunctionResultContent():
new_args: dict[str, Any] = {}
new_args.update(self._openai_content_parser(message.role, content, call_id_to_id))
all_messages.append(new_args)
case FunctionCallContent():
function_call = self._openai_content_parser(message.role, content, call_id_to_id)
all_messages.append(function_call) # type: ignore
case _:
if "content" not in args:
args["content"] = []
args["content"].append(self._openai_content_parser(message.role, content, call_id_to_id)) # type: ignore
if "content" in args or "tool_calls" in args:
all_messages.append(args)
return all_messages
def _openai_content_parser(
self,
role: ChatRole,
content: AIContents,
call_id_to_id: dict[str, str],
) -> dict[str, Any]:
"""Parse contents into the openai format."""
match content:
case FunctionCallContent():
return {
"call_id": content.call_id,
"id": call_id_to_id[content.call_id],
"type": "function_call",
"name": content.name,
"arguments": content.arguments,
}
case FunctionResultContent():
# call_id for the result needs to be the same as the call_id for the function call
args: dict[str, Any] = {
"call_id": content.call_id,
"id": call_id_to_id.get(content.call_id),
"type": "function_call_output",
}
if content.result:
args["output"] = prepare_function_call_results(content.result)
return args
case TextContent():
return {
"type": "output_text" if role == ChatRole.ASSISTANT else "input_text",
"text": content.text,
}
# TODO(peterychang): We'll probably need to specialize the other content types as well
case _:
return content.model_dump(exclude_none=True)
def _get_metadata_from_response(self, output: Any) -> dict[str, Any]:
"""Get metadata from a chat choice."""
return {
"logprobs": getattr(output, "logprobs", None),
}
if logprobs := getattr(output, "logprobs", None):
return {
"logprobs": logprobs,
}
return {}
TOpenAIResponsesClient = TypeVar("TOpenAIResponsesClient", bound="OpenAIResponsesClient")
@use_telemetry
@@ -489,8 +758,6 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
class OpenAIResponsesClient(OpenAIConfigBase, OpenAIResponsesClientBase):
"""OpenAI Responses client class."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
def __init__(
self,
ai_model_id: str | None = None,
@@ -540,25 +807,19 @@ class OpenAIResponsesClient(OpenAIConfigBase, OpenAIResponsesClientBase):
ai_model_id=openai_settings.responses_model_id,
api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None,
org_id=openai_settings.org_id,
ai_model_type=OpenAIModelTypes.RESPONSE,
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
)
@classmethod
def from_dict(cls, settings: dict[str, Any]) -> "OpenAIResponsesClient":
def from_dict(cls: type[TOpenAIResponsesClient], settings: dict[str, Any]) -> TOpenAIResponsesClient:
"""Initialize an Open AI service from a dictionary of settings.
Args:
settings: A dictionary of settings for the service.
"""
return OpenAIResponsesClient(
ai_model_id=settings["ai_model_id"],
default_headers=settings.get("default_headers"),
api_key=settings.get("api_key"),
org_id=settings.get("org_id"),
)
return cls(**settings)
# endregion
python/packages/main/agent_framework/openai/_shared.py
+25 -174
View File
@@ -1,19 +1,16 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
from abc import ABC
from collections.abc import Mapping
from copy import copy
from enum import Enum
from typing import Annotated, Any, ClassVar, Union
from openai import (
AsyncOpenAI,
AsyncStream,
BadRequestError,
_legacy_response, # type: ignore
)
from openai.lib._parsing._completions import type_to_response_format_param
from openai.types import Completion
from openai.types.audio import Transcription
from openai.types.chat import ChatCompletion, ChatCompletionChunk
@@ -25,10 +22,9 @@ from pydantic.types import StringConstraints
from .._logging import get_logger
from .._pydantic import AFBaseModel, AFBaseSettings
from .._types import ChatOptions, SpeechToTextOptions, TextToSpeechOptions
from ..exceptions import ServiceInitializationError, ServiceInvalidRequestError, ServiceResponseException
from .._types import AIContents, ChatOptions, SpeechToTextOptions, TextToSpeechOptions
from ..exceptions import ServiceInitializationError
from ..telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from ._exceptions import OpenAIContentFilterException
logger: logging.Logger = get_logger("agent_framework.openai")
@@ -46,11 +42,7 @@ RESPONSE_TYPE = Union[
_legacy_response.HttpxBinaryResponseContent,
]
OPTION_TYPE = Union[
ChatOptions,
SpeechToTextOptions,
TextToSpeechOptions,
]
OPTION_TYPE = Union[ChatOptions, SpeechToTextOptions, TextToSpeechOptions, dict[str, Any]]
__all__ = [
@@ -58,6 +50,23 @@ __all__ = [
]
def prepare_function_call_results(content: AIContents | Any | list[AIContents | Any]) -> str | list[str]:
"""Prepare the values of the function call results."""
if isinstance(content, list):
results: list[str] = []
for item in content:
res = prepare_function_call_results(item)
if isinstance(res, list):
results.extend(res)
else:
results.append(res)
return results[0] if len(results) == 1 else results
if isinstance(content, BaseModel):
return content.model_dump_json(exclude_none=True, exclude={"raw_representation", "additional_properties"})
# fallback
return json.dumps(content)
class OpenAISettings(AFBaseSettings):
"""OpenAI environment settings.
@@ -108,177 +117,23 @@ class OpenAISettings(AFBaseSettings):
realtime_model_id: str | None = None
class OpenAIModelTypes(Enum):
"""OpenAI model types, can be text, chat or embedding."""
CHAT = "chat"
EMBEDDING = "embedding"
TEXT_TO_IMAGE = "text-to-image"
SPEECH_TO_TEXT = "speech-to-text"
TEXT_TO_SPEECH = "text-to-speech"
REALTIME = "realtime"
RESPONSE = "response"
class OpenAIHandler(AFBaseModel, ABC):
"""Internal class for calls to OpenAI API's."""
class OpenAIHandler(AFBaseModel):
"""Base class for OpenAI Clients."""
client: AsyncOpenAI
ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
ai_model_type: OpenAIModelTypes = OpenAIModelTypes.CHAT
async def _send_request(self, options: OPTION_TYPE, messages: list[dict[str, Any]] | None = None) -> RESPONSE_TYPE:
"""Send a request to the OpenAI API."""
if self.ai_model_type == OpenAIModelTypes.CHAT:
assert isinstance(options, ChatOptions) # nosec # noqa: S101
return await self._send_completion_request(options, messages)
# TODO(evmattso): move other PromptExecutionSettings to a common options class
if self.ai_model_type == OpenAIModelTypes.EMBEDDING:
raise NotImplementedError("Embedding generation is not yet implemented in OpenAIHandler")
if self.ai_model_type == OpenAIModelTypes.TEXT_TO_IMAGE:
raise NotImplementedError("Text to image generation is not yet implemented in OpenAIHandler")
if self.ai_model_type == OpenAIModelTypes.SPEECH_TO_TEXT:
assert isinstance(options, SpeechToTextOptions) # nosec # noqa: S101
return await self._send_audio_to_text_request(options)
if self.ai_model_type == OpenAIModelTypes.TEXT_TO_SPEECH:
assert isinstance(options, TextToSpeechOptions) # nosec # noqa: S101
return await self._send_text_to_audio_request(options)
if self.ai_model_type == OpenAIModelTypes.RESPONSE:
assert isinstance(options, ChatOptions) # nosec # noqa: S101
return await self._send_response_request(options, messages)
raise NotImplementedError(f"Model type {self.ai_model_type} is not supported")
async def _send_completion_request(
self,
chat_options: "ChatOptions",
messages: list[dict[str, Any]] | None = None,
) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
"""Execute the appropriate call to OpenAI models."""
try:
options_dict = chat_options.to_provider_settings()
if messages and "messages" not in options_dict:
options_dict["messages"] = messages
if "messages" not in options_dict:
raise ServiceInvalidRequestError("Messages are required for chat completions")
self._handle_structured_outputs(chat_options, options_dict)
if chat_options.tools is None:
options_dict.pop("parallel_tool_calls", None)
return await self.client.chat.completions.create(**options_dict) # type: ignore
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
async def _send_audio_to_text_request(self, options: SpeechToTextOptions) -> Transcription:
"""Send a request to the OpenAI audio to text endpoint."""
if not options.additional_properties["filename"]:
raise ServiceInvalidRequestError("Audio file is required for audio to text service")
try:
# TODO(peterychang): open isn't async safe
with open(options.additional_properties["filename"], "rb") as audio_file: # noqa: ASYNC230
return await self.client.audio.transcriptions.create( # type: ignore
file=audio_file,
**options.to_provider_settings(exclude={"filename"}),
)
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to transcribe audio",
ex,
) from ex
async def _send_text_to_audio_request(
self, options: TextToSpeechOptions
) -> _legacy_response.HttpxBinaryResponseContent:
"""Send a request to the OpenAI text to audio endpoint.
The OpenAI API returns the content of the generated audio file.
"""
try:
return await self.client.audio.speech.create(
**options.to_provider_settings(),
)
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to generate audio",
ex,
) from ex
async def _send_response_request(
self,
chat_options: "ChatOptions",
messages: list[dict[str, Any]] | None = None,
) -> Response | AsyncStream[ResponseStreamEvent]:
try:
options_dict = chat_options.to_provider_settings()
if messages and "input" not in options_dict:
options_dict["input"] = messages
if "input" not in options_dict:
raise ServiceInvalidRequestError("Messages are required for chat completions")
if chat_options.tools is None:
options_dict.pop("parallel_tool_calls", None)
else:
options_dict["tools"] = options_dict["response_tools"]
options_dict.pop("response_tools", None)
if chat_options.response_format:
# create call does not support response_format, so we need to handle it via parse call
resp_format = options_dict.pop("response_format", None)
return await self.client.responses.parse(
**options_dict,
text_format=resp_format,
)
if "store" not in options_dict:
options_dict["store"] = False
if "conversation_id" in options_dict:
options_dict["previous_response_id"] = options_dict["conversation_id"]
options_dict.pop("conversation_id")
return await self.client.responses.create(**options_dict) # type: ignore
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error",
ex,
) from ex
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
except Exception as ex:
raise ServiceResponseException(
f"{type(self)} service failed to complete the prompt",
ex,
) from ex
def _handle_structured_outputs(self, chat_options: "ChatOptions", options_dict: dict[str, Any]) -> None:
if (
chat_options.response_format
and isinstance(chat_options.response_format, type)
and issubclass(chat_options.response_format, BaseModel)
):
options_dict["response_format"] = type_to_response_format_param(chat_options.response_format)
class OpenAIConfigBase(OpenAIHandler):
"""Internal class for configuring a connection to an OpenAI service."""
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def __init__(
self,
ai_model_id: str = Field(min_length=1),
api_key: str | None = Field(min_length=1),
ai_model_type: OpenAIModelTypes | None = OpenAIModelTypes.CHAT,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncOpenAI | None = None,
@@ -295,8 +150,6 @@ class OpenAIConfigBase(OpenAIHandler):
Default to a preset value.
api_key (str): OpenAI API key for authentication.
Must be non-empty. (Optional)
ai_model_type (OpenAIModelTypes): The type of OpenAI
model to interact with. Defaults to CHAT.
org_id (str): OpenAI organization ID. This is optional
unless the account belongs to multiple organizations.
default_headers (Mapping[str, str]): Default headers
@@ -324,7 +177,6 @@ class OpenAIConfigBase(OpenAIHandler):
args = {
"ai_model_id": ai_model_id,
"client": client,
"ai_model_type": ai_model_type,
}
if instruction_role:
args["instruction_role"] = instruction_role
@@ -344,7 +196,6 @@ class OpenAIConfigBase(OpenAIHandler):
"completion_tokens",
"total_tokens",
"api_type",
"ai_model_type",
"client",
},
by_alias=True,
python/packages/main/agent_framework/telemetry.py
+1
View File
@@ -138,6 +138,7 @@ class GenAIAttributes(str, Enum):
# Agent Framework specific attributes
MEASUREMENT_FUNCTION_TAG_NAME = "agent_framework.function.name"
MEASUREMENT_FUNCTION_INVOCATION_DURATION = "agent_framework.function.invocation.duration"
AGENT_FRAMEWORK_GEN_AI_SYSTEM = "microsoft.agent_framework"
python/packages/main/tests/main/test_clients.py
-45
View File
@@ -305,48 +305,3 @@ async def test_base_client_with_streaming_function_calling_disabled(chat_client_
updates.append(update)
assert len(updates) == 1
assert exec_counter == 0
def test_chat_options_parsing_tools(chat_client_base, ai_function_tool) -> None:
"""Test that chat options can parse tools correctly."""
def echo() -> str:
"""Echo the input."""
return "Echo"
dict_function = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Retrieves current weather for the given location.",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"},
"units": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "Units the temperature will be returned in.",
},
},
"required": ["location", "units"],
"additionalProperties": False,
},
"strict": True,
},
}
options = ChatOptions(tools=[ai_function_tool, echo, dict_function], tool_choice="auto")
assert len(options.tools) == 3
assert options.tools[0] == ai_function_tool
assert options.tools[1] != echo
assert options.tools[2] == dict_function
# after prepare, the tools should be represented as dicts
# while ai_tools is still the same.
chat_client_base._prepare_tools_and_tool_choice(chat_options=options)
assert options._ai_tools[0] == ai_function_tool
assert options._ai_tools[2] == dict_function
assert len(options.tools) == 3
assert options.tools[0]["function"]["name"] == "simple_function"
assert options.tools[1]["function"]["name"] == "echo"
assert options.tools[2]["function"]["name"] == "get_weather"
python/packages/main/tests/main/test_tools.py
+7 -17
View File
@@ -95,7 +95,7 @@ async def test_ai_function_invoke_telemetry_enabled():
# Mock the histogram
mock_histogram = Mock()
telemetry_test_tool.invocation_duration_histogram = mock_histogram
telemetry_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke
result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
@@ -139,7 +139,7 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args():
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
pydantic_test_tool.invocation_duration_histogram = mock_histogram
pydantic_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke with Pydantic model
result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
@@ -172,7 +172,7 @@ async def test_ai_function_invoke_telemetry_with_exception():
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
exception_test_tool.invocation_duration_histogram = mock_histogram
exception_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke and expect exception
with pytest.raises(ValueError, match="Test exception for telemetry"):
@@ -212,7 +212,7 @@ async def test_ai_function_invoke_telemetry_async_function():
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
async_telemetry_test.invocation_duration_histogram = mock_histogram
async_telemetry_test._invocation_duration_histogram = mock_histogram
# Call invoke
result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
@@ -251,7 +251,7 @@ async def test_ai_function_invoke_telemetry_no_tool_call_id():
mock_start_span.return_value = mock_context_manager
mock_histogram = Mock()
no_id_test_tool.invocation_duration_histogram = mock_histogram
no_id_test_tool._invocation_duration_histogram = mock_histogram
# Call invoke without tool_call_id
result = await no_id_test_tool.invoke(x=5)
@@ -300,29 +300,19 @@ def test_hosted_code_interpreter_tool_default():
assert tool.name == "code_interpreter"
assert tool.inputs == []
assert tool.description is None
assert tool.description == ""
assert tool.additional_properties is None
assert str(tool) == "HostedCodeInterpreterTool(name=code_interpreter)"
def test_hosted_code_interpreter_tool_custom_name():
"""Test HostedCodeInterpreterTool with custom name."""
tool = HostedCodeInterpreterTool(name="custom_interpreter")
assert tool.name == "custom_interpreter"
assert tool.inputs == []
assert str(tool) == "HostedCodeInterpreterTool(name=custom_interpreter)"
def test_hosted_code_interpreter_tool_with_description():
"""Test HostedCodeInterpreterTool with description and additional properties."""
tool = HostedCodeInterpreterTool(
name="test_interpreter",
description="A test code interpreter",
additional_properties={"version": "1.0", "language": "python"},
)
assert tool.name == "test_interpreter"
assert tool.name == "code_interpreter"
assert tool.description == "A test code interpreter"
assert tool.additional_properties == {"version": "1.0", "language": "python"}
python/packages/main/tests/main/test_types.py
+4 -4
View File
@@ -12,6 +12,7 @@ from agent_framework import (
AIAnnotation,
AIContent,
AIContents,
AIFunction,
AITool,
AnnotatedRegion,
ChatFinishReason,
@@ -723,11 +724,12 @@ def test_chat_options_init_with_args(ai_function_tool, ai_tool) -> None:
assert options.presence_penalty == 0.0
assert options.frequency_penalty == 0.0
assert options.user == "user-123"
for tool in options._ai_tools:
for tool in options.tools:
assert isinstance(tool, AITool)
assert tool.name is not None
assert tool.description is not None
assert tool.parameters() is not None
if isinstance(tool, AIFunction):
assert tool.parameters() is not None
settings = options.to_provider_settings()
assert settings["model"] == "gpt-4" # uses alias
@@ -754,8 +756,6 @@ def test_chat_options_and(ai_function_tool, ai_tool) -> None:
options3 = options1 & options2
assert options3.ai_model_id == "gpt-4.1"
assert len(options3._ai_tools) == 2
assert options3._ai_tools == [ai_function_tool, ai_tool]
assert options3.tools == [ai_function_tool, ai_tool]
assert options3.logit_bias == {"x": 1}
assert options3.metadata == {"a": "b"}
python/packages/main/tests/openai/test_openai_chat_client_base.py
+6 -7
View File
@@ -15,7 +15,6 @@ from pydantic import BaseModel
from agent_framework import ChatMessage, ChatResponseUpdate
from agent_framework.exceptions import (
ServiceInvalidResponseError,
ServiceResponseException,
)
from agent_framework.openai import OpenAIChatClient
@@ -192,7 +191,7 @@ async def test_cmc_general_exception(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_scmc(
async def test_get_streaming(
mock_create: AsyncMock,
chat_history: list[ChatMessage],
openai_unit_test_env: dict[str, str],
@@ -231,7 +230,7 @@ async def test_scmc(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_scmc_singular(
async def test_get_streaming_singular(
mock_create: AsyncMock,
chat_history: list[ChatMessage],
openai_unit_test_env: dict[str, str],
@@ -270,7 +269,7 @@ async def test_scmc_singular(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_scmc_structured_output_no_fcc(
async def test_get_streaming_structured_output_no_fcc(
mock_create: AsyncMock,
chat_history: list[ChatMessage],
openai_unit_test_env: dict[str, str],
@@ -308,7 +307,7 @@ async def test_scmc_structured_output_no_fcc(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_scmc_no_fcc_in_response(
async def test_get_streaming_no_fcc_in_response(
mock_create: AsyncMock,
chat_history: list[ChatMessage],
mock_streaming_chat_completion_response: ChatCompletion,
@@ -334,7 +333,7 @@ async def test_scmc_no_fcc_in_response(
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_scmc_no_stream(
async def test_get_streaming_no_stream(
mock_create: AsyncMock,
chat_history: list[ChatMessage],
openai_unit_test_env: dict[str, str],
@@ -344,7 +343,7 @@ async def test_scmc_no_stream(
chat_history.append(ChatMessage(role="user", text="hello world"))
openai_chat_completion = OpenAIChatClient()
with pytest.raises(ServiceInvalidResponseError):
with pytest.raises(ServiceResponseException):
[
msg
async for msg in openai_chat_completion.get_streaming_response(
python/packages/main/tests/openai/test_openai_responses_client.py
+31 -35
View File
@@ -7,7 +7,7 @@ import pytest
from pydantic import BaseModel
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai import OpenAIResponsesClient
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
@@ -247,23 +247,21 @@ async def test_openai_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/openai/openai-python/issues/2305
with pytest.raises(ServiceResponseException):
response = openai_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 = openai_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
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather
@skip_if_openai_integration_tests_disabled
@@ -294,22 +292,20 @@ async def test_openai_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/openai/openai-python/issues/2305
with pytest.raises(ServiceResponseException):
response = openai_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 = openai_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
output = OutputStruct.model_validate_json(full_message)
assert "Seattle" in output.location
assert "sunny" in output.weather