mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Azure chat client (#185)
* updated openai, fcc works, with sample * reduced files in openai * Add azure chat client * fix tests * Update python/packages/main/tests/unit/test_openai_chat_completion_base.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/azure/agent_framework/azure/__init__.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/azure/agent_framework/azure/_azure_openai_settings.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * PR comments * fix bad merge * disable tests for now * actually disable tests for azure * fix tests, align test files with merge changes * update code for new project structure * PR comments * add streaming integration tests. Fix flakiness --------- Co-authored-by: eavanvalkenburg <github@vanvalkenburg.eu> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
fa96f74ee9
commit
f0dc661c3e
@@ -2,9 +2,16 @@
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._chat_client import AzureChatClient
|
||||
from ._entra_id_authentication import get_entra_auth_token
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
__all__ = ["__version__"]
|
||||
__all__ = [
|
||||
"AzureChatClient",
|
||||
"__version__",
|
||||
"get_entra_auth_token",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import logging
|
||||
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,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIModelTypes
|
||||
from agent_framework.openai._chat_client import OpenAIChatClientBase
|
||||
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_chunk import Choice as ChunkChoice
|
||||
from pydantic import SecretStr, ValidationError
|
||||
from pydantic.networks import AnyUrl
|
||||
|
||||
from ._shared import (
|
||||
DEFAULT_AZURE_API_VERSION,
|
||||
DEFAULT_AZURE_TOKEN_ENDPOINT,
|
||||
AzureOpenAIConfigBase,
|
||||
AzureOpenAISettings,
|
||||
)
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
TChatResponse = TypeVar("TChatResponse", ChatResponse, ChatResponseUpdate)
|
||||
|
||||
|
||||
class AzureChatClient(AzureOpenAIConfigBase, OpenAIChatClientBase):
|
||||
"""Azure Chat completion class."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
deployment_name: str | None = None,
|
||||
endpoint: str | None = None,
|
||||
base_url: str | None = None,
|
||||
api_version: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncAzureOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
instruction_role: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize an AzureChatCompletion service.
|
||||
|
||||
Args:
|
||||
api_key (str | None): The optional api key. If provided, will override the value in the
|
||||
env vars or .env file.
|
||||
deployment_name (str | None): The optional deployment. If provided, will override the value
|
||||
(chat_deployment_name) in the env vars or .env file.
|
||||
endpoint (str | None): The optional deployment endpoint. If provided will override the value
|
||||
in the env vars or .env file.
|
||||
base_url (str | None): The optional deployment base_url. If provided will override the value
|
||||
in the env vars or .env file.
|
||||
api_version (str | None): The optional deployment api version. If provided will override the value
|
||||
in the env vars or .env file.
|
||||
ad_token (str | None): The Azure Active Directory token. (Optional)
|
||||
ad_token_provider (AsyncAzureADTokenProvider): The Azure Active Directory token provider. (Optional)
|
||||
token_endpoint (str | None): The token endpoint to request an Azure token. (Optional)
|
||||
default_headers (Mapping[str, str]): The default headers mapping of string keys to
|
||||
string values for HTTP requests. (Optional)
|
||||
async_client (AsyncAzureOpenAI | None): An existing client to use. (Optional)
|
||||
env_file_path (str | None): Use the environment settings file as a fallback to using env vars.
|
||||
env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'.
|
||||
instruction_role (str | None): The role to use for 'instruction' messages, for example, summarization
|
||||
prompts could use `developer` or `system`. (Optional)
|
||||
"""
|
||||
try:
|
||||
# Filter out any None values from the arguments
|
||||
azure_openai_settings = AzureOpenAISettings(
|
||||
api_key=SecretStr(api_key) if api_key else None,
|
||||
base_url=AnyUrl(base_url) if base_url else None,
|
||||
endpoint=AnyUrl(endpoint) if endpoint else None,
|
||||
chat_deployment_name=deployment_name,
|
||||
api_version=api_version or DEFAULT_AZURE_API_VERSION,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
token_endpoint=token_endpoint or DEFAULT_AZURE_TOKEN_ENDPOINT,
|
||||
)
|
||||
except ValidationError as exc:
|
||||
raise ServiceInitializationError(f"Failed to validate settings: {exc}") from exc
|
||||
|
||||
if not azure_openai_settings.chat_deployment_name:
|
||||
raise ServiceInitializationError("chat_deployment_name is required.")
|
||||
|
||||
super().__init__(
|
||||
deployment_name=azure_openai_settings.chat_deployment_name,
|
||||
endpoint=azure_openai_settings.endpoint,
|
||||
base_url=azure_openai_settings.base_url,
|
||||
api_version=azure_openai_settings.api_version,
|
||||
api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None,
|
||||
ad_token=ad_token,
|
||||
ad_token_provider=ad_token_provider,
|
||||
token_endpoint=azure_openai_settings.token_endpoint,
|
||||
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":
|
||||
"""Initialize an Azure OpenAI service from a dictionary of settings.
|
||||
|
||||
Args:
|
||||
settings: A dictionary of settings for the service.
|
||||
should contain keys: service_id, and optionally:
|
||||
ad_auth, ad_token_provider, default_headers
|
||||
"""
|
||||
return AzureChatClient(
|
||||
api_key=settings.get("api_key"),
|
||||
deployment_name=settings.get("deployment_name"),
|
||||
endpoint=settings.get("endpoint"),
|
||||
base_url=settings.get("base_url"),
|
||||
api_version=settings.get("api_version"),
|
||||
ad_token=settings.get("ad_token"),
|
||||
ad_token_provider=settings.get("ad_token_provider"),
|
||||
default_headers=settings.get("default_headers"),
|
||||
env_file_path=settings.get("env_file_path"),
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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.
|
||||
"""
|
||||
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)]
|
||||
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
|
||||
from agent_framework.exceptions import ServiceInvalidAuthError
|
||||
from azure.core.exceptions import ClientAuthenticationError
|
||||
from azure.identity import DefaultAzureCredential
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_entra_auth_token(token_endpoint: str) -> 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:
|
||||
token_endpoint: The token endpoint to use to retrieve the authentication token.
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
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}`.")
|
||||
return None
|
||||
|
||||
return auth_token.token if auth_token else None
|
||||
@@ -0,0 +1,261 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from copy import copy
|
||||
from typing import Any, ClassVar, Final
|
||||
|
||||
from agent_framework import AFBaseSettings, HttpsUrl
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIHandler, OpenAIModelTypes
|
||||
from agent_framework.telemetry import USER_AGENT_KEY
|
||||
from openai.lib.azure import AsyncAzureOpenAI
|
||||
from pydantic import ConfigDict, SecretStr, validate_call
|
||||
|
||||
from ._entra_id_authentication import get_entra_auth_token
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
DEFAULT_AZURE_API_VERSION: Final[str] = "2024-10-21"
|
||||
DEFAULT_AZURE_TOKEN_ENDPOINT: Final[str] = "https://cognitiveservices.azure.com/.default" # noqa: S105
|
||||
|
||||
|
||||
class AzureOpenAISettings(AFBaseSettings):
|
||||
"""AzureOpenAI model settings.
|
||||
|
||||
The settings are first loaded from environment variables with the prefix 'AZURE_OPENAI_'.
|
||||
If the environment variables are not found, the settings can be loaded from a .env file
|
||||
with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
|
||||
are ignored; however, validation will fail alerting that the settings are missing.
|
||||
|
||||
Optional settings for prefix 'AZURE_OPENAI_' are:
|
||||
- chat_deployment_name: str - The name of the Azure Chat deployment. This value
|
||||
will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_CHAT_DEPLOYMENT_NAME)
|
||||
- responses_deployment_name: str - The name of the Azure Responses deployment. This value
|
||||
will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME)
|
||||
- text_deployment_name: str - The name of the Azure Text deployment. This value
|
||||
will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_TEXT_DEPLOYMENT_NAME)
|
||||
- embedding_deployment_name: str - The name of the Azure Embedding deployment. This value
|
||||
will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME)
|
||||
- text_to_image_deployment_name: str - The name of the Azure Text to Image deployment. This
|
||||
value will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME)
|
||||
- audio_to_text_deployment_name: str - The name of the Azure Audio to Text deployment. This
|
||||
value will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME)
|
||||
- text_to_audio_deployment_name: str - The name of the Azure Text to Audio deployment. This
|
||||
value will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME)
|
||||
- realtime_deployment_name: str - The name of the Azure Realtime deployment. This value
|
||||
will correspond to the custom name you chose for your deployment
|
||||
when you deployed a model. This value can be found under
|
||||
Resource Management > Deployments in the Azure portal or, alternatively,
|
||||
under Management > Deployments in Azure AI Foundry.
|
||||
(Env var AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME)
|
||||
- api_key: SecretStr - The API key for the Azure deployment. This value can be
|
||||
found in the Keys & Endpoint section when examining your resource in
|
||||
the Azure portal. You can use either KEY1 or KEY2.
|
||||
(Env var AZURE_OPENAI_API_KEY)
|
||||
- base_url: HttpsUrl | None - base_url: The url of the Azure deployment. This value
|
||||
can be found in the Keys & Endpoint section when examining
|
||||
your resource from the Azure portal, the base_url consists of the endpoint,
|
||||
followed by /openai/deployments/{deployment_name}/,
|
||||
use endpoint if you only want to supply the endpoint.
|
||||
(Env var AZURE_OPENAI_BASE_URL)
|
||||
- endpoint: HttpsUrl - The endpoint of the Azure deployment. This value
|
||||
can be found in the Keys & Endpoint section when examining
|
||||
your resource from the Azure portal, the endpoint should end in openai.azure.com.
|
||||
If both base_url and endpoint are supplied, base_url will be used.
|
||||
(Env var AZURE_OPENAI_ENDPOINT)
|
||||
- api_version: str | None - The API version to use. The default value is "2024-02-01".
|
||||
(Env var AZURE_OPENAI_API_VERSION)
|
||||
- token_endpoint: str - The token endpoint to use to retrieve the authentication token.
|
||||
The default value is "https://cognitiveservices.azure.com/.default".
|
||||
(Env var AZURE_OPENAI_TOKEN_ENDPOINT)
|
||||
"""
|
||||
|
||||
env_prefix: ClassVar[str] = "AZURE_OPENAI_"
|
||||
|
||||
chat_deployment_name: str | None = None
|
||||
responses_deployment_name: str | None = None
|
||||
text_deployment_name: str | None = None
|
||||
embedding_deployment_name: str | None = None
|
||||
text_to_image_deployment_name: str | None = None
|
||||
audio_to_text_deployment_name: str | None = None
|
||||
text_to_audio_deployment_name: str | None = None
|
||||
realtime_deployment_name: str | None = None
|
||||
endpoint: HttpsUrl | None = None
|
||||
base_url: HttpsUrl | None = None
|
||||
api_key: SecretStr | None = None
|
||||
api_version: str = DEFAULT_AZURE_API_VERSION
|
||||
token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT
|
||||
|
||||
def get_azure_openai_auth_token(self, token_endpoint: str | None = None) -> str | None:
|
||||
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint for the use with Azure OpenAI.
|
||||
|
||||
The required role for the token is `Cognitive Services OpenAI Contributor`.
|
||||
The token endpoint may be specified as an environment variable, via the .env
|
||||
file or as an argument. If the token endpoint is not provided, the default is None.
|
||||
The `token_endpoint` argument takes precedence over the `token_endpoint` attribute.
|
||||
|
||||
Args:
|
||||
token_endpoint: The token endpoint to use. Defaults to `https://cognitiveservices.azure.com/.default`.
|
||||
|
||||
Returns:
|
||||
The Azure token or None if the token could not be retrieved.
|
||||
|
||||
Raises:
|
||||
ServiceInitializationError: If the token endpoint is not provided.
|
||||
"""
|
||||
endpoint_to_use = token_endpoint or self.token_endpoint
|
||||
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)
|
||||
|
||||
|
||||
class AzureOpenAIConfigBase(OpenAIHandler):
|
||||
"""Internal class for configuring a connection to an Azure OpenAI service."""
|
||||
|
||||
@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,
|
||||
api_key: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: Callable[[], str | Awaitable[str]] | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
client: AsyncAzureOpenAI | None = None,
|
||||
instruction_role: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Internal class for configuring a connection to an Azure OpenAI service.
|
||||
|
||||
The `validate_call` decorator is used with a configuration that allows arbitrary types.
|
||||
This is necessary for types like `HttpsUrl` and `OpenAIModelTypes`.
|
||||
|
||||
Args:
|
||||
deployment_name (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)
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
"""
|
||||
# Merge APP_INFO into the headers if it exists
|
||||
merged_headers = dict(copy(default_headers)) if default_headers else {}
|
||||
|
||||
if not client:
|
||||
# If the client is None, the api_key is none, the ad_token is none, and the ad_token_provider is none,
|
||||
# then we will attempt to get the ad_token using the default endpoint specified in the Azure OpenAI
|
||||
# settings.
|
||||
if not api_key and not ad_token_provider and not ad_token and token_endpoint:
|
||||
ad_token = get_entra_auth_token(token_endpoint)
|
||||
|
||||
if not api_key and not ad_token and not ad_token_provider:
|
||||
raise ServiceInitializationError(
|
||||
"Please provide either api_key, ad_token or ad_token_provider or a client."
|
||||
)
|
||||
|
||||
if not endpoint and not base_url:
|
||||
raise ServiceInitializationError("Please provide an endpoint or a base_url")
|
||||
|
||||
args: dict[str, Any] = {
|
||||
"default_headers": merged_headers,
|
||||
}
|
||||
if api_version:
|
||||
args["api_version"] = api_version
|
||||
if ad_token:
|
||||
args["azure_ad_token"] = ad_token
|
||||
if ad_token_provider:
|
||||
args["azure_ad_token_provider"] = ad_token_provider
|
||||
if api_key:
|
||||
args["api_key"] = api_key
|
||||
if base_url:
|
||||
args["base_url"] = str(base_url)
|
||||
if endpoint and not base_url:
|
||||
args["azure_endpoint"] = str(endpoint)
|
||||
# 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:
|
||||
args["azure_deployment"] = deployment_name
|
||||
|
||||
if "websocket_base_url" in kwargs:
|
||||
args["websocket_base_url"] = kwargs.pop("websocket_base_url")
|
||||
|
||||
client = AsyncAzureOpenAI(**args)
|
||||
args = {
|
||||
"ai_model_id": deployment_name,
|
||||
"client": client,
|
||||
"ai_model_type": ai_model_type,
|
||||
}
|
||||
if instruction_role:
|
||||
args["instruction_role"] = instruction_role
|
||||
super().__init__(**args, **kwargs)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert the configuration to a dictionary."""
|
||||
client_settings = {
|
||||
"base_url": str(self.client.base_url),
|
||||
"api_version": self.client._custom_query["api-version"], # type: ignore
|
||||
"api_key": self.client.api_key,
|
||||
"ad_token": getattr(self.client, "_azure_ad_token", None),
|
||||
"ad_token_provider": getattr(self.client, "_azure_ad_token_provider", None),
|
||||
"default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY},
|
||||
}
|
||||
base = self.model_dump(
|
||||
exclude={
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"total_tokens",
|
||||
"api_type",
|
||||
"org_id",
|
||||
"ai_model_type",
|
||||
"service_id",
|
||||
"client",
|
||||
},
|
||||
by_alias=True,
|
||||
exclude_none=True,
|
||||
)
|
||||
base.update(client_settings)
|
||||
return base
|
||||
@@ -24,6 +24,8 @@ classifiers = [
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework",
|
||||
"azure-identity >= 1.13",
|
||||
"openai>=1.94.0",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
# region: Connector Settings fixtures
|
||||
@fixture
|
||||
def exclude_list(request: Any) -> list[str]:
|
||||
"""Fixture that returns a list of environment variables to exclude."""
|
||||
return request.param if hasattr(request, "param") else []
|
||||
|
||||
|
||||
@fixture
|
||||
def override_env_param_dict(request: Any) -> dict[str, str]:
|
||||
"""Fixture that returns a dict of environment variables to override."""
|
||||
return request.param if hasattr(request, "param") else {}
|
||||
|
||||
|
||||
# These two fixtures are used for multiple things, also non-connector tests
|
||||
@fixture()
|
||||
def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
|
||||
"""Fixture to set environment variables for AzureOpenAISettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
|
||||
"AZURE_OPENAI_TEXT_DEPLOYMENT_NAME": "test_text_deployment",
|
||||
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
|
||||
"AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME": "test_text_to_image_deployment",
|
||||
"AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME": "test_audio_to_text_deployment",
|
||||
"AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME": "test_text_to_audio_deployment",
|
||||
"AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME": "test_realtime_deployment",
|
||||
"AZURE_OPENAI_API_KEY": "test_api_key",
|
||||
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.com",
|
||||
"AZURE_OPENAI_API_VERSION": "2023-03-15-preview",
|
||||
"AZURE_OPENAI_BASE_URL": "https://test_text_deployment.test-base-url.com",
|
||||
"AZURE_OPENAI_TOKEN_ENDPOINT": "https://test-token-endpoint.com",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict) # type: ignore
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value) # type: ignore
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False) # type: ignore
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def chat_history() -> list[ChatMessage]:
|
||||
return []
|
||||
@@ -0,0 +1,121 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
|
||||
|
||||
from agent_framework_azure import AzureChatClient
|
||||
|
||||
|
||||
@ai_function
|
||||
def get_story_text() -> str:
|
||||
"""Returns a story about Emily and David."""
|
||||
return (
|
||||
"Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change."
|
||||
)
|
||||
|
||||
|
||||
async def test_azure_openai_chat_client_response() -> None:
|
||||
"""Test Azure OpenAI chat completion responses."""
|
||||
azure_chat_client = AzureChatClient(deployment_name="gpt-4o")
|
||||
|
||||
assert isinstance(azure_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await azure_chat_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "scientists" in response.text
|
||||
|
||||
|
||||
async def test_azure_openai_chat_client_response_tools() -> None:
|
||||
"""Test AzureOpenAI chat completion responses."""
|
||||
azure_chat_client = AzureChatClient(deployment_name="gpt-4o")
|
||||
|
||||
assert isinstance(azure_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await azure_chat_client.get_response(
|
||||
messages=messages,
|
||||
tools=[get_story_text],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "scientists" in response.text
|
||||
|
||||
|
||||
async def test_azure_openai_chat_client_streaming() -> None:
|
||||
"""Test Azure OpenAI chat completion responses."""
|
||||
azure_chat_client = AzureChatClient(deployment_name="gpt-4o")
|
||||
|
||||
assert isinstance(azure_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = azure_chat_client.get_streaming_response(messages=messages)
|
||||
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert "scientists" in full_message
|
||||
|
||||
|
||||
async def test_azure_openai_chat_client_streaming_tools() -> None:
|
||||
"""Test AzureOpenAI chat completion responses."""
|
||||
azure_chat_client = AzureChatClient(deployment_name="gpt-4o")
|
||||
|
||||
assert isinstance(azure_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = azure_chat_client.get_streaming_response(
|
||||
messages=messages,
|
||||
tools=[get_story_text],
|
||||
tool_choice="auto",
|
||||
)
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert "scientists" in full_message
|
||||
@@ -0,0 +1,620 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
ChatClientBase,
|
||||
ChatMessage,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
|
||||
from agent_framework.openai.exceptions import (
|
||||
ContentFilterResultSeverity,
|
||||
OpenAIContentFilterException,
|
||||
)
|
||||
from agent_framework.telemetry import USER_AGENT_KEY
|
||||
from httpx import Request, Response
|
||||
from openai import AsyncAzureOpenAI, AsyncStream
|
||||
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionChunk
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
|
||||
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
|
||||
from openai.types.chat.chat_completion_message import ChatCompletionMessage
|
||||
|
||||
from agent_framework_azure import AzureChatClient
|
||||
|
||||
# region Service Setup
|
||||
|
||||
|
||||
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
assert azure_chat_client.client is not None
|
||||
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_client, ChatClientBase)
|
||||
|
||||
|
||||
def test_init_client(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization with client
|
||||
client = MagicMock(spec=AsyncAzureOpenAI)
|
||||
azure_chat_client = AzureChatClient(async_client=client)
|
||||
|
||||
assert azure_chat_client.client is not None
|
||||
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
|
||||
|
||||
|
||||
def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
# Custom header for testing
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
azure_chat_client = AzureChatClient(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert azure_chat_client.client is not None
|
||||
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_client, ChatClientBase)
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_chat_client.client.default_headers
|
||||
assert azure_chat_client.client.default_headers[key] == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
assert azure_chat_client.client is not None
|
||||
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_client, ChatClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_init_with_empty_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatClient(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatClient(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_init_with_invalid_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatClient()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
|
||||
default_headers = {"X-Test": "test"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_chat_client = AzureChatClient.from_dict(settings)
|
||||
dumped_settings = azure_chat_client.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["deployment_name"]
|
||||
assert str(settings["endpoint"]) in str(dumped_settings["base_url"])
|
||||
assert str(settings["deployment_name"]) in str(dumped_settings["base_url"])
|
||||
assert settings["api_key"] == dumped_settings["api_key"]
|
||||
assert settings["api_version"] == dumped_settings["api_version"]
|
||||
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
|
||||
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
|
||||
assert USER_AGENT_KEY not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
# endregion
|
||||
# region CMC
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_response() -> ChatCompletion:
|
||||
return ChatCompletion(
|
||||
id="test_id",
|
||||
choices=[
|
||||
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
|
||||
],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
|
||||
content = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content]
|
||||
return stream
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
await azure_chat_client.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=False,
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_with_logit_bias(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
token_bias: dict[str | int, float] = {"1": -100}
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
await azure_chat_client.get_response(messages=chat_history, logit_bias=token_bias)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
stream=False,
|
||||
logit_bias=token_bias,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_with_stop(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
stop = ["!"]
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
await azure_chat_client.get_response(messages=chat_history, stop=stop)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
stream=False,
|
||||
stop=stop,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context={ # type: ignore
|
||||
"citations": {
|
||||
"content": "test content",
|
||||
"title": "test title",
|
||||
"url": "test url",
|
||||
"filepath": "test filepath",
|
||||
"chunk_id": "test chunk_id",
|
||||
},
|
||||
"intent": "query used",
|
||||
},
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
messages_out: list[ChatMessage] = []
|
||||
messages_out.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
content = await azure_chat_client.get_response(
|
||||
messages=messages_in,
|
||||
additional_properties={"extra_body": expected_data_settings},
|
||||
)
|
||||
assert len(content.messages) == 1
|
||||
assert len(content.messages[0].contents) == 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"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data_string(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context=json.dumps({ # type: ignore
|
||||
"citations": {
|
||||
"content": "test content",
|
||||
"title": "test title",
|
||||
"url": "test url",
|
||||
"filepath": "test filepath",
|
||||
"chunk_id": "test chunk_id",
|
||||
},
|
||||
"intent": "query used",
|
||||
}),
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.append(ChatMessage(text=prompt, role="user"))
|
||||
messages_out: list[ChatMessage] = []
|
||||
messages_out.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
content = await azure_chat_client.get_response(
|
||||
messages=messages_in,
|
||||
additional_properties={"extra_body": expected_data_settings},
|
||||
)
|
||||
assert len(content.messages) == 1
|
||||
assert len(content.messages[0].contents) == 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"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data_fail(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context="not a dictionary", # type: ignore
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.append(ChatMessage(text=prompt, role="user"))
|
||||
messages_out: list[ChatMessage] = []
|
||||
messages_out.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
content = await azure_chat_client.get_response(
|
||||
messages=messages_in,
|
||||
additional_properties={"extra_body": expected_data_settings},
|
||||
)
|
||||
assert len(content.messages) == 1
|
||||
assert len(content.messages[0].contents) == 1
|
||||
assert isinstance(content.messages[0].contents[0], TextContent)
|
||||
assert content.messages[0].contents[0].text == "test"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(messages_out), # type: ignore
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@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 "
|
||||
"documentation: https://go.microsoft.com/fwlink/?linkid=2198766"
|
||||
)
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
|
||||
"Error code: 400 - {'error': {'message': \"%s\", 'type': null, 'param': 'prompt', 'code': 'content_filter', "
|
||||
"'status': 400, 'innererror': {'code': 'ResponsibleAIPolicyViolation', 'content_filter_result': {'hate': "
|
||||
"{'filtered': True, 'severity': 'high'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': "
|
||||
"{'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}}}"
|
||||
) % CONTENT_FILTERED_ERROR_MESSAGE
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_content_filtering_raises_correct_exception(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
assert test_endpoint is not None
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
|
||||
response=Response(400, request=Request("POST", test_endpoint)),
|
||||
body={
|
||||
"message": CONTENT_FILTERED_ERROR_MESSAGE,
|
||||
"type": None,
|
||||
"param": "prompt",
|
||||
"code": "content_filter",
|
||||
"status": 400,
|
||||
"innererror": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_result": {
|
||||
"hate": {"filtered": True, "severity": "high"},
|
||||
"self_harm": {"filtered": False, "severity": "safe"},
|
||||
"sexual": {"filtered": False, "severity": "safe"},
|
||||
"violence": {"filtered": False, "severity": "safe"},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error") as exc_info:
|
||||
await azure_chat_client.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
|
||||
content_filter_exc = exc_info.value
|
||||
assert content_filter_exc.param == "prompt"
|
||||
assert content_filter_exc.content_filter_result["hate"].filtered
|
||||
assert content_filter_exc.content_filter_result["hate"].severity == ContentFilterResultSeverity.HIGH
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_content_filtering_without_response_code_raises_with_default_code(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
assert test_endpoint is not None
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
|
||||
response=Response(400, request=Request("POST", test_endpoint)),
|
||||
body={
|
||||
"message": CONTENT_FILTERED_ERROR_MESSAGE,
|
||||
"type": None,
|
||||
"param": "prompt",
|
||||
"code": "content_filter",
|
||||
"status": 400,
|
||||
"innererror": {
|
||||
"content_filter_result": {
|
||||
"hate": {"filtered": True, "severity": "high"},
|
||||
"self_harm": {"filtered": False, "severity": "safe"},
|
||||
"sexual": {"filtered": False, "severity": "safe"},
|
||||
"violence": {"filtered": False, "severity": "safe"},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error"):
|
||||
await azure_chat_client.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_bad_request_non_content_filter(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.append(ChatMessage(text=prompt, role="user"))
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
assert test_endpoint is not None
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
|
||||
)
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
|
||||
with pytest.raises(ServiceResponseException, match="service failed to complete the prompt"):
|
||||
await azure_chat_client.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming(
|
||||
mock_create: AsyncMock,
|
||||
azure_openai_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
|
||||
) -> None:
|
||||
mock_create.return_value = mock_streaming_chat_completion_response
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
azure_chat_client = AzureChatClient()
|
||||
async for msg in azure_chat_client.get_streaming_response(
|
||||
messages=chat_history,
|
||||
):
|
||||
assert msg is not None
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=True,
|
||||
messages=azure_chat_client._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
# NOTE: The `stream_options={"include_usage": True}` is explicitly enforced in
|
||||
# `OpenAIChatCompletionBase._inner_get_streaming_response`.
|
||||
# To ensure consistency, we align the arguments here accordingly.
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
@@ -10,7 +10,6 @@ except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
_IMPORTS = {
|
||||
"get_logger": "._logging",
|
||||
"AFBaseModel": "._pydantic",
|
||||
"AFBaseSettings": "._pydantic",
|
||||
"Agent": "._agents",
|
||||
@@ -18,39 +17,41 @@ _IMPORTS = {
|
||||
"AgentRunResponseUpdate": "._types",
|
||||
"AgentThread": "._agents",
|
||||
"AITool": "._tools",
|
||||
"ai_function": "._tools",
|
||||
"AIFunction": "._tools",
|
||||
"AIContent": "._types",
|
||||
"AIContents": "._types",
|
||||
"ChatClient": "._clients",
|
||||
"ChatClientAgent": "._agents",
|
||||
"ChatClientAgentThread": "._agents",
|
||||
"ChatClientAgentThreadType": "._agents",
|
||||
"ChatClientBase": "._clients",
|
||||
"ChatFinishReason": "._types",
|
||||
"ChatMessage": "._types",
|
||||
"ChatOptions": "._types",
|
||||
"ChatResponse": "._types",
|
||||
"ChatResponseUpdate": "._types",
|
||||
"ChatRole": "._types",
|
||||
"ChatToolMode": "._types",
|
||||
"DataContent": "._types",
|
||||
"EmbeddingGenerator": "._clients",
|
||||
"ErrorContent": "._types",
|
||||
"FunctionCallContent": "._types",
|
||||
"FunctionResultContent": "._types",
|
||||
"GeneratedEmbeddings": "._types",
|
||||
"HttpsUrl": "._pydantic",
|
||||
"InputGuardrail": ".guard_rails",
|
||||
"OutputGuardrail": ".guard_rails",
|
||||
"SpeechToTextOptions": "._types",
|
||||
"StructuredResponse": "._types",
|
||||
"TextContent": "._types",
|
||||
"TextReasoningContent": "._types",
|
||||
"DataContent": "._types",
|
||||
"TextToSpeechOptions": "._types",
|
||||
"UriContent": "._types",
|
||||
"UsageContent": "._types",
|
||||
"UsageDetails": "._types",
|
||||
"FunctionCallContent": "._types",
|
||||
"FunctionResultContent": "._types",
|
||||
"ChatFinishReason": "._types",
|
||||
"ChatMessage": "._types",
|
||||
"ChatResponse": "._types",
|
||||
"StructuredResponse": "._types",
|
||||
"ChatResponseUpdate": "._types",
|
||||
"ChatRole": "._types",
|
||||
"ErrorContent": "._types",
|
||||
"GeneratedEmbeddings": "._types",
|
||||
"ChatOptions": "._types",
|
||||
"ChatToolMode": "._types",
|
||||
"ChatClient": "._clients",
|
||||
"ChatClientBase": "._clients",
|
||||
"ai_function": "._tools",
|
||||
"get_logger": "._logging",
|
||||
"use_tool_calling": "._clients",
|
||||
"EmbeddingGenerator": "._clients",
|
||||
"InputGuardrail": ".guard_rails",
|
||||
"OutputGuardrail": ".guard_rails",
|
||||
"TextToSpeechOptions": "._types",
|
||||
"SpeechToTextOptions": "._types",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ from . import __version__ # type: ignore[attr-defined]
|
||||
from ._agents import Agent, AgentThread, ChatClientAgent, ChatClientAgentThread, ChatClientAgentThreadType
|
||||
from ._clients import ChatClient, ChatClientBase, EmbeddingGenerator, use_tool_calling
|
||||
from ._logging import get_logger
|
||||
from ._pydantic import AFBaseModel, AFBaseSettings
|
||||
from ._pydantic import AFBaseModel, AFBaseSettings, HttpsUrl
|
||||
from ._tools import AIFunction, AITool, ai_function
|
||||
from ._types import (
|
||||
AgentRunResponse,
|
||||
@@ -63,6 +63,7 @@ __all__ = [
|
||||
"FunctionCallContent",
|
||||
"FunctionResultContent",
|
||||
"GeneratedEmbeddings",
|
||||
"HttpsUrl",
|
||||
"InputGuardrail",
|
||||
"OutputGuardrail",
|
||||
"SpeechToTextOptions",
|
||||
|
||||
@@ -112,7 +112,14 @@ def _tool_call_non_streaming(func: TInnerGetResponse) -> TInnerGetResponse:
|
||||
for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
|
||||
response = await func(self, messages=messages, chat_options=chat_options)
|
||||
# if there are function calls, we will handle them first
|
||||
function_calls = [it for it in response.messages[0].contents if isinstance(it, FunctionCallContent)]
|
||||
function_results = {
|
||||
it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
|
||||
}
|
||||
function_calls = [
|
||||
it
|
||||
for it in response.messages[0].contents
|
||||
if isinstance(it, FunctionCallContent) and it.call_id not in function_results
|
||||
]
|
||||
if function_calls:
|
||||
# Run all function calls concurrently
|
||||
results = await asyncio.gather(*[
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from typing import Any, ClassVar, TypeVar
|
||||
from typing import Annotated, Any, ClassVar, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, UrlConstraints
|
||||
from pydantic.networks import AnyUrl
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
HttpsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
|
||||
|
||||
|
||||
class AFBaseModel(BaseModel):
|
||||
"""Base class for all pydantic models in the Agent Framework."""
|
||||
|
||||
@@ -16,7 +16,16 @@ from collections.abc import (
|
||||
)
|
||||
from typing import Annotated, Any, ClassVar, Generic, Literal, TypeVar, overload
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError, field_validator, model_validator
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
PrivateAttr,
|
||||
ValidationError,
|
||||
field_validator,
|
||||
model_serializer,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
from ._pydantic import AFBaseModel
|
||||
from ._tools import AITool, ai_function
|
||||
@@ -1398,6 +1407,11 @@ class ChatToolMode(AFBaseModel):
|
||||
return self.mode == other.mode and self.required_function_name == other.required_function_name
|
||||
return False
|
||||
|
||||
@model_serializer
|
||||
def serialize_model(self) -> str:
|
||||
"""Serializes the ChatToolMode to just the mode string."""
|
||||
return self.mode
|
||||
|
||||
|
||||
ChatToolMode.AUTO = ChatToolMode(mode="auto") # type: ignore[assignment]
|
||||
ChatToolMode.REQUIRED_ANY = ChatToolMode(mode="required") # type: ignore[assignment]
|
||||
@@ -1496,10 +1510,13 @@ class ChatOptions(AFBaseModel):
|
||||
Dictionary of settings for provider.
|
||||
"""
|
||||
default_exclude = {"additional_properties"}
|
||||
# No tool choice if no tools are defined
|
||||
if self.tools is None or len(self.tools) == 0:
|
||||
default_exclude.add("tool_choice")
|
||||
merged_exclude = default_exclude if exclude is None else default_exclude | set(exclude)
|
||||
|
||||
settings = self.model_dump(exclude_none=True, by_alias=by_alias, exclude=merged_exclude)
|
||||
settings = {k: v for k, v in settings.items() if v and not isinstance(v, dict)}
|
||||
settings = {k: v for k, v in settings.items() if v}
|
||||
settings.update(self.additional_properties)
|
||||
for key in merged_exclude:
|
||||
settings.pop(key, None)
|
||||
|
||||
@@ -7,6 +7,8 @@ PACKAGE_EXTRA = "azure"
|
||||
|
||||
_IMPORTS = {
|
||||
"__version__": "agent_framework_azure",
|
||||
"AzureChatClient": "agent_framework_azure",
|
||||
"get_entra_auth_token": "agent_framework_azure",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -46,6 +46,12 @@ class ServiceContentFilterException(ServiceResponseException):
|
||||
pass
|
||||
|
||||
|
||||
class ServiceInvalidAuthError(ServiceException):
|
||||
"""An error occurred while authenticating the service."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class ServiceInvalidExecutionSettingsError(ServiceResponseException):
|
||||
"""An error occurred while validating the execution settings of the service."""
|
||||
|
||||
|
||||
@@ -2,9 +2,11 @@
|
||||
|
||||
|
||||
from ._chat_client import OpenAIChatClient
|
||||
from ._shared import OpenAISettings
|
||||
from ._shared import OpenAIHandler, OpenAIModelTypes, OpenAISettings
|
||||
|
||||
__all__ = [
|
||||
"OpenAIChatClient",
|
||||
"OpenAIHandler",
|
||||
"OpenAIModelTypes",
|
||||
"OpenAISettings",
|
||||
]
|
||||
|
||||
@@ -239,11 +239,11 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
|
||||
function_call = self._openai_content_parser(content)
|
||||
if "tool_calls" not in args:
|
||||
args["tool_calls"] = []
|
||||
args["tool_calls"].append(function_call)
|
||||
args["tool_calls"].append(function_call) # type: ignore
|
||||
case _:
|
||||
if "content" not in args:
|
||||
args["content"] = []
|
||||
args["content"].append(self._openai_content_parser(content))
|
||||
args["content"].append(self._openai_content_parser(content)) # type: ignore
|
||||
if "content" in args or "tool_calls" in args:
|
||||
all_messages.append(args)
|
||||
return all_messages
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from typing import Any
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
|
||||
# region: Connector Settings fixtures
|
||||
@fixture
|
||||
def exclude_list(request: Any) -> list[str]:
|
||||
"""Fixture that returns a list of environment variables to exclude."""
|
||||
return request.param if hasattr(request, "param") else []
|
||||
|
||||
|
||||
@fixture
|
||||
def override_env_param_dict(request: Any) -> dict[str, str]:
|
||||
"""Fixture that returns a dict of environment variables to override."""
|
||||
return request.param if hasattr(request, "param") else {}
|
||||
|
||||
|
||||
@fixture()
|
||||
def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
|
||||
"""Fixture to set environment variables for OpenAISettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"OPENAI_API_KEY": "test_api_key",
|
||||
"OPENAI_ORG_ID": "test_org_id",
|
||||
"OPENAI_RESPONSES_MODEL_ID": "test_responses_model_id",
|
||||
"OPENAI_CHAT_MODEL_ID": "test_chat_model_id",
|
||||
"OPENAI_TEXT_MODEL_ID": "test_text_model_id",
|
||||
"OPENAI_EMBEDDING_MODEL_ID": "test_embedding_model_id",
|
||||
"OPENAI_TEXT_TO_IMAGE_MODEL_ID": "test_text_to_image_model_id",
|
||||
"OPENAI_AUDIO_TO_TEXT_MODEL_ID": "test_audio_to_text_model_id",
|
||||
"OPENAI_TEXT_TO_AUDIO_MODEL_ID": "test_text_to_audio_model_id",
|
||||
"OPENAI_REALTIME_MODEL_ID": "test_realtime_model_id",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict) # type: ignore
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value) # type: ignore
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False) # type: ignore
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def chat_history() -> list[ChatMessage]:
|
||||
return []
|
||||
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
@ai_function
|
||||
def get_story_text() -> str:
|
||||
"""Returns a story about Emily and David."""
|
||||
return (
|
||||
"Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change."
|
||||
)
|
||||
|
||||
|
||||
async def test_openai_chat_completion_response() -> None:
|
||||
"""Test OpenAI chat completion responses."""
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-mini")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_chat_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "scientists" in response.text
|
||||
|
||||
|
||||
async def test_openai_chat_completion_response_tools() -> None:
|
||||
"""Test OpenAI chat completion responses."""
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-mini")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_chat_client.get_response(
|
||||
messages=messages,
|
||||
tools=[get_story_text],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "scientists" in response.text
|
||||
|
||||
|
||||
async def test_openai_chat_client_streaming() -> None:
|
||||
"""Test Azure OpenAI chat completion responses."""
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-mini")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
|
||||
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
|
||||
"groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
|
||||
"of climate change.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_chat_client.get_streaming_response(messages=messages)
|
||||
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert "scientists" in full_message
|
||||
|
||||
|
||||
async def test_openai_chat_client_streaming_tools() -> None:
|
||||
"""Test AzureOpenAI chat completion responses."""
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-mini")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_chat_client.get_streaming_response(
|
||||
messages=messages,
|
||||
tools=[get_story_text],
|
||||
tool_choice="auto",
|
||||
)
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert "scientists" in full_message
|
||||
@@ -0,0 +1,103 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import ChatClient
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization
|
||||
open_ai_chat_completion = OpenAIChatClient()
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert isinstance(open_ai_chat_completion, ChatClient)
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
# Test successful initialization
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatClient(api_key="34523", ai_model_id={"test": "dict"}) # type: ignore
|
||||
|
||||
|
||||
def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization
|
||||
ai_model_id = "test_model_id"
|
||||
open_ai_chat_completion = OpenAIChatClient(ai_model_id=ai_model_id)
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == ai_model_id
|
||||
assert isinstance(open_ai_chat_completion, ChatClient)
|
||||
|
||||
|
||||
def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test successful initialization
|
||||
open_ai_chat_completion = OpenAIChatClient(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert isinstance(open_ai_chat_completion, ChatClient)
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in open_ai_chat_completion.client.default_headers
|
||||
assert open_ai_chat_completion.client.default_headers[key] == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatClient(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatClient(
|
||||
ai_model_id=ai_model_id,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_serialize(openai_unit_test_env: dict[str, str]) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
open_ai_chat_completion = OpenAIChatClient.from_dict(settings)
|
||||
dumped_settings = open_ai_chat_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
|
||||
assert "User-Agent" not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None:
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"org_id": openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
}
|
||||
|
||||
open_ai_chat_completion = OpenAIChatClient.from_dict(settings)
|
||||
dumped_settings = open_ai_chat_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
|
||||
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
|
||||
assert "User-Agent" not in dumped_settings["default_headers"]
|
||||
@@ -0,0 +1,348 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from copy import deepcopy
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncStream
|
||||
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionChunk
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
|
||||
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
|
||||
from openai.types.chat.chat_completion_message import ChatCompletionMessage
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework import ChatMessage, ChatResponseUpdate
|
||||
from agent_framework.exceptions import (
|
||||
ServiceInvalidResponseError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
|
||||
async def mock_async_process_chat_stream_response(_):
|
||||
mock_content = MagicMock(spec=ChatResponseUpdate)
|
||||
yield mock_content, None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_response() -> ChatCompletion:
|
||||
return ChatCompletion(
|
||||
id="test_id",
|
||||
choices=[
|
||||
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
|
||||
],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
|
||||
content = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content]
|
||||
return stream
|
||||
|
||||
|
||||
# region Chat Message Content
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
await openai_chat_completion.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=False,
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_chat_options(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
await openai_chat_completion.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=False,
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_no_fcc_in_response(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
await openai_chat_completion.get_response(
|
||||
messages=chat_history,
|
||||
arguments={},
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=False,
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_structured_output_no_fcc(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
# Define a mock response format
|
||||
class Test(BaseModel):
|
||||
name: str
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
await openai_chat_completion.get_response(
|
||||
messages=chat_history,
|
||||
response_format=Test,
|
||||
)
|
||||
mock_create.assert_awaited_once()
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc_chat_options(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_streaming_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
):
|
||||
assert isinstance(msg, ChatResponseUpdate)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock, side_effect=Exception)
|
||||
async def test_cmc_general_exception(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await openai_chat_completion.get_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
|
||||
|
||||
# region Streaming
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
content1 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
content2 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content1, content2]
|
||||
mock_create.return_value = stream
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
):
|
||||
assert isinstance(msg, ChatResponseUpdate)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc_singular(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
content1 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
content2 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content1, content2]
|
||||
mock_create.return_value = stream
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
):
|
||||
assert isinstance(msg, ChatResponseUpdate)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc_structured_output_no_fcc(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
content1 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
content2 = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content1, content2]
|
||||
mock_create.return_value = stream
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
# Define a mock response format
|
||||
class Test(BaseModel):
|
||||
name: str
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
response_format=Test,
|
||||
):
|
||||
assert isinstance(msg, ChatResponseUpdate)
|
||||
mock_create.assert_awaited_once()
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc_no_fcc_in_response(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: ChatCompletion,
|
||||
openai_unit_test_env: dict[str, str],
|
||||
):
|
||||
mock_create.return_value = mock_streaming_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
[
|
||||
msg
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
]
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history), # type: ignore
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_scmc_no_stream(
|
||||
mock_create: AsyncMock,
|
||||
chat_history: list[ChatMessage],
|
||||
openai_unit_test_env: dict[str, str],
|
||||
mock_chat_completion_response: ChatCompletion, # AsyncStream[ChatCompletionChunk]?
|
||||
):
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(role="user", text="hello world"))
|
||||
|
||||
openai_chat_completion = OpenAIChatClient()
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
[
|
||||
msg
|
||||
async for msg in openai_chat_completion.get_streaming_response(
|
||||
messages=chat_history,
|
||||
)
|
||||
]
|
||||
Reference in New Issue
Block a user