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

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

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

Major refactoring of the Python Agent Framework client architecture:

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

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

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

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

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

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

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

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

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

* fixes

* updated observabilitty

* reset azure init.pyi

* fix errors

* updated adr number

* fix foundry local

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

* fix tests and pyprojects

* fix test vars

* updated function tests

* update durable

* updated test setup for functions

* Fix Foundry auth in workflow samples

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

* Stabilize Python integration workflows

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

* Update hosting samples for Foundry

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

* Trigger full CI rerun

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

* Trigger CI rerun again

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

* trigger rerun

* trigger rerun

* fix for litellm

* undo durabletask changes

* Move Foundry APIs into foundry namespace

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

* Fix Foundry pyproject formatting

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

* Split provider samples by Foundry surface

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

* Restore hosting sample requirements

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

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

* updated tests

* udpated foundry integration tests

* removed dist from azurefunctions tests

* Use separate Foundry clients for concurrent agents

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

* fix client setup in azfunc and durable

* disabled two tests

* updated setup for some function and durable tests

* improved azure openai setup with new clients

* ignore deprecated

* fixes

* skip 11

* remove openai assistants int tests

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-25 09:56:29 +00:00
committed by GitHub
co-authored by Copilot
parent 4b533608b6
commit 5e056b672e
485 changed files with 9784 additions and 12084 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/README.md
+4 -4
View File
@@ -14,8 +14,8 @@ Highlights
```bash
pip install agent-framework-core --pre
# Optional: Add Azure AI integration
pip install agent-framework-azure-ai --pre
# Optional: Add Azure AI Foundry integration
pip install agent-framework-foundry --pre
```
Supported Platforms:
@@ -36,8 +36,8 @@ AZURE_OPENAI_API_KEY=...
AZURE_OPENAI_ENDPOINT=...
AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=...
...
AZURE_AI_PROJECT_ENDPOINT=...
AZURE_AI_MODEL_DEPLOYMENT_NAME=...
FOUNDRY_PROJECT_ENDPOINT=...
FOUNDRY_MODEL=...
```
You can also override environment variables by explicitly passing configuration parameters to the chat client constructor:
python/packages/core/agent_framework/_types.py
+46 -8
View File
@@ -1995,6 +1995,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
messages: Message | Sequence[Message] | None = None,
response_id: str | None = None,
conversation_id: str | None = None,
model: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: FinishReasonLiteral | FinishReason | None = None,
@@ -2011,8 +2012,9 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
messages: A single Message or sequence of Message objects to include in the response.
response_id: Optional ID of the chat response.
conversation_id: Optional identifier for the state of the conversation.
model_id: Optional model ID used in the creation of the chat response.
created_at: Optional timestamp for the chat response.
model: Optional model used in the creation of the chat response.
model_id: Deprecated alias for ``model``.
created_at: Optional timestamp for when the response was created.
finish_reason: Optional reason for the chat response (e.g., "stop", "length", "tool_calls").
usage_details: Optional usage details for the chat response.
value: Optional value of the structured output.
@@ -2022,6 +2024,8 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
additional_properties: Optional additional properties associated with the chat response.
raw_representation: Optional raw representation of the chat response from an underlying implementation.
"""
if model_id is not None and model is None:
model = model_id
if messages is None:
self.messages: list[Message] = []
elif isinstance(messages, Message):
@@ -2039,7 +2043,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
self.messages = processed_messages
self.response_id = response_id
self.conversation_id = conversation_id
self.model_id = model_id
self.model = model
self.created_at = created_at
self.finish_reason = finish_reason
self.usage_details = usage_details
@@ -2052,6 +2056,15 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
self.continuation_token = continuation_token
self.raw_representation: Any | list[Any] | None = raw_representation
@property
def model_id(self) -> str | None:
"""Deprecated alias for :attr:`model`."""
return self.model
@model_id.setter
def model_id(self, value: str | None) -> None:
self.model = value
@overload
@classmethod
def from_updates(
@@ -2249,6 +2262,7 @@ class ChatResponseUpdate(SerializationMixin):
response_id: str | None = None,
message_id: str | None = None,
conversation_id: str | None = None,
model: str | None = None,
model_id: str | None = None,
created_at: CreatedAtT | None = None,
finish_reason: FinishReasonLiteral | FinishReason | None = None,
@@ -2265,7 +2279,8 @@ class ChatResponseUpdate(SerializationMixin):
response_id: Optional ID of the response of which this update is a part.
message_id: Optional ID of the message of which this update is a part.
conversation_id: Optional identifier for the state of the conversation of which this update is a part
model_id: Optional model ID associated with this response update.
model: Optional model associated with this response update.
model_id: Deprecated alias for ``model``.
created_at: Optional timestamp for the chat response update.
finish_reason: Optional finish reason for the operation.
continuation_token: Optional token for resuming a long-running background operation.
@@ -2275,6 +2290,8 @@ class ChatResponseUpdate(SerializationMixin):
from an underlying implementation.
"""
if model_id is not None and model is None:
model = model_id
# Handle contents - support dict conversion for from_dict
if contents is None:
self.contents: list[Content] = []
@@ -2294,7 +2311,7 @@ class ChatResponseUpdate(SerializationMixin):
self.response_id = response_id
self.message_id = message_id
self.conversation_id = conversation_id
self.model_id = model_id
self.model = model
self.created_at = created_at
self.finish_reason = finish_reason
self.continuation_token = continuation_token
@@ -2304,6 +2321,15 @@ class ChatResponseUpdate(SerializationMixin):
)
self.raw_representation = raw_representation
@property
def model_id(self) -> str | None:
"""Deprecated alias for :attr:`model`."""
return self.model
@model_id.setter
def model_id(self, value: str | None) -> None:
self.model = value
@property
def text(self) -> str:
"""Returns the concatenated text of all contents in the update."""
@@ -3418,7 +3444,7 @@ class Embedding(Generic[EmbeddingT]):
Args:
vector: The embedding vector data.
model_id: The model used to generate this embedding.
model: The model used to generate this embedding.
dimensions: Explicit dimension count (computed from vector length if omitted).
created_at: Timestamp of when the embedding was generated.
additional_properties: Additional metadata.
@@ -3430,7 +3456,7 @@ class Embedding(Generic[EmbeddingT]):
embedding = Embedding(
vector=[0.1, 0.2, 0.3],
model_id="text-embedding-3-small",
model="text-embedding-3-small",
)
assert embedding.dimensions == 3
"""
@@ -3439,19 +3465,31 @@ class Embedding(Generic[EmbeddingT]):
self,
vector: EmbeddingT,
*,
model: str | None = None,
model_id: str | None = None,
dimensions: int | None = None,
created_at: datetime | None = None,
additional_properties: dict[str, Any] | None = None,
) -> None:
if model_id is not None and model is None:
model = model_id
self.vector = vector
self._dimensions = dimensions
self.model_id = model_id
self.model = model
self.created_at = created_at
self.additional_properties = (
_restore_compaction_annotation_in_additional_properties(additional_properties) or {}
)
@property
def model_id(self) -> str | None:
"""Deprecated alias for :attr:`model`."""
return self.model
@model_id.setter
def model_id(self, value: str | None) -> None:
self.model = value
@property
def dimensions(self) -> int | None:
"""Return the number of dimensions in the embedding vector.
python/packages/core/agent_framework/_workflows/_events.py
+1 -1
View File
@@ -121,7 +121,7 @@ WorkflowEventType = Literal[
"executor_completed", # Executor handler completed (use .executor_id, .data)
"executor_failed", # Executor handler raised error (use .executor_id, .details)
# Orchestration event types (use .data for typed payload)
"group_chat", # Group chat orchestrator events (use .data as GroupChatRequestSentEvent | GroupChatResponseReceivedEvent) # noqa: E501
"group_chat", # Group chat orchestrator events (use .data as GroupChatRequestSentEvent | GroupChatResponseReceivedEvent) # noqa: E501
"handoff_sent", # Handoff routing events (use .data as HandoffSentEvent)
"magentic_orchestrator", # Magentic orchestrator events (use .data as MagenticOrchestratorEvent)
]
python/packages/core/agent_framework/azure/__init__.py
+12 -22
View File
@@ -2,16 +2,7 @@
"""Azure integration namespace for optional Agent Framework connectors.
This module lazily re-exports objects from optional Azure connector packages and
built-in core Azure OpenAI modules.
Supported classes include:
- AzureAIClient
- AzureAIAgentClient
- AzureOpenAIChatClient
- AzureOpenAIResponsesClient
- AzureAISearchContextProvider
- DurableAIAgent
This module lazily re-exports objects from optional Azure connector packages.
"""
import importlib
@@ -30,18 +21,17 @@ _IMPORTS: dict[str, tuple[str, str]] = {
"AzureAISearchSettings": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"),
"AzureAISettings": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureAIAgentsProvider": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureCredentialTypes": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
"AzureTokenProvider": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
"FoundryMemoryProvider": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "agent-framework-core"),
"AzureOpenAIAssistantsOptions": ("agent_framework.azure._assistants_client", "agent-framework-core"),
"AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "agent-framework-core"),
"AzureOpenAIChatOptions": ("agent_framework.azure._chat_client", "agent-framework-core"),
"AzureOpenAIEmbeddingClient": ("agent_framework.azure._embedding_client", "agent-framework-core"),
"AzureOpenAIResponsesClient": ("agent_framework.azure._responses_client", "agent-framework-core"),
"AzureOpenAIResponsesOptions": ("agent_framework.azure._responses_client", "agent-framework-core"),
"AzureOpenAISettings": ("agent_framework.azure._shared", "agent-framework-core"),
"AzureUserSecurityContext": ("agent_framework.azure._chat_client", "agent-framework-core"),
"AzureCredentialTypes": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureTokenProvider": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIAssistantsClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIAssistantsOptions": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIChatClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIChatOptions": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIEmbeddingClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIResponsesClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAIResponsesOptions": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureOpenAISettings": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"AzureUserSecurityContext": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
"DurableAIAgent": ("agent_framework_durabletask", "agent-framework-durabletask"),
"DurableAIAgentClient": ("agent_framework_durabletask", "agent-framework-durabletask"),
"DurableAIAgentOrchestrationContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
python/packages/core/agent_framework/azure/__init__.pyi
+24 -10
View File
@@ -1,5 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
# Type stubs for the agent_framework.azure lazy-loading namespace.
# Install the relevant packages for full type support.
from agent_framework_azure_ai import (
AzureAIAgentClient,
AzureAIAgentsProvider,
@@ -7,9 +10,23 @@ from agent_framework_azure_ai import (
AzureAIProjectAgentOptions,
AzureAIProjectAgentProvider,
AzureAISettings,
FoundryMemoryProvider,
AzureCredentialTypes,
AzureOpenAIAssistantsClient,
AzureOpenAIAssistantsOptions,
AzureOpenAIChatClient,
AzureOpenAIChatOptions,
AzureOpenAIEmbeddingClient,
AzureOpenAIResponsesClient,
AzureOpenAIResponsesOptions,
AzureOpenAISettings,
AzureTokenProvider,
AzureUserSecurityContext,
RawAzureAIClient,
)
from agent_framework_azure_ai_search import (
AzureAISearchContextProvider,
AzureAISearchSettings,
)
from agent_framework_azure_ai_search import AzureAISearchContextProvider, AzureAISearchSettings
from agent_framework_azurefunctions import AgentFunctionApp
from agent_framework_durabletask import (
AgentCallbackContext,
@@ -20,13 +37,6 @@ from agent_framework_durabletask import (
DurableAIAgentWorker,
)
from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
from agent_framework.azure._chat_client import AzureOpenAIChatClient
from agent_framework.azure._embedding_client import AzureOpenAIEmbeddingClient
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from agent_framework.azure._responses_client import AzureOpenAIResponsesClient
from agent_framework.azure._shared import AzureOpenAISettings
__all__ = [
"AgentCallbackContext",
"AgentFunctionApp",
@@ -41,14 +51,18 @@ __all__ = [
"AzureAISettings",
"AzureCredentialTypes",
"AzureOpenAIAssistantsClient",
"AzureOpenAIAssistantsOptions",
"AzureOpenAIChatClient",
"AzureOpenAIChatOptions",
"AzureOpenAIEmbeddingClient",
"AzureOpenAIResponsesClient",
"AzureOpenAIResponsesOptions",
"AzureOpenAISettings",
"AzureTokenProvider",
"AzureUserSecurityContext",
"DurableAIAgent",
"DurableAIAgentClient",
"DurableAIAgentOrchestrationContext",
"DurableAIAgentWorker",
"FoundryMemoryProvider",
"RawAzureAIClient",
]
python/packages/core/agent_framework/azure/_assistants_client.py
-194
View File
@@ -1,194 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Mapping
from typing import Any, ClassVar, Generic
from openai.lib.azure import AsyncAzureOpenAI
from .._settings import load_settings
from ..openai import OpenAIAssistantsClient
from ..openai._assistants_client import OpenAIAssistantsOptions
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
from ._shared import AzureOpenAISettings, _apply_azure_defaults # pyright: ignore[reportPrivateUsage]
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
# region Azure OpenAI Assistants Options TypedDict
AzureOpenAIAssistantsOptionsT = TypeVar(
"AzureOpenAIAssistantsOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIAssistantsOptions",
covariant=True,
)
# endregion
class AzureOpenAIAssistantsClient(
OpenAIAssistantsClient[AzureOpenAIAssistantsOptionsT], Generic[AzureOpenAIAssistantsOptionsT]
):
"""Azure OpenAI Assistants client."""
DEFAULT_AZURE_API_VERSION: ClassVar[str] = "2024-05-01-preview"
def __init__(
self,
*,
deployment_name: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
assistant_description: str | None = None,
thread_id: str | None = None,
api_key: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an Azure OpenAI Assistants client.
Keyword Args:
deployment_name: The Azure OpenAI deployment name for the model to use.
Can also be set via environment variable AZURE_OPENAI_CHAT_DEPLOYMENT_NAME.
assistant_id: The ID of an Azure OpenAI assistant to use.
If not provided, a new assistant will be created (and deleted after the request).
assistant_name: The name to use when creating new assistants.
assistant_description: The description to use when creating new assistants.
thread_id: Default thread ID to use for conversations. Can be overridden by
conversation_id property when making a request.
If not provided, a new thread will be created (and deleted after the request).
api_key: The API key to use. If provided will override the env vars or .env file value.
Can also be set via environment variable AZURE_OPENAI_API_KEY.
endpoint: The deployment endpoint. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_ENDPOINT.
base_url: The deployment base URL. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_BASE_URL.
api_version: The deployment API version. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
token_endpoint: The token endpoint to request an Azure token.
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
credential: Azure credential or token provider for authentication. Accepts a
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
bearer token string (sync or async), for example from
``azure.identity.get_bearer_token_provider()``.
default_headers: The default headers mapping of string keys to
string values for HTTP requests.
async_client: An existing client to use.
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAIAssistantsClient
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=gpt-4
# Set AZURE_OPENAI_API_KEY=your-key
client = AzureOpenAIAssistantsClient()
# Or passing parameters directly
client = AzureOpenAIAssistantsClient(
endpoint="https://your-endpoint.openai.azure.com", deployment_name="gpt-4", api_key="your-key"
)
# Or loading from a .env file
client = AzureOpenAIAssistantsClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.azure import AzureOpenAIAssistantsOptions
class MyOptions(AzureOpenAIAssistantsOptions, total=False):
my_custom_option: str
client: AzureOpenAIAssistantsClient[MyOptions] = AzureOpenAIAssistantsClient()
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
chat_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings, default_api_version=self.DEFAULT_AZURE_API_VERSION)
chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
if not chat_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
api_key_secret = azure_openai_settings.get("api_key")
token_scope = azure_openai_settings.get("token_endpoint")
# Resolve credential to token provider
ad_token_provider = None
if not async_client and not api_key_secret and credential:
ad_token_provider = resolve_credential_to_token_provider(credential, token_scope)
if not async_client and not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
# Create Azure client if not provided
if not async_client:
client_params: dict[str, Any] = {
"default_headers": default_headers,
}
if resolved_api_version := azure_openai_settings.get("api_version"):
client_params["api_version"] = resolved_api_version
if api_key_secret:
client_params["api_key"] = api_key_secret.get_secret_value()
elif ad_token_provider:
client_params["azure_ad_token_provider"] = ad_token_provider
if resolved_base_url := azure_openai_settings.get("base_url"):
client_params["base_url"] = str(resolved_base_url)
elif resolved_endpoint := azure_openai_settings.get("endpoint"):
client_params["azure_endpoint"] = str(resolved_endpoint)
async_client = AsyncAzureOpenAI(**client_params)
super().__init__(
model_id=chat_deployment_name,
assistant_id=assistant_id,
assistant_name=assistant_name,
assistant_description=assistant_description,
thread_id=thread_id,
async_client=async_client, # type: ignore[reportArgumentType]
default_headers=default_headers,
)
python/packages/core/agent_framework/azure/_chat_client.py
-349
View File
@@ -1,349 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import logging
import sys
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic, cast
from pydantic import BaseModel
from agent_framework import (
Annotation,
ChatMiddlewareLayer,
ChatResponse,
ChatResponseUpdate,
Content,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
)
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from .._settings import load_settings
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
_apply_azure_defaults, # pyright: ignore[reportPrivateUsage]
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from openai.lib.azure import AsyncAzureOpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from agent_framework._middleware import MiddlewareTypes
logger: logging.Logger = logging.getLogger(__name__)
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
# region Azure OpenAI Chat Options TypedDict
class AzureUserSecurityContext(TypedDict, total=False):
"""User security context for Azure AI applications.
These fields help security operations teams investigate and mitigate security
incidents by providing context about the application and end user.
Learn more: https://learn.microsoft.com/azure/well-architected/service-guides/cosmos-db
"""
application_name: str
"""Name of the application making the request."""
end_user_id: str
"""Unique identifier for the end user (recommend hashing username/email)."""
end_user_tenant_id: str
"""Microsoft 365 tenant ID the end user belongs to. Required for multi-tenant apps."""
source_ip: str
"""The original client's IP address."""
class AzureOpenAIChatOptions(OpenAIChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""Azure OpenAI-specific chat options dict.
Extends OpenAIChatOptions with Azure-specific options including
the "On Your Data" feature and enhanced security context.
See: https://learn.microsoft.com/azure/ai-foundry/openai/reference-preview-latest
Keys:
# Inherited from OpenAIChatOptions/ChatOptions:
model_id: The model to use for the request,
translates to ``model`` in Azure OpenAI API.
temperature: Sampling temperature between 0 and 2.
top_p: Nucleus sampling parameter.
max_tokens: Maximum number of tokens to generate,
translates to ``max_completion_tokens`` in Azure OpenAI API.
stop: Stop sequences.
seed: Random seed for reproducibility.
frequency_penalty: Frequency penalty between -2.0 and 2.0.
presence_penalty: Presence penalty between -2.0 and 2.0.
tools: List of tools (functions) available to the model.
tool_choice: How the model should use tools.
allow_multiple_tool_calls: Whether to allow parallel tool calls,
translates to ``parallel_tool_calls`` in Azure OpenAI API.
response_format: Structured output schema.
metadata: Request metadata for tracking.
user: End-user identifier for abuse monitoring.
store: Whether to store the conversation.
instructions: System instructions for the model.
logit_bias: Token bias values (-100 to 100).
logprobs: Whether to return log probabilities.
top_logprobs: Number of top log probabilities to return (0-20).
# Azure-specific options:
data_sources: Azure "On Your Data" data sources configuration.
user_security_context: Enhanced security context for Azure Defender.
n: Number of chat completions to generate (not recommended, incurs costs).
"""
# Azure-specific options
data_sources: list[dict[str, Any]]
"""Azure "On Your Data" data sources for retrieval-augmented generation.
Supported types: azure_search, azure_cosmos_db, elasticsearch, pinecone, mongo_db.
See: https://learn.microsoft.com/azure/ai-foundry/openai/references/on-your-data
"""
user_security_context: AzureUserSecurityContext
"""Enhanced security context for Azure Defender integration."""
n: int
"""Number of chat completion choices to generate for each input message.
Note: You will be charged based on tokens across all choices. Keep n=1 to minimize costs."""
AzureOpenAIChatOptionsT = TypeVar(
"AzureOpenAIChatOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="AzureOpenAIChatOptions",
covariant=True,
)
# endregion
ChatResponseT = TypeVar("ChatResponseT", ChatResponse, ChatResponseUpdate)
AzureOpenAIChatClientT = TypeVar("AzureOpenAIChatClientT", bound="AzureOpenAIChatClient")
class AzureOpenAIChatClient( # type: ignore[misc]
AzureOpenAIConfigMixin,
FunctionInvocationLayer[AzureOpenAIChatOptionsT],
ChatMiddlewareLayer[AzureOpenAIChatOptionsT],
ChatTelemetryLayer[AzureOpenAIChatOptionsT],
RawOpenAIChatClient[AzureOpenAIChatOptionsT],
Generic[AzureOpenAIChatOptionsT],
):
"""Azure OpenAI Chat completion class with middleware, telemetry, and function invocation support."""
def __init__(
self,
*,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
additional_properties: dict[str, Any] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None:
"""Initialize an Azure OpenAI Chat completion client.
Keyword Args:
api_key: The API key. If provided, will override the value in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_KEY.
deployment_name: The deployment name. If provided, will override the value
(chat_deployment_name) in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_CHAT_DEPLOYMENT_NAME.
endpoint: The deployment endpoint. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_ENDPOINT.
base_url: The deployment base URL. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_BASE_URL.
api_version: The deployment API version. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
token_endpoint: The token endpoint to request an Azure token.
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
credential: Azure credential or token provider for authentication. Accepts a
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
bearer token string (sync or async), for example from
``azure.identity.get_bearer_token_provider()``.
default_headers: The default headers mapping of string keys to
string values for HTTP requests.
async_client: An existing client to use.
additional_properties: Additional properties stored on the client instance.
env_file_path: Use the environment settings file as a fallback to using env vars.
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAIChatClient
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=<model name>
# Set AZURE_OPENAI_API_KEY=your-key
client = AzureOpenAIChatClient()
# Or passing parameters directly
client = AzureOpenAIChatClient(
endpoint="https://your-endpoint.openai.azure.com",
deployment_name="<model name>",
api_key="your-key",
)
# Or loading from a .env file
client = AzureOpenAIChatClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.azure import AzureOpenAIChatOptions
class MyOptions(AzureOpenAIChatOptions, total=False):
my_custom_option: str
client: AzureOpenAIChatClient[MyOptions] = AzureOpenAIChatClient()
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
chat_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings)
chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
if not chat_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
api_version_value = cast(str, azure_openai_settings.get("api_version"))
api_key_value = azure_openai_settings.get("api_key")
token_endpoint_value = azure_openai_settings.get("token_endpoint")
super().__init__(
deployment_name=chat_deployment_name,
endpoint=endpoint_value,
base_url=base_url_value,
api_version=api_version_value,
api_key=api_key_value.get_secret_value() if api_key_value else None,
token_endpoint=token_endpoint_value,
credential=credential,
default_headers=default_headers,
client=async_client,
additional_properties=additional_properties,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
@override
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
"""Parse the choice into a Content object with type='text'.
Overwritten from RawOpenAIChatClient to deal with Azure On Your Data function.
For docs see:
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
"""
message = getattr(choice, "message", None)
if message is None:
message = getattr(choice, "delta", None)
# When you enable asynchronous content filtering in Azure OpenAI, you may receive empty deltas
if message is None: # type: ignore
return None
if hasattr(message, "refusal") and message.refusal:
return Content.from_text(text=message.refusal, raw_representation=choice)
if not message.content:
return None
text_content = Content.from_text(text=message.content, raw_representation=choice)
if not message.model_extra or "context" not in message.model_extra:
return text_content
context_raw: object = cast(object, message.context) # type: ignore[union-attr]
if isinstance(context_raw, str):
try:
context_raw = json.loads(context_raw)
except json.JSONDecodeError:
logger.warning("Context is not a valid JSON string, ignoring context.")
return text_content
if not isinstance(context_raw, dict):
logger.warning("Context is not a valid dictionary, ignoring context.")
return text_content
context = cast(dict[str, Any], context_raw)
# `all_retrieved_documents` is currently not used, but can be retrieved
# through the raw_representation in the text content.
if intent := context.get("intent"):
text_content.additional_properties = {"intent": intent}
citations = context.get("citations")
if isinstance(citations, list) and citations:
annotations: list[Annotation] = []
for citation_raw in cast(list[object], citations):
if not isinstance(citation_raw, dict):
continue
citation = cast(dict[str, Any], citation_raw)
annotations.append(
Annotation(
type="citation",
title=citation.get("title", ""),
url=citation.get("url", ""),
snippet=citation.get("content", ""),
file_id=citation.get("filepath", ""),
tool_name="Azure-on-your-Data",
additional_properties={"chunk_id": citation.get("chunk_id", "")},
raw_representation=citation,
)
)
text_content.annotations = annotations
return text_content
python/packages/core/agent_framework/azure/_embedding_client.py
-141
View File
@@ -1,141 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Mapping
from typing import Generic
from openai.lib.azure import AsyncAzureOpenAI
from agent_framework.observability import EmbeddingTelemetryLayer
from agent_framework.openai import OpenAIEmbeddingOptions
from agent_framework.openai._embedding_client import RawOpenAIEmbeddingClient
from .._settings import load_settings
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
_apply_azure_defaults, # pyright: ignore[reportPrivateUsage]
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
AzureOpenAIEmbeddingOptionsT = TypeVar(
"AzureOpenAIEmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIEmbeddingOptions",
covariant=True,
)
class AzureOpenAIEmbeddingClient(
AzureOpenAIConfigMixin,
EmbeddingTelemetryLayer[str, list[float], AzureOpenAIEmbeddingOptionsT],
RawOpenAIEmbeddingClient[AzureOpenAIEmbeddingOptionsT],
Generic[AzureOpenAIEmbeddingOptionsT],
):
"""Azure OpenAI embedding client with telemetry support.
Keyword Args:
api_key: The API key. If provided, will override the value in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_KEY.
deployment_name: The deployment name. If provided, will override the value
(embedding_deployment_name) in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME.
endpoint: The deployment endpoint.
Can also be set via environment variable AZURE_OPENAI_ENDPOINT.
base_url: The deployment base URL.
Can also be set via environment variable AZURE_OPENAI_BASE_URL.
api_version: The deployment API version.
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
token_endpoint: The token endpoint to request an Azure token.
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAIEmbeddingClient
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-small
# Set AZURE_OPENAI_API_KEY=your-key
client = AzureOpenAIEmbeddingClient()
# Or passing parameters directly
client = AzureOpenAIEmbeddingClient(
endpoint="https://your-endpoint.openai.azure.com",
deployment_name="text-embedding-3-small",
api_key="your-key",
)
result = await client.get_embeddings(["Hello, world!"])
"""
def __init__(
self,
*,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
otel_provider_name: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an Azure OpenAI embedding client."""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
embedding_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings)
embedding_deployment_name = azure_openai_settings.get("embedding_deployment_name")
if not embedding_deployment_name:
raise ValueError(
"Azure OpenAI embedding deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME' environment variable."
)
api_key_secret = azure_openai_settings.get("api_key")
super().__init__(
deployment_name=embedding_deployment_name,
endpoint=azure_openai_settings.get("endpoint"),
base_url=azure_openai_settings.get("base_url"),
api_version=azure_openai_settings.get("api_version") or "",
api_key=api_key_secret.get_secret_value() if api_key_secret else None,
token_endpoint=azure_openai_settings.get("token_endpoint"),
credential=credential,
default_headers=default_headers,
client=async_client,
otel_provider_name=otel_provider_name,
)
python/packages/core/agent_framework/azure/_entra_id_authentication.py
-68
View File
@@ -1,68 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from typing import Union
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from ..exceptions import ChatClientInvalidAuthException
logger: logging.Logger = logging.getLogger(__name__)
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]
"""A callable that returns a bearer token string, either synchronously or asynchronously."""
AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
"""Union of Azure credential types.
Accepts:
- ``TokenCredential`` — synchronous Azure credential (e.g. ``DefaultAzureCredential()``)
- ``AsyncTokenCredential`` — asynchronous Azure credential (e.g. ``azure.identity.aio.DefaultAzureCredential()``)
"""
def resolve_credential_to_token_provider(
credential: AzureCredentialTypes | AzureTokenProvider,
token_endpoint: str | None,
) -> AzureTokenProvider:
"""Convert an Azure credential or token provider into an ``ad_token_provider`` callable.
If the credential is already a callable token provider, it is returned as-is
(``token_endpoint`` is not required in this case).
If it is a ``TokenCredential`` or ``AsyncTokenCredential``, it is wrapped using
``azure.identity.get_bearer_token_provider`` (sync or async variant) which
handles token caching and automatic refresh.
Args:
credential: An Azure credential or token provider callable.
token_endpoint: The token scope/endpoint
(e.g. ``"https://cognitiveservices.azure.com/.default"``).
Required when ``credential`` is a ``TokenCredential`` or ``AsyncTokenCredential``.
Returns:
A callable that returns a bearer token string (sync or async).
Raises:
ServiceInvalidAuthError: If the token endpoint is empty when needed for credential wrapping.
"""
# Already a token provider callable (not a credential object) — use directly
if callable(credential) and not isinstance(credential, (TokenCredential, AsyncTokenCredential)):
return credential
if not token_endpoint:
raise ChatClientInvalidAuthException(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
if isinstance(credential, AsyncTokenCredential):
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
return get_async_bearer_token_provider(credential, token_endpoint)
from azure.identity import get_bearer_token_provider
return get_bearer_token_provider(credential, token_endpoint) # type: ignore[arg-type]
python/packages/core/agent_framework/azure/_responses_client.py
-277
View File
@@ -1,277 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic
from urllib.parse import urljoin, urlparse
from azure.ai.projects.aio import AIProjectClient
from openai import AsyncOpenAI
from .._middleware import ChatMiddlewareLayer
from .._settings import load_settings
from .._telemetry import AGENT_FRAMEWORK_USER_AGENT
from .._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
from ..observability import ChatTelemetryLayer
from ..openai._responses_client import RawOpenAIResponsesClient
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
_apply_azure_defaults, # pyright: ignore[reportPrivateUsage]
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from .._middleware import MiddlewareTypes
from ..openai._responses_client import OpenAIResponsesOptions
AzureOpenAIResponsesOptionsT = TypeVar(
"AzureOpenAIResponsesOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIResponsesOptions",
covariant=True,
)
class AzureOpenAIResponsesClient( # type: ignore[misc]
AzureOpenAIConfigMixin,
FunctionInvocationLayer[AzureOpenAIResponsesOptionsT],
ChatMiddlewareLayer[AzureOpenAIResponsesOptionsT],
ChatTelemetryLayer[AzureOpenAIResponsesOptionsT],
RawOpenAIResponsesClient[AzureOpenAIResponsesOptionsT],
Generic[AzureOpenAIResponsesOptionsT],
):
"""Azure Responses completion class with middleware, telemetry, and function invocation support."""
def __init__(
self,
*,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
project_client: Any | None = None,
project_endpoint: str | None = None,
allow_preview: bool | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure OpenAI Responses client.
The client can be created in two ways:
1. **Direct Azure OpenAI** (default): Provide endpoint, api_key, or credential
to connect directly to an Azure OpenAI deployment.
2. **Foundry project endpoint**: Provide a ``project_client`` or ``project_endpoint``
(with ``credential``) to create the client via an Azure AI Foundry project.
This requires the ``azure-ai-projects`` package to be installed.
Keyword Args:
api_key: The API key. If provided, will override the value in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_KEY.
deployment_name: The deployment name. If provided, will override the value
(responses_deployment_name) in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME.
endpoint: The deployment endpoint. If provided will override the value
in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_ENDPOINT.
base_url: The deployment base URL. If provided will override the value
in the env vars or .env file. Currently, the base_url must end with "/openai/v1/".
Can also be set via environment variable AZURE_OPENAI_BASE_URL.
api_version: The deployment API version. If provided will override the value
in the env vars or .env file. Currently, the api_version must be "preview".
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
token_endpoint: The token endpoint to request an Azure token.
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
credential: Azure credential or token provider for authentication. Accepts a
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
bearer token string (sync or async), for example from
``azure.identity.get_bearer_token_provider()``.
default_headers: The default headers mapping of string keys to
string values for HTTP requests.
async_client: An existing client to use.
project_client: An existing ``AIProjectClient`` (from ``azure.ai.projects.aio``) to use.
The OpenAI client will be obtained via ``project_client.get_openai_client()``.
Requires the ``azure-ai-projects`` package.
project_endpoint: The Azure AI Foundry project endpoint URL.
When provided with ``credential``, an ``AIProjectClient`` will be created
and used to obtain the OpenAI client. Requires the ``azure-ai-projects`` package.
allow_preview: Enables preview opt-in on internally-created ``AIProjectClient``.
env_file_path: Use the environment settings file as a fallback to using env vars.
env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
kwargs: Additional keyword arguments.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAIResponsesClient
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME=gpt-4o
# Set AZURE_OPENAI_API_KEY=your-key
client = AzureOpenAIResponsesClient()
# Or passing parameters directly
client = AzureOpenAIResponsesClient(
endpoint="https://your-endpoint.openai.azure.com", deployment_name="gpt-4o", api_key="your-key"
)
# Or loading from a .env file
client = AzureOpenAIResponsesClient(env_file_path="path/to/.env")
# Using a Foundry project endpoint
from azure.identity import DefaultAzureCredential
client = AzureOpenAIResponsesClient(
project_endpoint="https://your-project.services.ai.azure.com",
deployment_name="gpt-4o",
credential=DefaultAzureCredential(),
)
# Or using an existing AIProjectClient
from azure.ai.projects.aio import AIProjectClient
project_client = AIProjectClient(
endpoint="https://your-project.services.ai.azure.com",
credential=DefaultAzureCredential(),
)
client = AzureOpenAIResponsesClient(
project_client=project_client,
deployment_name="gpt-4o",
)
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.azure import AzureOpenAIResponsesOptions
class MyOptions(AzureOpenAIResponsesOptions, total=False):
my_custom_option: str
client: AzureOpenAIResponsesClient[MyOptions] = AzureOpenAIResponsesClient()
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
if (model_id := kwargs.pop("model_id", None)) and not deployment_name:
deployment_name = str(model_id)
# Project client path: create OpenAI client from an Azure AI Foundry project
if async_client is None and (project_client is not None or project_endpoint is not None):
async_client = self._create_client_from_project(
project_client=project_client,
project_endpoint=project_endpoint,
credential=credential,
allow_preview=allow_preview,
)
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
responses_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings, default_api_version="preview")
# TODO(peterychang): This is a temporary hack to ensure that the base_url is set correctly
# while this feature is in preview.
# But we should only do this if we're on azure. Private deployments may not need this.
endpoint_value = azure_openai_settings.get("endpoint")
if (
not azure_openai_settings.get("base_url")
and endpoint_value
and (hostname := urlparse(str(endpoint_value)).hostname)
and hostname.endswith(".openai.azure.com")
):
azure_openai_settings["base_url"] = urljoin(str(endpoint_value), "/openai/v1/")
responses_deployment_name = azure_openai_settings.get("responses_deployment_name")
if not responses_deployment_name:
raise ValueError(
"Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
)
api_key_secret = azure_openai_settings.get("api_key")
super().__init__(
deployment_name=responses_deployment_name,
endpoint=azure_openai_settings.get("endpoint"),
base_url=azure_openai_settings.get("base_url"),
api_version=azure_openai_settings.get("api_version") or "",
api_key=api_key_secret.get_secret_value() if api_key_secret else None,
token_endpoint=azure_openai_settings.get("token_endpoint"),
credential=credential,
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
@staticmethod
def _create_client_from_project(
*,
project_client: AIProjectClient | None,
project_endpoint: str | None,
credential: AzureCredentialTypes | AzureTokenProvider | None,
allow_preview: bool | None = None,
) -> AsyncOpenAI:
"""Create an AsyncOpenAI client from an Azure AI Foundry project."""
if project_client is not None:
return project_client.get_openai_client()
if not project_endpoint:
raise ValueError("Azure AI project endpoint is required when project_client is not provided.")
if not credential:
raise ValueError("Azure credential is required when using project_endpoint without a project_client.")
project_client_kwargs: dict[str, Any] = {
"endpoint": project_endpoint,
"credential": credential, # type: ignore[arg-type]
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
project_client = AIProjectClient(**project_client_kwargs)
return project_client.get_openai_client()
@override
def _check_model_presence(self, options: dict[str, Any]) -> None:
if not options.get("model"):
if not self.model_id:
raise ValueError("deployment_name must be a non-empty string")
options["model"] = self.model_id
python/packages/core/agent_framework/azure/_shared.py
-223
View File
@@ -1,223 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
import sys
from collections.abc import Mapping
from copy import copy
from typing import Any, ClassVar, Final
from openai import AsyncOpenAI
from openai.lib.azure import AsyncAzureOpenAI
from .._settings import SecretString
from .._telemetry import APP_INFO, prepend_agent_framework_to_user_agent
from ..openai._shared import OpenAIBase
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
logger: logging.Logger = logging.getLogger(__name__)
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
DEFAULT_AZURE_API_VERSION: Final[str] = "2024-10-21"
DEFAULT_AZURE_TOKEN_ENDPOINT: Final[str] = "https://cognitiveservices.azure.com/.default" # noqa: S105
class AzureOpenAISettings(TypedDict, total=False):
"""AzureOpenAI model settings.
Settings are resolved in this order: explicit keyword arguments, values from an
explicitly provided .env file, then environment variables with the prefix
'AZURE_OPENAI_'. If settings are missing after resolution, validation will fail.
Keyword Args:
endpoint: The endpoint of the Azure deployment. 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.
Can be set via environment variable AZURE_OPENAI_ENDPOINT.
chat_deployment_name: 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.
Can be set via environment variable AZURE_OPENAI_CHAT_DEPLOYMENT_NAME.
responses_deployment_name: 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.
Can be set via environment variable AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME.
embedding_deployment_name: The name of the Azure Embedding deployment.
Can be set via environment variable AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME.
api_key: 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.
Can be set via environment variable AZURE_OPENAI_API_KEY.
api_version: The API version to use. The default value is `DEFAULT_AZURE_API_VERSION`.
Can be set via environment variable AZURE_OPENAI_API_VERSION.
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.
Can be set via environment variable AZURE_OPENAI_BASE_URL.
token_endpoint: The token endpoint to use to retrieve the authentication token.
The default value is `DEFAULT_AZURE_TOKEN_ENDPOINT`.
Can be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAISettings
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=gpt-4
# Set AZURE_OPENAI_API_KEY=your-key
settings = load_settings(AzureOpenAISettings, env_prefix="AZURE_OPENAI_")
# Or passing parameters directly
settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
endpoint="https://your-endpoint.openai.azure.com",
chat_deployment_name="gpt-4",
api_key="your-key",
)
# Or loading from a .env file
settings = load_settings(AzureOpenAISettings, env_prefix="AZURE_OPENAI_", env_file_path="path/to/.env")
"""
chat_deployment_name: str | None
responses_deployment_name: str | None
embedding_deployment_name: str | None
endpoint: str | None
base_url: str | None
api_key: SecretString | None
api_version: str | None
token_endpoint: str | None
def _apply_azure_defaults(
settings: AzureOpenAISettings,
default_api_version: str = DEFAULT_AZURE_API_VERSION,
default_token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT,
) -> None:
"""Apply default values for api_version and token_endpoint after loading settings.
Args:
settings: The loaded Azure OpenAI settings dict.
default_api_version: The default API version to use if not set.
default_token_endpoint: The default token endpoint to use if not set.
"""
if not settings.get("api_version"):
settings["api_version"] = default_api_version
if not settings.get("token_endpoint"):
settings["token_endpoint"] = default_token_endpoint
_AZURE_DEFAULTS_APPLIER = _apply_azure_defaults
class AzureOpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an Azure OpenAI service."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.openai"
# Note: INJECTABLE = {"client"} is inherited from OpenAIBase
def __init__(
self,
deployment_name: str,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str = DEFAULT_AZURE_API_VERSION,
api_key: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncOpenAI | None = None,
instruction_role: str | None = None,
**kwargs: Any,
) -> None:
"""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 `str` and `OpenAIModelTypes`.
Args:
deployment_name: Name of the deployment.
endpoint: The specific endpoint URL for the deployment.
base_url: The base URL for Azure services.
api_version: Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
api_key: API key for Azure services.
token_endpoint: Azure AD token scope used to obtain a bearer token from a credential.
credential: Azure credential or token provider for authentication. Accepts a
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
bearer token string (sync or async).
default_headers: Default headers for HTTP requests.
client: An existing client to use.
instruction_role: The role to use for 'instruction' messages, for example, summarization
prompts could use `developer` or `system`.
kwargs: Additional keyword arguments.
"""
# Merge APP_INFO into the headers if it exists
merged_headers = dict(copy(default_headers)) if default_headers else {}
if APP_INFO:
merged_headers.update(APP_INFO)
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
if not client:
# Resolve credential to a token provider if needed
ad_token_provider = None
if not api_key and credential:
ad_token_provider = resolve_credential_to_token_provider(credential, token_endpoint)
if not api_key and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
if not endpoint and not base_url:
raise ValueError("Please provide an endpoint or a base_url")
args: dict[str, Any] = {
"default_headers": merged_headers,
}
if api_version:
args["api_version"] = api_version
if ad_token_provider:
args["azure_ad_token_provider"] = ad_token_provider
if api_key:
args["api_key"] = api_key
if base_url:
args["base_url"] = str(base_url)
if endpoint and not base_url:
args["azure_endpoint"] = str(endpoint)
if deployment_name:
args["azure_deployment"] = deployment_name
if "websocket_base_url" in kwargs:
args["websocket_base_url"] = kwargs.pop("websocket_base_url")
client = AsyncAzureOpenAI(**args)
# Store configuration as instance attributes for serialization
self.endpoint = str(endpoint)
self.base_url = str(base_url)
self.api_version = api_version
self.deployment_name = deployment_name
self.instruction_role = instruction_role
# Store default_headers but filter out USER_AGENT_KEY for serialization
if default_headers:
from .._telemetry import USER_AGENT_KEY
def_headers = {k: v for k, v in default_headers.items() if k != USER_AGENT_KEY}
else:
def_headers = None
self.default_headers = def_headers
super().__init__(model_id=deployment_name, client=client, **kwargs)
python/packages/core/agent_framework/foundry/__init__.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright (c) Microsoft. All rights reserved.
"""Foundry integration namespace for optional Agent Framework connectors.
This module lazily re-exports objects from cloud Foundry and Foundry Local connector packages.
"""
import importlib
from typing import Any
_IMPORTS: dict[str, tuple[str, str]] = {
"FoundryAgent": ("agent_framework_foundry", "agent-framework-foundry"),
"FoundryChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
"FoundryChatOptions": ("agent_framework_foundry", "agent-framework-foundry"),
"FoundryMemoryProvider": ("agent_framework_foundry", "agent-framework-foundry"),
"FoundryLocalChatOptions": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
"FoundryLocalClient": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
"FoundryLocalSettings": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
"RawFoundryAgent": ("agent_framework_foundry", "agent-framework-foundry"),
"RawFoundryAgentChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
"RawFoundryChatClient": ("agent_framework_foundry", "agent-framework-foundry"),
}
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
import_path, package_name = _IMPORTS[name]
try:
return getattr(importlib.import_module(import_path), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The package {package_name} is required to use `{name}`. "
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
) from exc
raise AttributeError(f"Module `foundry` has no attribute {name}.")
def __dir__() -> list[str]:
return list(_IMPORTS.keys())
python/packages/core/agent_framework/foundry/__init__.pyi
+32
View File
@@ -0,0 +1,32 @@
# Copyright (c) Microsoft. All rights reserved.
# Type stubs for the agent_framework.foundry lazy-loading namespace.
# Install the relevant packages for full type support.
from agent_framework_foundry import (
FoundryAgent,
FoundryChatClient,
FoundryChatOptions,
FoundryMemoryProvider,
RawFoundryAgent,
RawFoundryAgentChatClient,
RawFoundryChatClient,
)
from agent_framework_foundry_local import (
FoundryLocalChatOptions,
FoundryLocalClient,
FoundryLocalSettings,
)
__all__ = [
"FoundryAgent",
"FoundryChatClient",
"FoundryChatOptions",
"FoundryLocalChatOptions",
"FoundryLocalClient",
"FoundryLocalSettings",
"FoundryMemoryProvider",
"RawFoundryAgent",
"RawFoundryAgentChatClient",
"RawFoundryChatClient",
]
python/packages/core/agent_framework/openai/__init__.py
+47 -40
View File
@@ -1,48 +1,55 @@
# Copyright (c) Microsoft. All rights reserved.
"""OpenAI namespace for built-in Agent Framework clients.
"""OpenAI namespace for Agent Framework clients.
This module re-exports objects from the core OpenAI implementation modules in
``agent_framework.openai``.
This module lazily re-exports objects from the ``agent-framework-openai`` package.
Install it with: ``pip install agent-framework-openai``
Supported classes include:
- OpenAIChatClient
- OpenAIResponsesClient
- OpenAIAssistantsClient
- OpenAIAssistantProvider
- OpenAIChatClient (Responses API)
- OpenAIChatCompletionClient (Chat Completions API)
- OpenAIEmbeddingClient
- OpenAIAssistantsClient (deprecated)
"""
from ._assistant_provider import OpenAIAssistantProvider
from ._assistants_client import (
AssistantToolResources,
OpenAIAssistantsClient,
OpenAIAssistantsOptions,
)
from ._chat_client import OpenAIChatClient, OpenAIChatOptions
from ._embedding_client import OpenAIEmbeddingClient, OpenAIEmbeddingOptions
from ._exceptions import ContentFilterResultSeverity, OpenAIContentFilterException
from ._responses_client import (
OpenAIContinuationToken,
OpenAIResponsesClient,
OpenAIResponsesOptions,
RawOpenAIResponsesClient,
)
from ._shared import OpenAISettings
import importlib
from typing import Any
__all__ = [
"AssistantToolResources",
"ContentFilterResultSeverity",
"OpenAIAssistantProvider",
"OpenAIAssistantsClient",
"OpenAIAssistantsOptions",
"OpenAIChatClient",
"OpenAIChatOptions",
"OpenAIContentFilterException",
"OpenAIContinuationToken",
"OpenAIEmbeddingClient",
"OpenAIEmbeddingOptions",
"OpenAIResponsesClient",
"OpenAIResponsesOptions",
"OpenAISettings",
"RawOpenAIResponsesClient",
]
_IMPORTS: dict[str, tuple[str, str]] = {
"OpenAIChatClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIChatOptions": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIContinuationToken": ("agent_framework_openai", "agent-framework-openai"),
"RawOpenAIChatClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIChatCompletionOptions": ("agent_framework_openai", "agent-framework-openai"),
"RawOpenAIChatCompletionClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIEmbeddingClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIEmbeddingOptions": ("agent_framework_openai", "agent-framework-openai"),
"OpenAISettings": ("agent_framework_openai", "agent-framework-openai"),
"ContentFilterResultSeverity": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIContentFilterException": ("agent_framework_openai", "agent-framework-openai"),
"AssistantToolResources": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIAssistantProvider": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIAssistantsClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIAssistantsOptions": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIResponsesClient": ("agent_framework_openai", "agent-framework-openai"),
"OpenAIResponsesOptions": ("agent_framework_openai", "agent-framework-openai"),
"RawOpenAIResponsesClient": ("agent_framework_openai", "agent-framework-openai"),
}
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
import_path, package_name = _IMPORTS[name]
try:
return getattr(importlib.import_module(import_path), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The package {package_name} is required to use `{name}`. "
f"Please use `pip install {package_name}`, or update your requirements.txt or pyproject.toml file."
) from exc
raise AttributeError(f"Module `openai` has no attribute {name}.")
def __dir__() -> list[str]:
return list(_IMPORTS.keys())
python/packages/core/agent_framework/openai/__init__.pyi
+48
View File
@@ -0,0 +1,48 @@
# Copyright (c) Microsoft. All rights reserved.
# Type stubs for the agent_framework.openai lazy-loading namespace.
# Install agent-framework-openai for full type support.
from agent_framework_openai import (
AssistantToolResources,
ContentFilterResultSeverity,
OpenAIAssistantProvider,
OpenAIAssistantsClient,
OpenAIAssistantsOptions,
OpenAIChatClient,
OpenAIChatCompletionClient,
OpenAIChatCompletionOptions,
OpenAIChatOptions,
OpenAIContentFilterException,
OpenAIContinuationToken,
OpenAIEmbeddingClient,
OpenAIEmbeddingOptions,
OpenAIResponsesClient,
OpenAIResponsesOptions,
OpenAISettings,
RawOpenAIChatClient,
RawOpenAIChatCompletionClient,
RawOpenAIResponsesClient,
)
__all__ = [
"AssistantToolResources",
"ContentFilterResultSeverity",
"OpenAIAssistantProvider",
"OpenAIAssistantsClient",
"OpenAIAssistantsOptions",
"OpenAIChatClient",
"OpenAIChatCompletionClient",
"OpenAIChatCompletionOptions",
"OpenAIChatOptions",
"OpenAIContentFilterException",
"OpenAIContinuationToken",
"OpenAIEmbeddingClient",
"OpenAIEmbeddingOptions",
"OpenAIResponsesClient",
"OpenAIResponsesOptions",
"OpenAISettings",
"RawOpenAIChatClient",
"RawOpenAIChatCompletionClient",
"RawOpenAIResponsesClient",
]
python/packages/core/agent_framework/openai/_assistant_provider.py
-565
View File
@@ -1,565 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from typing import TYPE_CHECKING, Any, Generic, cast
from openai import AsyncOpenAI
from openai.types.beta.assistant import Assistant
from pydantic import BaseModel
from agent_framework._settings import SecretString, load_settings
from .._agents import Agent
from .._middleware import MiddlewareTypes
from .._sessions import BaseContextProvider
from .._tools import FunctionTool, ToolTypes, normalize_tools
from ._assistants_client import OpenAIAssistantsClient
from ._shared import OpenAISettings, from_assistant_tools, to_assistant_tools
if TYPE_CHECKING:
from ._assistants_client import OpenAIAssistantsOptions
if sys.version_info >= (3, 13):
from typing import TypeVar # type:ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type:ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self, TypedDict # type:ignore # pragma: no cover
else:
from typing_extensions import Self, TypedDict # type:ignore # pragma: no cover
# Type variable for options - allows typed OpenAIAssistantProvider[OptionsCoT] returns
# Default matches OpenAIAssistantsClient's default options type
OptionsCoT = TypeVar(
"OptionsCoT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIAssistantsOptions",
covariant=True,
)
class OpenAIAssistantProvider(Generic[OptionsCoT]):
"""Provider for creating Agent instances from OpenAI Assistants API.
This provider allows you to create, retrieve, and wrap OpenAI Assistants
as Agent instances for use in the agent framework.
Examples:
Basic usage with automatic client creation:
.. code-block:: python
from agent_framework.openai import OpenAIAssistantProvider
# Uses OPENAI_API_KEY environment variable
provider = OpenAIAssistantProvider()
# Create a new assistant
agent = await provider.create_agent(
name="MyAssistant",
model="gpt-4",
instructions="You are a helpful assistant.",
tools=[my_function],
)
result = await agent.run("Hello!")
Using an existing client:
.. code-block:: python
from openai import AsyncOpenAI
from agent_framework.openai import OpenAIAssistantProvider
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Get an existing assistant by ID
agent = await provider.get_agent(
assistant_id="asst_123",
tools=[my_function], # Provide implementations for function tools
)
Wrapping an SDK Assistant object:
.. code-block:: python
# Fetch assistant directly via SDK
assistant = await client.beta.assistants.retrieve("asst_123")
# Wrap without additional HTTP call
agent = provider.as_agent(assistant, tools=[my_function])
"""
def __init__(
self,
client: AsyncOpenAI | None = None,
*,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize the OpenAI Assistant Provider.
Args:
client: An existing AsyncOpenAI client to use. If not provided,
a new client will be created using the other parameters.
Keyword Args:
api_key: OpenAI API key. Can also be set via OPENAI_API_KEY env var.
org_id: OpenAI organization ID. Can also be set via OPENAI_ORG_ID env var.
base_url: Base URL for the OpenAI API. Can also be set via OPENAI_BASE_URL env var.
env_file_path: Path to .env file for configuration.
env_file_encoding: Encoding of the .env file.
Raises:
ValueError: If no client is provided and API key is missing.
Examples:
.. code-block:: python
# Using environment variables
provider = OpenAIAssistantProvider()
# Using explicit API key
provider = OpenAIAssistantProvider(api_key="sk-...")
# Using existing client
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
"""
self._client: AsyncOpenAI | None = client
self._should_close_client: bool = client is None
if client is None:
# Load settings and create client
settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
org_id=org_id,
base_url=base_url,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
api_key_setting = settings.get("api_key")
if not api_key_setting:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
# Get API key value
api_key_value: str | Callable[[], str | Awaitable[str]]
if isinstance(api_key_setting, SecretString):
api_key_value = api_key_setting.get_secret_value()
else:
api_key_value = api_key_setting
# Create client
client_args: dict[str, Any] = {"api_key": api_key_value}
if org_id_value := settings.get("org_id"):
client_args["organization"] = org_id_value
if base_url_value := settings.get("base_url"):
client_args["base_url"] = base_url_value
self._client = AsyncOpenAI(**client_args)
async def __aenter__(self) -> Self:
"""Async context manager entry."""
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Async context manager exit."""
await self.close()
async def close(self) -> None:
"""Close the provider and clean up resources.
If the provider created its own client, it will be closed.
If an external client was provided, it will not be closed.
"""
if self._should_close_client and self._client is not None:
await self._client.close()
async def create_agent(
self,
*,
name: str,
model: str,
instructions: str | None = None,
description: str | None = None,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
metadata: dict[str, str] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a new assistant on OpenAI and return a Agent.
This method creates a new assistant on the OpenAI service and wraps it
in a Agent instance. The assistant will persist on OpenAI until deleted.
Keyword Args:
name: The name of the assistant (required).
model: The model ID to use, e.g., "gpt-4", "gpt-4o" (required).
instructions: System instructions for the assistant.
description: A description of the assistant.
tools: Tools available to the assistant. Can include:
- FunctionTool instances or callables decorated with @tool
- Dict-based tools from OpenAIAssistantsClient.get_code_interpreter_tool()
- Dict-based tools from OpenAIAssistantsClient.get_file_search_tool()
- Raw tool dictionaries
metadata: Metadata to attach to the assistant (max 16 key-value pairs).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
Include ``response_format`` here for structured output responses.
middleware: MiddlewareTypes for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the created assistant.
Raises:
ValueError: If assistant creation fails.
Examples:
.. code-block:: python
provider = OpenAIAssistantProvider()
# Create with function tools
agent = await provider.create_agent(
name="WeatherBot",
model="gpt-4",
instructions="You are a helpful weather assistant.",
tools=[get_weather],
)
# Create with structured output
agent = await provider.create_agent(
name="StructuredBot",
model="gpt-4",
default_options={"response_format": MyPydanticModel},
)
"""
# Normalize tools
normalized_tools = normalize_tools(tools)
assistant_tools: list[FunctionTool | MutableMapping[str, Any]] = [
tool for tool in normalized_tools if isinstance(tool, (FunctionTool, MutableMapping))
]
api_tools = to_assistant_tools(assistant_tools) if assistant_tools else []
# Extract response_format from default_options if present
opts = dict(default_options) if default_options else {}
response_format = opts.get("response_format")
# Build assistant creation parameters
create_params: dict[str, Any] = {
"model": model,
"name": name,
}
if instructions is not None:
create_params["instructions"] = instructions
if description is not None:
create_params["description"] = description
if api_tools:
create_params["tools"] = api_tools
if metadata is not None:
create_params["metadata"] = metadata
# Handle response format for OpenAI API
if response_format is not None and isinstance(response_format, type) and issubclass(response_format, BaseModel):
create_params["response_format"] = {
"type": "json_schema",
"json_schema": {
"name": response_format.__name__,
"schema": response_format.model_json_schema(),
"strict": True,
},
}
# Create the assistant
if not self._client:
raise RuntimeError("OpenAI client is not initialized.")
assistant = await self._client.beta.assistants.create(**create_params) # type: ignore[reportDeprecated]
# Create Agent - pass default_options which contains response_format
return self._create_chat_agent_from_assistant(
assistant=assistant,
tools=normalized_tools,
instructions=instructions,
middleware=middleware,
context_providers=context_providers,
default_options=default_options,
)
async def get_agent(
self,
assistant_id: str,
*,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
instructions: str | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Retrieve an existing assistant by ID and return a Agent.
This method fetches an existing assistant from OpenAI by its ID
and wraps it in a Agent instance.
Args:
assistant_id: The ID of the assistant to retrieve (e.g., "asst_123").
Keyword Args:
tools: Function tools to make available. IMPORTANT: If the assistant
was created with function tools, you MUST provide matching
implementations here. Hosted tools (code_interpreter, file_search)
are automatically included.
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the retrieved assistant.
Raises:
RuntimeError: If the assistant cannot be retrieved.
ValueError: If required function tools are missing.
Examples:
.. code-block:: python
provider = OpenAIAssistantProvider()
# Get assistant without function tools
agent = await provider.get_agent(assistant_id="asst_123")
# Get assistant with function tools
agent = await provider.get_agent(
assistant_id="asst_456",
tools=[get_weather, search_database], # Implementations required!
)
"""
# Fetch the assistant
if not self._client:
raise RuntimeError("OpenAI client is not initialized.")
assistant = await self._client.beta.assistants.retrieve(assistant_id) # type: ignore[reportDeprecated]
# Use as_agent to wrap it
return self.as_agent(
assistant=assistant,
tools=tools,
instructions=instructions,
default_options=default_options,
middleware=middleware,
context_providers=context_providers,
)
def as_agent(
self,
assistant: Assistant,
*,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
instructions: str | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Wrap an existing SDK Assistant object as a Agent.
This method does NOT make any HTTP calls. It simply wraps an already-
fetched Assistant object in a Agent.
Args:
assistant: The OpenAI Assistant SDK object to wrap.
Keyword Args:
tools: Function tools to make available. If the assistant has
function tools defined, you MUST provide matching implementations.
Hosted tools (code_interpreter, file_search) are automatically included.
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the assistant.
Raises:
ValueError: If required function tools are missing.
Examples:
.. code-block:: python
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Fetch assistant via SDK
assistant = await client.beta.assistants.retrieve("asst_123")
# Wrap without additional HTTP call
agent = provider.as_agent(
assistant,
tools=[my_function],
instructions="Custom instructions override",
)
"""
# Validate that required function tools are provided
self._validate_function_tools(assistant.tools or [], tools)
# Merge hosted tools with user-provided function tools
merged_tools = self._merge_tools(assistant.tools or [], tools)
# Create Agent
return self._create_chat_agent_from_assistant(
assistant=assistant,
tools=merged_tools,
instructions=instructions,
default_options=default_options,
middleware=middleware,
context_providers=context_providers,
)
def _validate_function_tools(
self,
assistant_tools: list[Any],
provided_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
) -> None:
"""Validate that required function tools are provided.
Args:
assistant_tools: Tools defined on the assistant.
provided_tools: Tools provided by the user.
Raises:
ValueError: If a required function tool is missing.
"""
# Get function tool names from assistant
required_functions: set[str] = set()
for tool in assistant_tools:
if (
hasattr(tool, "type")
and tool.type == "function"
and hasattr(tool, "function")
and hasattr(tool.function, "name")
):
required_functions.add(tool.function.name)
if not required_functions:
return # No function tools required
# Get provided function names using normalize_tools
provided_functions: set[str] = set()
if provided_tools is not None:
normalized = normalize_tools(provided_tools)
for tool in normalized:
if isinstance(tool, FunctionTool):
provided_functions.add(tool.name)
elif isinstance(tool, Mapping):
typed_tool = cast(Mapping[str, Any], tool)
raw_func_spec = typed_tool.get("function")
if isinstance(raw_func_spec, Mapping):
typed_func_spec = cast(Mapping[str, Any], raw_func_spec)
raw_name = typed_func_spec.get("name")
if isinstance(raw_name, str) and raw_name:
provided_functions.add(raw_name)
# Check for missing functions
missing = required_functions - provided_functions
if missing:
missing_list = ", ".join(sorted(missing))
raise ValueError(
f"Assistant requires function tool(s) '{missing_list}' but no implementation was provided. "
f"Please pass the function implementation(s) in the 'tools' parameter."
)
def _merge_tools(
self,
assistant_tools: list[Any],
user_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
) -> list[FunctionTool | MutableMapping[str, Any] | Any]:
"""Merge hosted tools from assistant with user-provided function tools.
Args:
assistant_tools: Tools defined on the assistant.
user_tools: Tools provided by the user.
Returns:
A list of all tools (hosted tools + user function implementations).
"""
merged: list[FunctionTool | MutableMapping[str, Any] | Any] = []
# Add hosted tools from assistant using shared conversion
hosted_tools = from_assistant_tools(assistant_tools)
merged.extend(hosted_tools)
# Add user-provided tools (normalized)
if user_tools is not None:
normalized_user_tools = normalize_tools(user_tools)
merged.extend(normalized_user_tools)
return merged
def _create_chat_agent_from_assistant(
self,
assistant: Assistant,
tools: list[FunctionTool | MutableMapping[str, Any] | Any] | None,
instructions: str | None,
middleware: Sequence[MiddlewareTypes] | None,
context_providers: Sequence[BaseContextProvider] | None,
default_options: OptionsCoT | None = None,
**kwargs: Any,
) -> Agent[OptionsCoT]:
"""Create a Agent from an Assistant.
Args:
assistant: The OpenAI Assistant object.
tools: Tools for the agent.
instructions: Instructions override.
middleware: MiddlewareTypes for the agent.
context_providers: Context providers for the agent.
default_options: Default chat options for the agent (may include response_format).
**kwargs: Additional arguments passed to Agent.
Returns:
A configured Agent instance.
"""
# Create the chat client with the assistant
client = OpenAIAssistantsClient(
model_id=assistant.model,
assistant_id=assistant.id,
assistant_name=assistant.name,
assistant_description=assistant.description,
async_client=self._client,
)
# Use instructions from assistant if not overridden
final_instructions = instructions if instructions is not None else assistant.instructions
# Create and return Agent
return Agent(
client=client,
id=assistant.id,
name=assistant.name,
description=assistant.description,
instructions=final_instructions,
tools=tools if tools else None,
middleware=middleware,
context_providers=context_providers,
default_options=default_options, # type: ignore[arg-type]
**kwargs,
)
python/packages/core/agent_framework/openai/_assistants_client.py
-952
View File
@@ -1,952 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import logging
import sys
from collections.abc import (
AsyncIterable,
Awaitable,
Callable,
Mapping,
MutableMapping,
Sequence,
)
from typing import TYPE_CHECKING, Any, Generic, Literal, TypedDict, cast
from openai import AsyncOpenAI
from openai.types.beta.threads import (
FileCitationAnnotation,
FileCitationDeltaAnnotation,
FilePathAnnotation,
FilePathDeltaAnnotation,
ImageURLContentBlockParam,
ImageURLParam,
MessageContentPartParam,
MessageDeltaEvent,
Run,
TextContentBlockParam,
TextDeltaBlock,
)
from openai.types.beta.threads import (
Message as ThreadMessage,
)
from openai.types.beta.threads.run_create_params import AdditionalMessage
from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
from openai.types.beta.threads.runs import RunStep
from pydantic import BaseModel
from .._clients import BaseChatClient
from .._middleware import ChatMiddlewareLayer
from .._settings import load_settings
from .._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
normalize_tools,
)
from .._types import (
Annotation,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
ResponseStream,
TextSpanRegion,
UsageDetails,
)
from ..observability import ChatTelemetryLayer
from ._shared import OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self, TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from .._middleware import MiddlewareTypes
logger = logging.getLogger("agent_framework.openai")
# region OpenAI Assistants Options TypedDict
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
class VectorStoreToolResource(TypedDict, total=False):
"""Vector store configuration for file search tool resources."""
vector_store_ids: list[str]
"""IDs of vector stores attached to this assistant."""
class CodeInterpreterToolResource(TypedDict, total=False):
"""Code interpreter tool resource configuration."""
file_ids: list[str]
"""File IDs accessible by the code interpreter tool. Max 20 files per assistant."""
class AssistantToolResources(TypedDict, total=False):
"""Tool resources attached to the assistant.
See: https://platform.openai.com/docs/api-reference/assistants/createAssistant#assistants-createassistant-tool_resources
"""
code_interpreter: CodeInterpreterToolResource
"""Resources for code interpreter tool, including file IDs."""
file_search: VectorStoreToolResource
"""Resources for file search tool, including vector store IDs."""
class OpenAIAssistantsOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""OpenAI Assistants API-specific options dict.
Extends base ChatOptions with Assistants API-specific parameters
for creating and running assistants.
See: https://platform.openai.com/docs/api-reference/assistants
Keys:
# Inherited from ChatOptions:
model_id: The model to use for the assistant,
translates to ``model`` in OpenAI API.
temperature: Sampling temperature between 0 and 2.
top_p: Nucleus sampling parameter.
max_tokens: Maximum number of tokens to generate,
translates to ``max_completion_tokens`` in OpenAI API.
tools: List of tools (functions, code_interpreter, file_search).
tool_choice: How the model should use tools.
allow_multiple_tool_calls: Whether to allow parallel tool calls,
translates to ``parallel_tool_calls`` in OpenAI API.
response_format: Structured output schema.
metadata: Request metadata for tracking.
# Options not supported in Assistants API (inherited but unused):
stop: Not supported.
seed: Not supported (use assistant-level configuration instead).
frequency_penalty: Not supported.
presence_penalty: Not supported.
user: Not supported.
store: Not supported.
# Assistants-specific options:
name: Name of the assistant.
description: Description of the assistant.
instructions: System instructions for the assistant.
tool_resources: Resources for tools (file IDs, vector stores).
reasoning_effort: Effort level for o-series reasoning models.
conversation_id: Thread ID to continue conversation in.
"""
# Assistants-specific options
name: str
"""Name of the assistant (max 256 characters)."""
description: str
"""Description of the assistant (max 512 characters)."""
tool_resources: AssistantToolResources
"""Tool-specific resources like file IDs and vector stores."""
reasoning_effort: Literal["low", "medium", "high"]
"""Effort level for o-series reasoning models (o1, o3-mini).
Higher effort = more reasoning time and potentially better results."""
conversation_id: str # type: ignore[misc]
"""Thread ID to continue a conversation in an existing thread."""
# OpenAI/ChatOptions fields not supported in Assistants API
stop: None # type: ignore[misc]
"""Not supported in Assistants API."""
seed: None # type: ignore[misc]
"""Not supported in Assistants API (use assistant-level configuration)."""
frequency_penalty: None # type: ignore[misc]
"""Not supported in Assistants API."""
presence_penalty: None # type: ignore[misc]
"""Not supported in Assistants API."""
user: None # type: ignore[misc]
"""Not supported in Assistants API."""
store: None # type: ignore[misc]
"""Not supported in Assistants API."""
ASSISTANTS_OPTION_TRANSLATIONS: dict[str, str] = {
"model_id": "model",
"max_tokens": "max_completion_tokens",
"allow_multiple_tool_calls": "parallel_tool_calls",
}
"""Maps ChatOptions keys to OpenAI Assistants API parameter names."""
OpenAIAssistantsOptionsT = TypeVar(
"OpenAIAssistantsOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIAssistantsOptions",
covariant=True,
)
# endregion
class OpenAIAssistantsClient( # type: ignore[misc]
OpenAIConfigMixin,
FunctionInvocationLayer[OpenAIAssistantsOptionsT],
ChatMiddlewareLayer[OpenAIAssistantsOptionsT],
ChatTelemetryLayer[OpenAIAssistantsOptionsT],
BaseChatClient[OpenAIAssistantsOptionsT],
Generic[OpenAIAssistantsOptionsT],
):
"""OpenAI Assistants client with middleware, telemetry, and function invocation support."""
# region Hosted Tool Factory Methods
@staticmethod
def get_code_interpreter_tool() -> dict[str, Any]:
"""Create a code interpreter tool configuration for the Assistants API.
Returns:
A dict tool configuration ready to pass to ChatAgent.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIAssistantsClient
# Enable code interpreter
tool = OpenAIAssistantsClient.get_code_interpreter_tool()
agent = ChatAgent(client, tools=[tool])
"""
return {"type": "code_interpreter"}
@staticmethod
def get_file_search_tool(
*,
max_num_results: int | None = None,
) -> dict[str, Any]:
"""Create a file search tool configuration for the Assistants API.
Keyword Args:
max_num_results: Maximum number of results to return from file search.
Returns:
A dict tool configuration ready to pass to ChatAgent.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIAssistantsClient
# Basic file search
tool = OpenAIAssistantsClient.get_file_search_tool()
# With result limit
tool = OpenAIAssistantsClient.get_file_search_tool(max_num_results=10)
agent = ChatAgent(client, tools=[tool])
"""
tool: dict[str, Any] = {"type": "file_search"}
if max_num_results is not None:
tool["file_search"] = {"max_num_results": max_num_results}
return tool
# endregion
def __init__(
self,
*,
model_id: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
assistant_description: str | None = None,
thread_id: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an OpenAI Assistants client.
Keyword Args:
model_id: OpenAI model name, see https://platform.openai.com/docs/models.
Can also be set via environment variable OPENAI_CHAT_MODEL_ID.
assistant_id: The ID of an OpenAI assistant to use.
If not provided, a new assistant will be created (and deleted after the request).
assistant_name: The name to use when creating new assistants.
assistant_description: The description to use when creating new assistants.
thread_id: Default thread ID to use for conversations. Can be overridden by
conversation_id property when making a request.
If not provided, a new thread will be created (and deleted after the request).
api_key: The API key to use. If provided will override the env vars or .env file value.
Can also be set via environment variable OPENAI_API_KEY.
org_id: The org ID to use. If provided will override the env vars or .env file value.
Can also be set via environment variable OPENAI_ORG_ID.
base_url: The base URL to use. If provided will override the standard value.
Can also be set via environment variable OPENAI_BASE_URL.
default_headers: The default headers mapping of string keys to
string values for HTTP requests.
async_client: An existing client to use.
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
middleware: Optional sequence of middleware to apply to requests.
function_invocation_configuration: Optional configuration for function invocation behavior.
kwargs: Other keyword parameters.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIAssistantsClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_CHAT_MODEL_ID=gpt-4
client = OpenAIAssistantsClient()
# Or passing parameters directly
client = OpenAIAssistantsClient(model_id="gpt-4", api_key="sk-...")
# Or loading from a .env file
client = OpenAIAssistantsClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.openai import OpenAIAssistantsOptions
class MyOptions(OpenAIAssistantsOptions, total=False):
my_custom_option: str
client: OpenAIAssistantsClient[MyOptions] = OpenAIAssistantsClient(model_id="gpt-4")
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
base_url=base_url,
org_id=org_id,
chat_model_id=model_id,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
api_key_value = openai_settings.get("api_key")
if not async_client and not api_key_value:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
chat_model_id = openai_settings.get("chat_model_id")
if not chat_model_id:
raise ValueError(
"OpenAI model ID is required. "
"Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
)
super().__init__(
model_id=chat_model_id,
api_key=self._get_api_key(api_key_value),
org_id=openai_settings.get("org_id"),
default_headers=default_headers,
client=async_client,
base_url=openai_settings.get("base_url"),
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
self.assistant_id: str | None = assistant_id
self.assistant_name: str | None = assistant_name
self.assistant_description: str | None = assistant_description
self.thread_id: str | None = thread_id
self._should_delete_assistant: bool = False
async def __aenter__(self) -> Self:
"""Async context manager entry."""
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: Any,
) -> None:
"""Async context manager exit - clean up any assistants we created."""
await self.close()
async def close(self) -> None:
"""Clean up any assistants we created."""
if self._should_delete_assistant and self.assistant_id is not None:
client = await self._ensure_client()
await client.beta.assistants.delete(self.assistant_id) # type: ignore[reportDeprecated]
object.__setattr__(self, "assistant_id", None)
object.__setattr__(self, "_should_delete_assistant", False)
@override
def _inner_get_response(
self,
*,
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
if stream:
# Streaming mode - return the async generator directly
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
# prepare
run_options, tool_results = self._prepare_options(messages, options, **kwargs)
# Get the thread ID
thread_id: str | None = options.get(
"conversation_id", run_options.get("conversation_id", self.thread_id)
)
if thread_id is None and tool_results is not None:
raise ValueError("No thread ID was provided, but chat messages includes tool results.")
# Determine which assistant to use and create if needed
assistant_id = await self._get_assistant_id_or_create()
# execute
stream_obj, thread_id = await self._create_assistant_stream(
thread_id, assistant_id, run_options, tool_results
)
# process
async for update in self._process_stream_events(stream_obj, thread_id):
yield update
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode - collect updates and convert to response
async def _get_response() -> ChatResponse:
stream_result = self._inner_get_response(messages=messages, options=options, stream=True, **kwargs)
return await ChatResponse.from_update_generator(
updates=stream_result, # type: ignore[arg-type]
output_format_type=options.get("response_format"), # type: ignore[arg-type]
)
return _get_response()
async def _get_assistant_id_or_create(self) -> str:
"""Determine which assistant to use and create if needed.
Returns:
str: The assistant_id to use.
"""
# If no assistant is provided, create a temporary assistant
if self.assistant_id is None:
if not self.model_id:
raise ValueError("Parameter 'model_id' is required for assistant creation.")
client = await self._ensure_client()
created_assistant = await client.beta.assistants.create( # type: ignore[reportDeprecated]
model=self.model_id,
description=self.assistant_description,
name=self.assistant_name,
)
self.assistant_id = created_assistant.id
self._should_delete_assistant = True
return self.assistant_id
async def _create_assistant_stream(
self,
thread_id: str | None,
assistant_id: str,
run_options: dict[str, Any],
tool_results: list[Content] | None,
) -> tuple[Any, str]:
"""Create the assistant stream for processing.
Returns:
tuple: (stream, final_thread_id)
"""
client = await self._ensure_client()
# Get any active run for this thread
thread_run = await self._get_active_thread_run(thread_id)
tool_run_id, tool_outputs = self._prepare_tool_outputs_for_assistants(tool_results)
if thread_run is not None and tool_run_id is not None and tool_run_id == thread_run.id and tool_outputs:
# There's an active run and we have tool results to submit, so submit the results.
stream = client.beta.threads.runs.submit_tool_outputs_stream( # type: ignore[reportDeprecated]
run_id=tool_run_id,
thread_id=thread_run.thread_id,
tool_outputs=tool_outputs,
)
final_thread_id = thread_run.thread_id
else:
# Handle thread creation or cancellation
final_thread_id = await self._prepare_thread(thread_id, thread_run, run_options)
# Now create a new run and stream the results.
stream = client.beta.threads.runs.stream( # type: ignore[reportDeprecated]
assistant_id=assistant_id, thread_id=final_thread_id, **run_options
)
return stream, final_thread_id
async def _get_active_thread_run(self, thread_id: str | None) -> Run | None:
"""Get any active run for the given thread."""
client = await self._ensure_client()
if thread_id is None:
return None
async for run in client.beta.threads.runs.list(thread_id=thread_id, limit=1, order="desc"): # type: ignore[reportDeprecated]
if run.status not in ["completed", "cancelled", "failed", "expired"]:
return run
return None
async def _prepare_thread(self, thread_id: str | None, thread_run: Run | None, run_options: dict[str, Any]) -> str:
"""Prepare the thread for a new run, creating or cleaning up as needed."""
client = await self._ensure_client()
if thread_id is None:
# No thread ID was provided, so create a new thread.
thread = await client.beta.threads.create( # type: ignore[reportDeprecated]
messages=run_options["additional_messages"],
tool_resources=run_options.get("tool_resources"),
metadata=run_options.get("metadata"),
)
run_options["additional_messages"] = []
run_options.pop("tool_resources", None)
return thread.id
if thread_run is not None:
# There was an active run; we need to cancel it before starting a new run.
await client.beta.threads.runs.cancel(run_id=thread_run.id, thread_id=thread_id) # type: ignore[reportDeprecated]
return thread_id
async def _process_stream_events(self, stream: Any, thread_id: str) -> AsyncIterable[ChatResponseUpdate]:
response_id: str | None = None
async with stream as response_stream:
async for response in response_stream:
if response.event == "thread.run.created":
yield ChatResponseUpdate(
contents=[],
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
role="assistant",
)
elif response.event == "thread.run.step.created" and isinstance(response.data, RunStep):
response_id = response.data.run_id
elif response.event == "thread.message.delta" and isinstance(response.data, MessageDeltaEvent):
delta = response.data.delta
role = "user" if delta.role == "user" else "assistant"
for delta_block in delta.content or []:
if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
text_content = Content.from_text(delta_block.text.value)
if delta_block.text.annotations:
annotations: list[Annotation] = []
text_content.annotations = annotations
for annotation in delta_block.text.annotations:
if isinstance(annotation, FileCitationDeltaAnnotation):
ann: Annotation = Annotation(
type="citation",
additional_properties={
"text": annotation.text,
"index": annotation.index,
},
raw_representation=annotation,
)
if annotation.file_citation and annotation.file_citation.file_id:
ann["file_id"] = annotation.file_citation.file_id
if annotation.start_index is not None and annotation.end_index is not None:
ann["annotated_regions"] = [
TextSpanRegion(
type="text_span",
start_index=annotation.start_index,
end_index=annotation.end_index,
)
]
annotations.append(ann)
elif isinstance(annotation, FilePathDeltaAnnotation):
ann = Annotation(
type="citation",
additional_properties={
"text": annotation.text,
"index": annotation.index,
},
raw_representation=annotation,
)
if annotation.file_path and annotation.file_path.file_id:
ann["file_id"] = annotation.file_path.file_id
if annotation.start_index is not None and annotation.end_index is not None:
ann["annotated_regions"] = [
TextSpanRegion(
type="text_span",
start_index=annotation.start_index,
end_index=annotation.end_index,
)
]
annotations.append(ann)
yield ChatResponseUpdate(
role=role, # type: ignore[arg-type]
contents=[text_content],
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
)
elif response.event == "thread.message.completed" and isinstance(response.data, ThreadMessage):
# Process completed message to extract fully resolved annotations.
# Delta events may carry partial/empty annotation data; the completed
# message contains the final text with all citation details populated.
completed_contents: list[Content] = []
for block in response.data.content:
if block.type != "text":
continue
text_content = Content.from_text(block.text.value)
if block.text.annotations:
completed_annotations: list[Annotation] = []
text_content.annotations = completed_annotations
for completed_annotation in block.text.annotations:
if isinstance(completed_annotation, FileCitationAnnotation):
props: dict[str, Any] = {
"text": completed_annotation.text,
}
ann = Annotation(
type="citation",
additional_properties=props,
raw_representation=completed_annotation,
)
if (
completed_annotation.file_citation
and completed_annotation.file_citation.file_id
):
ann["file_id"] = completed_annotation.file_citation.file_id
ann["annotated_regions"] = [
TextSpanRegion(
type="text_span",
start_index=completed_annotation.start_index,
end_index=completed_annotation.end_index,
)
]
text_content.annotations.append(ann)
elif isinstance(completed_annotation, FilePathAnnotation):
ann = Annotation(
type="citation",
additional_properties={
"text": completed_annotation.text,
},
raw_representation=completed_annotation,
)
if completed_annotation.file_path and completed_annotation.file_path.file_id:
ann["file_id"] = completed_annotation.file_path.file_id
ann["annotated_regions"] = [
TextSpanRegion(
type="text_span",
start_index=completed_annotation.start_index,
end_index=completed_annotation.end_index,
)
]
text_content.annotations.append(ann)
else:
logger.debug("Unparsed annotation type: %s", completed_annotation.type)
completed_contents.append(text_content)
if completed_contents:
yield ChatResponseUpdate(
role="assistant",
contents=completed_contents,
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
)
elif response.event == "thread.run.requires_action" and isinstance(response.data, Run):
contents = self._parse_function_calls_from_assistants(response.data, response_id)
if contents:
yield ChatResponseUpdate(
role="assistant",
contents=contents,
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
)
elif (
response.event == "thread.run.completed"
and isinstance(response.data, Run)
and response.data.usage is not None
):
usage = response.data.usage
usage_content = Content.from_usage(
UsageDetails(
input_token_count=usage.prompt_tokens,
output_token_count=usage.completion_tokens,
total_token_count=usage.total_tokens,
)
)
yield ChatResponseUpdate(
role="assistant",
contents=[usage_content],
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
)
else:
yield ChatResponseUpdate(
contents=[],
conversation_id=thread_id,
message_id=response_id,
raw_representation=response.data,
response_id=response_id,
role="assistant",
)
def _parse_function_calls_from_assistants(self, event_data: Run, response_id: str | None) -> list[Content]:
"""Parse function call contents from an assistants tool action event."""
contents: list[Content] = []
if event_data.required_action is not None:
for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
tool_call_any = cast(Any, tool_call)
call_id = json.dumps([response_id, tool_call.id])
tool_type = getattr(tool_call, "type", None)
if tool_type == "code_interpreter" and getattr(tool_call_any, "code_interpreter", None):
code_input = getattr(tool_call_any.code_interpreter, "input", None)
inputs = (
[Content.from_text(text=code_input, raw_representation=tool_call)]
if code_input is not None
else None
)
contents.append(
Content.from_code_interpreter_tool_call(
call_id=call_id,
inputs=inputs,
raw_representation=tool_call,
)
)
elif tool_type == "mcp":
contents.append(
Content.from_mcp_server_tool_call(
call_id=call_id,
tool_name=getattr(tool_call, "name", "") or "",
server_name=getattr(tool_call, "server_label", None),
arguments=getattr(tool_call, "args", None),
raw_representation=tool_call,
)
)
else:
function_name = tool_call.function.name
function_arguments = json.loads(tool_call.function.arguments)
contents.append(
Content.from_function_call(
call_id=call_id,
name=function_name,
arguments=function_arguments,
)
)
return contents
def _prepare_options(
self,
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[dict[str, Any], list[Content] | None]:
from .._types import validate_tool_mode
run_options: dict[str, Any] = {**kwargs}
# Extract options from the dict
max_tokens = options.get("max_tokens")
model_id = options.get("model_id")
top_p = options.get("top_p")
temperature = options.get("temperature")
allow_multiple_tool_calls = options.get("allow_multiple_tool_calls")
tool_choice = options.get("tool_choice")
tools = options.get("tools")
response_format = options.get("response_format")
tool_resources = options.get("tool_resources")
if max_tokens is not None:
run_options["max_completion_tokens"] = max_tokens
if model_id is not None:
run_options["model"] = model_id
if top_p is not None:
run_options["top_p"] = top_p
if temperature is not None:
run_options["temperature"] = temperature
if allow_multiple_tool_calls is not None:
run_options["parallel_tool_calls"] = allow_multiple_tool_calls
if tool_resources is not None:
run_options["tool_resources"] = tool_resources
tool_mode = validate_tool_mode(tool_choice)
tool_definitions: list[MutableMapping[str, Any]] = []
# Always include tools if provided, regardless of tool_choice
# tool_choice="none" means the model won't call tools, but tools should still be available
for tool in normalize_tools(tools):
if isinstance(tool, FunctionTool):
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
elif isinstance(tool, MutableMapping):
# Pass through dict-based tools directly (from static factory methods)
tool_definitions.append(cast(MutableMapping[str, Any], tool))
if len(tool_definitions) > 0:
run_options["tools"] = tool_definitions
if tool_mode is not None:
mode = tool_mode.get("mode")
if mode is None:
raise ValueError("tool_choice mode is required")
if mode == "required" and (func_name := tool_mode.get("required_function_name")) is not None:
run_options["tool_choice"] = {
"type": "function",
"function": {"name": func_name},
}
else:
run_options["tool_choice"] = mode
if response_format is not None:
if isinstance(response_format, dict):
run_options["response_format"] = response_format
else:
run_options["response_format"] = {
"type": "json_schema",
"json_schema": {
"name": response_format.__name__,
"schema": response_format.model_json_schema(),
"strict": True,
},
}
instructions: list[str] = []
tool_results: list[Content] | None = None
additional_messages: list[AdditionalMessage] | None = None
# System/developer messages are turned into instructions,
# since there is no such message roles in OpenAI Assistants.
# All other messages are added 1:1.
for chat_message in messages:
if chat_message.role in ["system", "developer"]:
for text_content in [content for content in chat_message.contents if content.type == "text"]:
text = getattr(text_content, "text", None)
if text:
instructions.append(text)
continue
message_contents: list[MessageContentPartParam] = []
for content in chat_message.contents:
if content.type == "text":
message_contents.append(TextContentBlockParam(type="text", text=content.text)) # type: ignore[attr-defined, typeddict-item]
elif content.type == "uri" and content.has_top_level_media_type("image"):
message_contents.append(
ImageURLContentBlockParam(type="image_url", image_url=ImageURLParam(url=content.uri)) # type: ignore[attr-defined, typeddict-item]
)
elif content.type == "function_result":
if tool_results is None:
tool_results = []
tool_results.append(content)
if len(message_contents) > 0:
if additional_messages is None:
additional_messages = []
additional_messages.append(
AdditionalMessage(
role="assistant" if chat_message.role == "assistant" else "user",
content=message_contents,
)
)
if additional_messages is not None:
run_options["additional_messages"] = additional_messages
if len(instructions) > 0:
run_options["instructions"] = "".join(instructions)
return run_options, tool_results
def _prepare_tool_outputs_for_assistants(
self,
tool_results: list[Content] | None,
) -> tuple[str | None, list[ToolOutput] | None]:
"""Prepare function results for submission to the assistants API."""
run_id: str | None = None
tool_outputs: list[ToolOutput] | None = None
if tool_results:
for function_result_content in tool_results:
# When creating the FunctionCallContent, we created it with a CallId == [runId, callId].
# We need to extract the run ID and ensure that the ToolOutput we send back to Azure
# is only the call ID.
run_and_call_ids: list[str] = json.loads(function_result_content.call_id) # type: ignore[arg-type]
if (
not run_and_call_ids
or len(run_and_call_ids) != 2
or not run_and_call_ids[0]
or not run_and_call_ids[1]
or (run_id is not None and run_id != run_and_call_ids[0])
):
continue
run_id = run_and_call_ids[0]
call_id = run_and_call_ids[1]
if tool_outputs is None:
tool_outputs = []
output = (
function_result_content.result
if function_result_content.result is not None
else "No output received."
)
tool_outputs.append(ToolOutput(tool_call_id=call_id, output=output))
return run_id, tool_outputs
def _update_agent_name_and_description(self, agent_name: str | None, description: str | None = None) -> None:
"""Update the agent name in the chat client.
Args:
agent_name: The new name for the agent.
description: The new description for the agent.
"""
# This is a no-op in the base class, but can be overridden by subclasses
# to update the agent name in the client.
if agent_name and not self.assistant_name:
self.assistant_name = agent_name
if description and not self.assistant_description:
self.assistant_description = description
python/packages/core/agent_framework/openai/_chat_client.py
-975
View File
@@ -1,975 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import logging
import sys
from collections.abc import (
AsyncIterable,
Awaitable,
Callable,
Mapping,
MutableMapping,
Sequence,
)
from datetime import datetime, timezone
from itertools import chain
from typing import Any, Generic, Literal, cast, overload
from openai import AsyncOpenAI, BadRequestError
from openai.lib._parsing._completions import type_to_response_format_param
from openai.types import CompletionUsage
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_message_custom_tool_call import (
ChatCompletionMessageCustomToolCall,
)
from openai.types.chat.completion_create_params import WebSearchOptions
from pydantic import BaseModel
from .._clients import BaseChatClient
from .._docstrings import apply_layered_docstring
from .._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from .._settings import load_settings
from .._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
ToolTypes,
normalize_tools,
)
from .._types import (
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReason,
Message,
ResponseStream,
UsageDetails,
)
from ..exceptions import (
ChatClientException,
ChatClientInvalidRequestException,
)
from ..observability import ChatTelemetryLayer
from ._exceptions import OpenAIContentFilterException
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
logger = logging.getLogger("agent_framework.openai")
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
# region OpenAI Chat Options TypedDict
class PredictionTextContent(TypedDict, total=False):
"""Prediction text content options for OpenAI Chat completions."""
type: Literal["text"]
text: str
class Prediction(TypedDict, total=False):
"""Prediction options for OpenAI Chat completions."""
type: Literal["content"]
content: str | list[PredictionTextContent]
class OpenAIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
"""OpenAI-specific chat options dict.
Extends ChatOptions with options specific to OpenAI's Chat Completions API.
Keys:
model_id: The model to use for the request,
translates to ``model`` in OpenAI API.
temperature: Sampling temperature between 0 and 2.
top_p: Nucleus sampling parameter.
max_tokens: Maximum number of tokens to generate,
translates to ``max_completion_tokens`` in OpenAI API.
stop: Stop sequences.
seed: Random seed for reproducibility.
frequency_penalty: Frequency penalty between -2.0 and 2.0.
presence_penalty: Presence penalty between -2.0 and 2.0.
tools: List of tools (functions) available to the model.
tool_choice: How the model should use tools.
allow_multiple_tool_calls: Whether to allow parallel tool calls,
translates to ``parallel_tool_calls`` in OpenAI API.
response_format: Structured output schema.
metadata: Request metadata for tracking.
user: End-user identifier for abuse monitoring.
store: Whether to store the conversation.
instructions: System instructions for the model (prepended as system message).
# OpenAI-specific options (supported by all models):
logit_bias: Token bias values (-100 to 100).
logprobs: Whether to return log probabilities.
top_logprobs: Number of top log probabilities to return (0-20).
prediction: Whether to use predicted return tokens.
"""
# OpenAI-specific generation parameters (supported by all models)
logit_bias: dict[str | int, float] # type: ignore[misc]
logprobs: bool
top_logprobs: int
prediction: Prediction
OpenAIChatOptionsT = TypeVar("OpenAIChatOptionsT", bound=TypedDict, default="OpenAIChatOptions", covariant=True) # type: ignore[valid-type]
OPTION_TRANSLATIONS: dict[str, str] = {
"model_id": "model",
"allow_multiple_tool_calls": "parallel_tool_calls",
"max_tokens": "max_completion_tokens",
}
# region Base Client
class RawOpenAIChatClient( # type: ignore[misc]
OpenAIBase,
BaseChatClient[OpenAIChatOptionsT],
Generic[OpenAIChatOptionsT],
):
"""Raw OpenAI Chat completion class without middleware, telemetry, or function invocation.
Warning:
**This class should not normally be used directly.** It does not include middleware,
telemetry, or function invocation support that you most likely need. If you do use it,
you should consider which additional layers to apply. There is a defined ordering that
you should follow:
1. **FunctionInvocationLayer** - Owns the tool/function calling loop and routes function middleware
2. **ChatMiddlewareLayer** - Applies chat middleware per model call and stays outside telemetry
3. **ChatTelemetryLayer** - Must stay inside chat middleware for correct per-call telemetry
Use ``OpenAIChatClient`` instead for a fully-featured client with all layers applied.
"""
# region Hosted Tool Factory Methods
@staticmethod
def get_web_search_tool(
*,
web_search_options: WebSearchOptions | None = None,
) -> dict[str, Any]:
"""Create a web search tool configuration for the Chat Completions API.
Note: For the Chat Completions API, web search is passed via the `web_search_options`
parameter rather than in the `tools` array. This method returns a dict that can be
passed as a tool to ChatAgent, which will handle it appropriately.
Keyword Args:
web_search_options: The full WebSearchOptions configuration. This TypedDict includes:
- user_location: Location context with "type" and "approximate" containing
"city", "country", "region", "timezone".
- search_context_size: One of "low", "medium", "high".
Returns:
A dict configuration that enables web search when passed to ChatAgent.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIChatClient
# Basic web search
tool = OpenAIChatClient.get_web_search_tool()
# With location context
tool = OpenAIChatClient.get_web_search_tool(
web_search_options={
"user_location": {
"type": "approximate",
"approximate": {"city": "Seattle", "country": "US"},
},
"search_context_size": "medium",
}
)
agent = ChatAgent(client, tools=[tool])
"""
tool: dict[str, Any] = {"type": "web_search"}
if web_search_options:
tool.update(web_search_options)
return tool
# endregion
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
**kwargs: Any,
) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ...
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[False] = ...,
options: OpenAIChatOptionsT | ChatOptions[None] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]]: ...
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[True],
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
**kwargs: Any,
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
@override
def get_response(
self,
messages: Sequence[Message],
*,
stream: bool = False,
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
"""Get a response from the raw OpenAI chat client."""
super_get_response = cast(
"Callable[..., Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]]",
super().get_response, # type: ignore[misc]
)
return super_get_response( # type: ignore[no-any-return]
messages=messages,
stream=stream,
options=options,
**kwargs,
)
@override
def _inner_get_response(
self,
*,
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
# prepare
options_dict = self._prepare_options(messages, options)
if stream:
# Streaming mode
options_dict["stream_options"] = {"include_usage": True}
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
client = await self._ensure_client()
try:
async for chunk in await client.chat.completions.create(stream=True, **options_dict):
if len(chunk.choices) == 0 and chunk.usage is None:
continue
yield self._parse_response_update_from_openai(chunk)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ChatClientException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ChatClientException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
# Non-streaming mode
async def _get_response() -> ChatResponse:
client = await self._ensure_client()
try:
return self._parse_response_from_openai(
await client.chat.completions.create(stream=False, **options_dict), options
)
except BadRequestError as ex:
if ex.code == "content_filter":
raise OpenAIContentFilterException(
f"{type(self)} service encountered a content error: {ex}",
inner_exception=ex,
) from ex
raise ChatClientException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
except Exception as ex:
raise ChatClientException(
f"{type(self)} service failed to complete the prompt: {ex}",
inner_exception=ex,
) from ex
return _get_response()
# region content creation
def _prepare_tools_for_openai(
self,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
) -> dict[str, Any]:
"""Prepare tools for the OpenAI Chat Completions API.
Converts FunctionTool to JSON schema format. Web search tools are routed
to web_search_options parameter. All other tools pass through unchanged.
Args:
tools: Tool(s) to prepare.
Returns:
Dict containing tools and optionally web_search_options.
"""
chat_tools: list[Any] = []
web_search_options: dict[str, Any] | None = None
for tool in normalize_tools(tools):
if isinstance(tool, FunctionTool):
chat_tools.append(tool.to_json_schema_spec())
elif isinstance(tool, MutableMapping):
typed_tool = cast(MutableMapping[str, Any], tool)
if typed_tool.get("type") == "web_search":
# Web search is handled via web_search_options, not tools array
web_search_options = {k: v for k, v in typed_tool.items() if k != "type"}
else:
# Pass through all other dict-based tools unchanged
chat_tools.append(typed_tool)
else:
# Pass through all other tools (SDK types) unchanged
chat_tools.append(tool)
result: dict[str, Any] = {}
if chat_tools:
result["tools"] = chat_tools
if web_search_options is not None:
result["web_search_options"] = web_search_options
return result
def _prepare_options(self, messages: Sequence[Message], options: Mapping[str, Any]) -> dict[str, Any]:
# Prepend instructions from options if they exist
from .._types import prepend_instructions_to_messages, validate_tool_mode
if instructions := options.get("instructions"):
messages = prepend_instructions_to_messages(list(messages), instructions, role="system")
# Start with a copy of options
run_options = {
k: v for k, v in options.items() if v is not None and k not in {"instructions", "tools", "conversation_id"}
}
# messages
if messages and "messages" not in run_options:
run_options["messages"] = self._prepare_messages_for_openai(messages)
if "messages" not in run_options:
raise ChatClientInvalidRequestException("Messages are required for chat completions")
# Translation between options keys and Chat Completion API
for old_key, new_key in OPTION_TRANSLATIONS.items():
if old_key in run_options and old_key != new_key:
run_options[new_key] = run_options.pop(old_key)
# model id
if not run_options.get("model"):
if not self.model_id:
raise ValueError("model_id must be a non-empty string")
run_options["model"] = self.model_id
# tools
tools = options.get("tools")
if tools is not None:
run_options.update(self._prepare_tools_for_openai(tools))
# Only include tool_choice and parallel_tool_calls if tools are present
if not run_options.get("tools"):
run_options.pop("parallel_tool_calls", None)
run_options.pop("tool_choice", None)
elif tool_choice := run_options.pop("tool_choice", None):
tool_mode = validate_tool_mode(tool_choice)
if tool_mode is not None:
if (mode := tool_mode.get("mode")) == "required" and (
func_name := tool_mode.get("required_function_name")
) is not None:
run_options["tool_choice"] = {
"type": "function",
"function": {"name": func_name},
}
else:
run_options["tool_choice"] = mode
# response format
if response_format := options.get("response_format"):
if isinstance(response_format, dict):
run_options["response_format"] = response_format
else:
run_options["response_format"] = type_to_response_format_param(response_format)
return run_options
def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping[str, Any]) -> ChatResponse:
"""Parse a response from OpenAI into a ChatResponse."""
response_metadata = self._get_metadata_from_chat_response(response)
messages: list[Message] = []
finish_reason: FinishReason | None = None
for choice in response.choices:
response_metadata.update(self._get_metadata_from_chat_choice(choice))
if choice.finish_reason:
finish_reason = choice.finish_reason # type: ignore[assignment]
contents: list[Content] = []
if text_content := self._parse_text_from_openai(choice):
contents.append(text_content)
if parsed_tool_calls := [tool for tool in self._parse_tool_calls_from_openai(choice)]:
contents.extend(parsed_tool_calls)
if reasoning_details := getattr(choice.message, "reasoning_details", None):
contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details)))
messages.append(Message(role="assistant", contents=contents))
return ChatResponse(
response_id=response.id,
created_at=datetime.fromtimestamp(response.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
usage_details=self._parse_usage_from_openai(response.usage) if response.usage else None,
messages=messages,
model_id=response.model,
additional_properties=response_metadata,
finish_reason=finish_reason,
response_format=options.get("response_format"),
)
def _parse_response_update_from_openai(
self,
chunk: ChatCompletionChunk,
) -> ChatResponseUpdate:
"""Parse a streaming response update from OpenAI."""
chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk)
contents: list[Content] = []
finish_reason: FinishReason | None = None
# Process usage data (may coexist with text/tool content in providers like Gemini).
# See https://github.com/microsoft/agent-framework/issues/3434
if chunk.usage:
contents.append(
Content.from_usage(usage_details=self._parse_usage_from_openai(chunk.usage), raw_representation=chunk)
)
for choice in chunk.choices:
chunk_metadata.update(self._get_metadata_from_chat_choice(choice))
contents.extend(self._parse_tool_calls_from_openai(choice))
if choice.finish_reason:
finish_reason = choice.finish_reason # type: ignore[assignment]
if text_content := self._parse_text_from_openai(choice):
contents.append(text_content)
if reasoning_details := getattr(choice.delta, "reasoning_details", None):
contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details)))
return ChatResponseUpdate(
created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
contents=contents,
role="assistant",
model_id=chunk.model,
additional_properties=chunk_metadata,
finish_reason=finish_reason,
raw_representation=chunk,
response_id=chunk.id,
message_id=chunk.id,
)
def _parse_usage_from_openai(self, usage: CompletionUsage) -> UsageDetails:
details = UsageDetails(
input_token_count=usage.prompt_tokens,
output_token_count=usage.completion_tokens,
total_token_count=usage.total_tokens,
)
if usage.completion_tokens_details:
if tokens := usage.completion_tokens_details.accepted_prediction_tokens:
details["completion/accepted_prediction_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if tokens := usage.completion_tokens_details.audio_tokens:
details["completion/audio_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if tokens := usage.completion_tokens_details.reasoning_tokens:
details["completion/reasoning_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if tokens := usage.completion_tokens_details.rejected_prediction_tokens:
details["completion/rejected_prediction_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if usage.prompt_tokens_details:
if tokens := usage.prompt_tokens_details.audio_tokens:
details["prompt/audio_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if tokens := usage.prompt_tokens_details.cached_tokens:
details["prompt/cached_tokens"] = tokens # type: ignore[typeddict-unknown-key]
return details
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
"""Parse the choice into a Content object with type='text'."""
message = choice.message if isinstance(choice, Choice) else choice.delta
if message.content:
return Content.from_text(text=message.content, raw_representation=choice)
if hasattr(message, "refusal") and message.refusal:
return Content.from_text(text=message.refusal, raw_representation=choice)
return None
def _get_metadata_from_chat_response(self, response: ChatCompletion) -> dict[str, Any]:
"""Get metadata from a chat response."""
return {
"system_fingerprint": response.system_fingerprint,
}
def _get_metadata_from_streaming_chat_response(self, response: ChatCompletionChunk) -> dict[str, Any]:
"""Get metadata from a streaming chat response."""
return {
"system_fingerprint": response.system_fingerprint,
}
def _get_metadata_from_chat_choice(self, choice: Choice | ChunkChoice) -> dict[str, Any]:
"""Get metadata from a chat choice."""
return {
"logprobs": getattr(choice, "logprobs", None),
}
def _parse_tool_calls_from_openai(self, choice: Choice | ChunkChoice) -> list[Content]:
"""Parse tool calls from an OpenAI response choice."""
resp: list[Content] = []
content = choice.message if isinstance(choice, Choice) else choice.delta
if content and content.tool_calls:
for tool in content.tool_calls:
if not isinstance(tool, ChatCompletionMessageCustomToolCall) and tool.function:
# ignoring tool.custom
fcc = Content.from_function_call(
call_id=tool.id if tool.id else "",
name=tool.function.name if tool.function.name else "",
arguments=tool.function.arguments if tool.function.arguments else "",
raw_representation=tool.function,
)
resp.append(fcc)
# When you enable asynchronous content filtering in Azure OpenAI, you may receive empty deltas
return resp
def _prepare_messages_for_openai(
self,
chat_messages: Sequence[Message],
role_key: str = "role",
content_key: str = "content",
) -> list[dict[str, Any]]:
"""Prepare the chat history for an OpenAI request.
Allowing customization of the key names for role/author, and optionally overriding the role.
"tool" messages need to be formatted different than system/user/assistant messages:
They require a "tool_call_id" and (function) "name" key, and the "metadata" key should
be removed. The "encoding" key should also be removed.
Override this method to customize the formatting of the chat history for a request.
Args:
chat_messages: The chat history to prepare.
role_key: The key name for the role/author.
content_key: The key name for the content/message.
Returns:
prepared_chat_history (Any): The prepared chat history for a request.
"""
list_of_list = [self._prepare_message_for_openai(message) for message in chat_messages]
# Flatten the list of lists into a single list
return list(chain.from_iterable(list_of_list))
# region Parsers
def _prepare_message_for_openai(self, message: Message) -> list[dict[str, Any]]:
"""Prepare a chat message for OpenAI."""
# System/developer messages must use plain string content because some
# OpenAI-compatible endpoints reject list content for non-user roles.
if message.role in ("system", "developer"):
texts = [content.text for content in message.contents if content.type == "text" and content.text]
if texts:
sys_args: dict[str, Any] = {"role": message.role, "content": "\n".join(texts)}
if message.author_name:
sys_args["name"] = message.author_name
return [sys_args]
return []
all_messages: list[dict[str, Any]] = []
pending_reasoning: Any = None
for content in message.contents:
# Skip approval content - it's internal framework state, not for the LLM
if content.type in ("function_approval_request", "function_approval_response"):
continue
args: dict[str, Any] = {
"role": message.role,
}
if message.author_name and message.role != "tool":
args["name"] = message.author_name
if "reasoning_details" in message.additional_properties and (
details := message.additional_properties["reasoning_details"]
):
args["reasoning_details"] = details
match content.type:
case "function_call":
if all_messages and "tool_calls" in all_messages[-1]:
# If the last message already has tool calls, append to it
all_messages[-1]["tool_calls"].append(self._prepare_content_for_openai(content))
else:
args["tool_calls"] = [self._prepare_content_for_openai(content)] # type: ignore
case "function_result":
args["tool_call_id"] = content.call_id
if content.items:
text_parts = [item.text or "" for item in content.items if item.type == "text"]
rich_items = [item for item in content.items if item.type in ("data", "uri")]
if rich_items:
logger.warning(
"OpenAI Chat Completions API does not support rich content (images, audio) "
"in tool results. Rich content items will be omitted. "
"Use the Responses API client for rich tool results."
)
args["content"] = "\n".join(text_parts) if text_parts else ""
else:
args["content"] = content.result if content.result is not None else ""
all_messages.append(args)
continue
case "text_reasoning" if (protected_data := content.protected_data) is not None:
# Buffer reasoning to attach to the next message with content/tool_calls
pending_reasoning = json.loads(protected_data)
case _:
if "content" not in args:
args["content"] = []
# this is a list to allow multi-modal content
args["content"].append(self._prepare_content_for_openai(content)) # type: ignore
if "content" in args or "tool_calls" in args:
if pending_reasoning is not None:
args["reasoning_details"] = pending_reasoning
pending_reasoning = None
all_messages.append(args)
# If reasoning was the only content, emit a valid message with empty content
if pending_reasoning is not None:
if all_messages:
all_messages[-1]["reasoning_details"] = pending_reasoning
else:
pending_args: dict[str, Any] = {
"role": message.role,
"content": "",
"reasoning_details": pending_reasoning,
}
if message.author_name and message.role != "tool":
pending_args["name"] = message.author_name
all_messages.append(pending_args)
# Flatten text-only content lists to plain strings for broader
# compatibility with OpenAI-like endpoints (e.g. Foundry Local).
# See https://github.com/microsoft/agent-framework/issues/4084
for msg in all_messages:
msg_content: Any = msg.get("content")
if isinstance(msg_content, list):
typed_msg_content = cast(list[object], msg_content)
text_items: list[Mapping[str, Any]] = []
for item in typed_msg_content:
if not isinstance(item, Mapping):
break
text_item = cast(Mapping[str, Any], item)
if text_item.get("type") != "text":
break
text_items.append(text_item)
else:
msg["content"] = "\n".join(
text_item.get("text", "") if isinstance(text_item.get("text", ""), str) else ""
for text_item in text_items
)
return all_messages
def _prepare_content_for_openai(self, content: Content) -> dict[str, Any]:
"""Prepare content for OpenAI."""
match content.type:
case "function_call":
args = json.dumps(content.arguments) if isinstance(content.arguments, Mapping) else content.arguments
return {
"id": content.call_id,
"type": "function",
"function": {"name": content.name, "arguments": args},
}
case "function_result":
return {
"tool_call_id": content.call_id,
"content": content.result if content.result is not None else "",
}
case "data" | "uri" if content.has_top_level_media_type("image"):
image_url_obj: dict[str, Any] = {"url": content.uri}
detail = content.additional_properties.get("detail")
if isinstance(detail, str):
image_url_obj["detail"] = detail
return {
"type": "image_url",
"image_url": image_url_obj,
}
case "data" | "uri" if content.has_top_level_media_type("audio"):
if content.media_type and "wav" in content.media_type:
audio_format = "wav"
elif content.media_type and "mp3" in content.media_type:
audio_format = "mp3"
else:
# Fallback to default to_dict for unsupported audio formats
return content.to_dict(exclude_none=True)
# Extract base64 data from data URI
audio_data = content.uri
if audio_data.startswith("data:"): # type: ignore[union-attr]
# Extract just the base64 part after "data:audio/format;base64,"
audio_data = audio_data.split(",", 1)[-1] # type: ignore[union-attr]
return {
"type": "input_audio",
"input_audio": {
"data": audio_data,
"format": audio_format,
},
}
case "data" | "uri" if content.has_top_level_media_type("application") and content.uri.startswith("data:"): # type: ignore[union-attr]
# All application/* media types should be treated as files for OpenAI
filename = getattr(content, "filename", None) or (
content.additional_properties.get("filename")
if hasattr(content, "additional_properties") and content.additional_properties
else None
)
file_obj = {"file_data": content.uri}
if filename:
file_obj["filename"] = filename
return {
"type": "file",
"file": file_obj,
}
case _:
# Default fallback for all other content types
return content.to_dict(exclude_none=True)
@override
def service_url(self) -> str:
"""Get the URL of the service.
Override this in the subclass to return the proper URL.
If the service does not have a URL, return None.
"""
return str(self.client.base_url) if self.client else "Unknown"
# region Public client
class OpenAIChatClient( # type: ignore[misc]
OpenAIConfigMixin,
FunctionInvocationLayer[OpenAIChatOptionsT],
ChatMiddlewareLayer[OpenAIChatOptionsT],
ChatTelemetryLayer[OpenAIChatOptionsT],
RawOpenAIChatClient[OpenAIChatOptionsT],
Generic[OpenAIChatOptionsT],
):
"""OpenAI Chat completion class with middleware, telemetry, and function invocation support."""
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
function_invocation_kwargs: Mapping[str, Any] | None = None,
client_kwargs: Mapping[str, Any] | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ...
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[False] = ...,
options: OpenAIChatOptionsT | ChatOptions[None] | None = None,
function_invocation_kwargs: Mapping[str, Any] | None = None,
client_kwargs: Mapping[str, Any] | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]]: ...
@overload
def get_response(
self,
messages: Sequence[Message],
*,
stream: Literal[True],
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
function_invocation_kwargs: Mapping[str, Any] | None = None,
client_kwargs: Mapping[str, Any] | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
**kwargs: Any,
) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
@override
def get_response(
self,
messages: Sequence[Message],
*,
stream: bool = False,
options: OpenAIChatOptionsT | ChatOptions[Any] | None = None,
function_invocation_kwargs: Mapping[str, Any] | None = None,
client_kwargs: Mapping[str, Any] | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
"""Get a response from the OpenAI chat client with all standard layers enabled."""
super_get_response = cast(
"Callable[..., Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]]",
super().get_response, # type: ignore[misc]
)
effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}
if middleware is not None:
effective_client_kwargs["middleware"] = middleware
return super_get_response( # type: ignore[no-any-return]
messages=messages,
stream=stream,
options=options,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=effective_client_kwargs,
**kwargs,
)
def __init__(
self,
*,
model_id: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
instruction_role: str | None = None,
base_url: str | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OpenAI Chat completion client.
Keyword Args:
model_id: OpenAI model name, see https://platform.openai.com/docs/models.
Can also be set via environment variable OPENAI_CHAT_MODEL_ID.
api_key: The API key to use. If provided will override the env vars or .env file value.
Can also be set via environment variable OPENAI_API_KEY.
org_id: The org ID to use. If provided will override the env vars or .env file value.
Can also be set via environment variable OPENAI_ORG_ID.
default_headers: The default headers mapping of string keys to
string values for HTTP requests.
async_client: An existing client to use.
instruction_role: The role to use for 'instruction' messages, for example,
"system" or "developer". If not provided, the default is "system".
base_url: The base URL to use. If provided will override
the standard value for an OpenAI connector, the env vars or .env file value.
Can also be set via environment variable OPENAI_BASE_URL.
middleware: Optional sequence of ChatAndFunctionMiddlewareTypes to apply to requests.
function_invocation_configuration: Optional configuration for function invocation support.
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIChatClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_CHAT_MODEL_ID=<model name>
client = OpenAIChatClient()
# Or passing parameters directly
client = OpenAIChatClient(model_id="<model name>", api_key="sk-...")
# Or loading from a .env file
client = OpenAIChatClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.openai import OpenAIChatOptions
class MyOptions(OpenAIChatOptions, total=False):
my_custom_option: str
client: OpenAIChatClient[MyOptions] = OpenAIChatClient(model_id="<model name>")
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
base_url=base_url,
org_id=org_id,
chat_model_id=model_id,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
api_key_value = openai_settings.get("api_key")
if not async_client and not api_key_value:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
chat_model_id = openai_settings.get("chat_model_id")
if not chat_model_id:
raise ValueError(
"OpenAI model ID is required. "
"Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
)
base_url_value = openai_settings.get("base_url")
super().__init__(
model_id=chat_model_id,
api_key=self._get_api_key(api_key_value),
base_url=base_url_value if base_url_value else None,
org_id=openai_settings.get("org_id"),
default_headers=default_headers,
client=async_client,
instruction_role=instruction_role,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
def _apply_openai_chat_client_docstrings() -> None:
"""Align OpenAI chat-client docstrings with the raw implementation."""
apply_layered_docstring(RawOpenAIChatClient.get_response, BaseChatClient.get_response)
apply_layered_docstring(
OpenAIChatClient.get_response,
RawOpenAIChatClient.get_response,
extra_keyword_args={
"middleware": """
Optional per-call chat and function middleware.
This is merged with any middleware configured on the client for the current request.
""",
},
)
_apply_openai_chat_client_docstrings()
python/packages/core/agent_framework/openai/_embedding_client.py
-219
View File
@@ -1,219 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import base64
import struct
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Any, Generic, Literal, TypedDict
from openai import AsyncOpenAI
from .._clients import BaseEmbeddingClient
from .._settings import load_settings
from .._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails
from ..observability import EmbeddingTelemetryLayer
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
class OpenAIEmbeddingOptions(EmbeddingGenerationOptions, total=False):
"""OpenAI-specific embedding options.
Extends EmbeddingGenerationOptions with OpenAI-specific fields.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIEmbeddingOptions
options: OpenAIEmbeddingOptions = {
"model_id": "text-embedding-3-small",
"dimensions": 1536,
"encoding_format": "float",
}
"""
encoding_format: Literal["float", "base64"]
user: str
OpenAIEmbeddingOptionsT = TypeVar(
"OpenAIEmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIEmbeddingOptions",
covariant=True,
)
class RawOpenAIEmbeddingClient(
OpenAIBase,
BaseEmbeddingClient[str, list[float], OpenAIEmbeddingOptionsT],
Generic[OpenAIEmbeddingOptionsT],
):
"""Raw OpenAI embedding client without telemetry."""
def service_url(self) -> str:
"""Get the URL of the service."""
return str(self.client.base_url) if self.client else "Unknown"
async def get_embeddings(
self,
values: Sequence[str],
*,
options: OpenAIEmbeddingOptionsT | None = None,
) -> GeneratedEmbeddings[list[float], OpenAIEmbeddingOptionsT]:
"""Call the OpenAI embeddings API.
Args:
values: The text values to generate embeddings for.
options: Optional embedding generation options.
Returns:
Generated embeddings with usage metadata.
Raises:
ValueError: If model_id is not provided or values is empty.
"""
if not values:
return GeneratedEmbeddings([], options=options) # type: ignore
opts: dict[str, Any] = options or {} # type: ignore
model = opts.get("model_id") or self.model_id
if not model:
raise ValueError("model_id is required")
kwargs: dict[str, Any] = {"input": list(values), "model": model}
if dimensions := opts.get("dimensions"):
kwargs["dimensions"] = dimensions
if encoding_format := opts.get("encoding_format"):
kwargs["encoding_format"] = encoding_format
if user := opts.get("user"):
kwargs["user"] = user
response = await (await self._ensure_client()).embeddings.create(**kwargs)
encoding = kwargs.get("encoding_format", "float")
embeddings: list[Embedding[list[float]]] = []
for item in response.data:
vector: list[float]
if encoding == "base64" and isinstance(item.embedding, str):
# Decode base64-encoded floats (little-endian IEEE 754)
raw = base64.b64decode(item.embedding)
vector = list(struct.unpack(f"<{len(raw) // 4}f", raw))
else:
vector = item.embedding # type: ignore[assignment]
embeddings.append(
Embedding(
vector=vector,
dimensions=len(vector),
model_id=response.model,
)
)
usage_dict: UsageDetails | None = None
if response.usage:
usage_dict = {
"input_token_count": response.usage.prompt_tokens,
"total_token_count": response.usage.total_tokens,
}
return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
class OpenAIEmbeddingClient(
OpenAIConfigMixin,
EmbeddingTelemetryLayer[str, list[float], OpenAIEmbeddingOptionsT],
RawOpenAIEmbeddingClient[OpenAIEmbeddingOptionsT],
Generic[OpenAIEmbeddingOptionsT],
):
"""OpenAI embedding client with telemetry support.
Keyword Args:
model_id: The embedding model ID (e.g. "text-embedding-3-small").
Can also be set via environment variable OPENAI_EMBEDDING_MODEL_ID.
api_key: OpenAI API key.
Can also be set via environment variable OPENAI_API_KEY.
org_id: OpenAI organization ID.
default_headers: Additional HTTP headers.
async_client: Pre-configured AsyncOpenAI client.
base_url: Custom API base URL.
otel_provider_name: Override the OpenTelemetry provider name for telemetry.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIEmbeddingClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_EMBEDDING_MODEL_ID=text-embedding-3-small
client = OpenAIEmbeddingClient()
# Or passing parameters directly
client = OpenAIEmbeddingClient(
model_id="text-embedding-3-small",
api_key="sk-...",
)
# Generate embeddings
result = await client.get_embeddings(["Hello, world!"])
print(result[0].vector)
"""
def __init__(
self,
*,
model_id: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
base_url: str | None = None,
otel_provider_name: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OpenAI embedding client."""
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
base_url=base_url,
org_id=org_id,
embedding_model_id=model_id,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
api_key_value = openai_settings.get("api_key")
if not async_client and not api_key_value:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
embedding_model_id = openai_settings.get("embedding_model_id")
if not embedding_model_id:
raise ValueError(
"OpenAI embedding model ID is required. "
"Set via 'model_id' parameter or 'OPENAI_EMBEDDING_MODEL_ID' environment variable."
)
base_url_value = openai_settings.get("base_url")
super().__init__(
model_id=embedding_model_id,
api_key=self._get_api_key(api_key_value),
base_url=base_url_value if base_url_value else None,
org_id=openai_settings.get("org_id"),
default_headers=default_headers,
client=async_client,
otel_provider_name=otel_provider_name,
)
python/packages/core/agent_framework/openai/_exceptions.py
-91
View File
@@ -1,91 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from typing import Any
from openai import BadRequestError
from ..exceptions import ChatClientContentFilterException
class ContentFilterResultSeverity(Enum):
"""The severity of the content filter result."""
HIGH = "high"
MEDIUM = "medium"
SAFE = "safe"
LOW = "low"
@dataclass
class ContentFilterResult:
"""The result of a content filter check."""
filtered: bool = False
detected: bool = False
severity: ContentFilterResultSeverity = ContentFilterResultSeverity.SAFE
@classmethod
def from_inner_error_result(cls, inner_error_results: dict[str, Any]) -> ContentFilterResult:
"""Creates a ContentFilterResult from the inner error results.
Args:
inner_error_results: The inner error results.
Returns:
ContentFilterResult: The ContentFilterResult.
"""
return cls(
filtered=inner_error_results.get("filtered", False),
detected=inner_error_results.get("detected", False),
severity=ContentFilterResultSeverity(
inner_error_results.get("severity", ContentFilterResultSeverity.SAFE.value)
),
)
class ContentFilterCodes(Enum):
"""Content filter codes."""
RESPONSIBLE_AI_POLICY_VIOLATION = "ResponsibleAIPolicyViolation"
@dataclass
class OpenAIContentFilterException(ChatClientContentFilterException):
"""AI exception for an error from Azure OpenAI's content filter."""
# The parameter that caused the error.
param: str | None
# The error code specific to the content filter.
content_filter_code: ContentFilterCodes
# The results of the different content filter checks.
content_filter_result: dict[str, ContentFilterResult]
def __init__(
self,
message: str,
inner_exception: BadRequestError,
) -> None:
"""Initializes a new instance of the ContentFilterAIException class.
Args:
message: The error message.
inner_exception: The inner exception.
"""
super().__init__(message)
self.param = inner_exception.param
if inner_exception.body is not None and isinstance(inner_exception.body, dict):
inner_error = inner_exception.body.get("innererror", {}) # type: ignore
self.content_filter_code = ContentFilterCodes(
inner_error.get("code", ContentFilterCodes.RESPONSIBLE_AI_POLICY_VIOLATION.value) # type: ignore
)
self.content_filter_result = {
key: ContentFilterResult.from_inner_error_result(values) # type: ignore
for key, values in inner_error.get("content_filter_result", {}).items() # type: ignore
}
python/packages/core/agent_framework/openai/_responses_client.py
-2373
View File
File diff suppressed because it is too large Load Diff
python/packages/core/agent_framework/openai/_shared.py
-348
View File
@@ -1,348 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
import sys
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from copy import copy
from typing import Any, ClassVar, Union, cast
import openai
from openai import (
AsyncOpenAI,
AsyncStream,
_legacy_response, # type: ignore
)
from openai.types import Completion
from openai.types.audio import Transcription
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from openai.types.images_response import ImagesResponse
from openai.types.responses.response import Response
from openai.types.responses.response_stream_event import ResponseStreamEvent
from packaging.version import parse
from .._serialization import SerializationMixin
from .._settings import SecretString
from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from .._tools import FunctionTool
logger: logging.Logger = logging.getLogger("agent_framework.openai")
RESPONSE_TYPE = Union[
ChatCompletion,
Completion,
AsyncStream[ChatCompletionChunk],
AsyncStream[Completion],
list[Any],
ImagesResponse,
Response,
AsyncStream[ResponseStreamEvent],
Transcription,
_legacy_response.HttpxBinaryResponseContent,
]
OPTION_TYPE = dict[str, Any]
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
def _check_openai_version_for_callable_api_key() -> None:
"""Check if OpenAI version supports callable API keys.
Callable API keys require OpenAI >= 1.106.0.
If the version is too old, raise a ValueError with helpful message.
"""
try:
current_version = parse(openai.__version__)
min_required_version = parse("1.106.0")
if current_version < min_required_version:
raise ValueError(
f"Callable API keys require OpenAI SDK >= 1.106.0, but you have {openai.__version__}. "
f"Please upgrade with 'pip install openai>=1.106.0' or provide a string API key instead. "
f"Note: If you're using mem0ai, you may need to upgrade to mem0ai>=1.0.0 "
f"to allow newer OpenAI versions."
)
except ValueError:
raise # Re-raise our own exception
except Exception as e:
logger.warning(f"Could not check OpenAI version for callable API key support: {e}")
class OpenAISettings(TypedDict, total=False):
"""OpenAI environment settings.
Settings are resolved in this order: explicit keyword arguments, values from an
explicitly provided .env file, then environment variables with the prefix
'OPENAI_'. If settings are missing after resolution, validation will fail.
Keyword Args:
api_key: OpenAI API key, see https://platform.openai.com/account/api-keys.
Can be set via environment variable OPENAI_API_KEY.
base_url: The base URL for the OpenAI API.
Can be set via environment variable OPENAI_BASE_URL.
org_id: This is usually optional unless your account belongs to multiple organizations.
Can be set via environment variable OPENAI_ORG_ID.
chat_model_id: The OpenAI chat model ID to use, for example, gpt-3.5-turbo or gpt-4.
Can be set via environment variable OPENAI_CHAT_MODEL_ID.
responses_model_id: The OpenAI responses model ID to use, for example, gpt-4o or o1.
Can be set via environment variable OPENAI_RESPONSES_MODEL_ID.
embedding_model_id: The OpenAI embedding model ID to use, for example, text-embedding-3-small.
Can be set via environment variable OPENAI_EMBEDDING_MODEL_ID.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAISettings
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_CHAT_MODEL_ID=gpt-4
settings = load_settings(OpenAISettings, env_prefix="OPENAI_")
# Or passing parameters directly
settings = load_settings(OpenAISettings, env_prefix="OPENAI_", api_key="sk-...", chat_model_id="gpt-4")
# Or loading from a .env file
settings = load_settings(OpenAISettings, env_prefix="OPENAI_", env_file_path="path/to/.env")
"""
api_key: SecretString | Callable[[], str | Awaitable[str]] | None
base_url: str | None
org_id: str | None
chat_model_id: str | None
responses_model_id: str | None
embedding_model_id: str | None
class OpenAIBase(SerializationMixin):
"""Base class for OpenAI Clients."""
INJECTABLE: ClassVar[set[str]] = {"client"}
def __init__(self, *, model_id: str | None = None, client: AsyncOpenAI | None = None, **kwargs: Any) -> None:
"""Initialize OpenAIBase.
Keyword Args:
client: The AsyncOpenAI client instance.
model_id: The AI model ID to use.
**kwargs: Additional keyword arguments.
"""
self.client = client
self.model_id = None
if model_id:
self.model_id = model_id.strip()
# Call super().__init__() to continue MRO chain (e.g., RawChatClient)
# Extract known kwargs that belong to other base classes
additional_properties = kwargs.pop("additional_properties", None)
middleware = kwargs.pop("middleware", None)
instruction_role = kwargs.pop("instruction_role", None)
function_invocation_configuration = kwargs.pop("function_invocation_configuration", None)
# Build super().__init__() args
super_kwargs = {}
if additional_properties is not None:
super_kwargs["additional_properties"] = additional_properties
if middleware is not None:
super_kwargs["middleware"] = middleware
if function_invocation_configuration is not None:
super_kwargs["function_invocation_configuration"] = function_invocation_configuration
# Call super().__init__() with filtered kwargs
super().__init__(**super_kwargs)
# Store instruction_role and any remaining kwargs as instance attributes
if instruction_role is not None:
self.instruction_role = instruction_role
for key, value in kwargs.items():
setattr(self, key, value)
async def _initialize_client(self) -> None:
"""Initialize OpenAI client asynchronously.
Override in subclasses to initialize the OpenAI client asynchronously.
"""
pass
async def _ensure_client(self) -> AsyncOpenAI:
"""Ensure OpenAI client is initialized."""
await self._initialize_client()
if self.client is None:
raise RuntimeError("OpenAI client is not initialized")
return self.client
def _get_api_key(
self, api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None
) -> str | Callable[[], str | Awaitable[str]] | None:
"""Get the appropriate API key value for client initialization.
Args:
api_key: The API key parameter which can be a string, SecretString, callable, or None.
Returns:
For callable API keys: returns the callable directly.
For SecretString/string/None API keys: returns as-is (SecretString is a str subclass).
"""
if isinstance(api_key, SecretString):
return api_key.get_secret_value()
# Check version compatibility for callable API keys
if callable(api_key):
_check_openai_version_for_callable_api_key()
return api_key # Pass callable, string, or None directly to OpenAI SDK
class OpenAIConfigMixin(OpenAIBase):
"""Internal class for configuring a connection to an OpenAI service."""
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
def __init__(
self,
model_id: str,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
client: AsyncOpenAI | None = None,
instruction_role: str | None = None,
base_url: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a client for OpenAI services.
This constructor sets up a client to interact with OpenAI's API, allowing for
different types of AI model interactions, like chat or text completion.
Args:
model_id: OpenAI model identifier. Must be non-empty.
Default to a preset value.
api_key: OpenAI API key for authentication, or a callable that returns an API key.
Must be non-empty. (Optional)
org_id: OpenAI organization ID. This is optional
unless the account belongs to multiple organizations.
default_headers: Default headers
for HTTP requests. (Optional)
client: An existing OpenAI client, optional.
instruction_role: The role to use for 'instruction'
messages, for example, summarization prompts could use `developer` or `system`. (Optional)
base_url: The optional base URL to use. If provided will override the standard value for a OpenAI connector.
Will not be used when supplying a custom client.
kwargs: Additional keyword arguments.
"""
# Merge APP_INFO into the headers if it exists
merged_headers = dict(copy(default_headers)) if default_headers else {}
if APP_INFO:
merged_headers.update(APP_INFO)
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
# Handle callable API key using base class method
api_key_value = self._get_api_key(api_key)
if not client:
if not api_key:
raise ValueError("Please provide an api_key")
args: dict[str, Any] = {"api_key": api_key_value, "default_headers": merged_headers}
if org_id:
args["organization"] = org_id
if base_url:
args["base_url"] = base_url
client = AsyncOpenAI(**args)
# Store configuration as instance attributes for serialization
self.org_id = org_id
self.base_url = str(base_url)
# Store default_headers but filter out USER_AGENT_KEY for serialization
if default_headers:
self.default_headers: dict[str, Any] | None = {
k: v for k, v in default_headers.items() if k != USER_AGENT_KEY
}
else:
self.default_headers = None
args = {
"model_id": model_id,
"client": client,
}
if instruction_role:
args["instruction_role"] = instruction_role
# Ensure additional_properties and middleware are passed through kwargs to RawChatClient
# These are consumed by RawChatClient.__init__ via kwargs
super().__init__(**args, **kwargs)
def to_assistant_tools(
tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Convert Agent Framework tools to OpenAI Assistants API format.
Handles FunctionTool instances and dict-based tools from static factory methods.
Args:
tools: Sequence of Agent Framework tools.
Returns:
List of tool definitions for OpenAI Assistants API.
"""
if not tools:
return []
tool_definitions: list[dict[str, Any]] = []
for tool in tools:
if isinstance(tool, FunctionTool):
tool_definitions.append(tool.to_json_schema_spec())
elif isinstance(tool, MutableMapping):
# Pass through dict-based tools directly (from static factory methods)
tool_definitions.append(dict(tool))
return tool_definitions
def from_assistant_tools(
assistant_tools: list[Any] | None,
) -> list[dict[str, Any]]:
"""Convert OpenAI Assistant tools to dict-based format.
This converts hosted tools (code_interpreter, file_search) from an OpenAI
Assistant definition back to dict-based tool definitions.
Note: Function tools are skipped - user must provide implementations separately.
Args:
assistant_tools: Tools from OpenAI Assistant object (assistant.tools).
Returns:
List of dict-based tool definitions for hosted tools.
"""
if not assistant_tools:
return []
tools: list[dict[str, Any]] = []
for tool in assistant_tools:
if hasattr(tool, "type"):
tool_type = tool.type
elif isinstance(tool, Mapping):
typed_tool = cast(Mapping[str, Any], tool)
tool_type_value: Any = typed_tool.get("type")
tool_type = tool_type_value if isinstance(tool_type_value, str) else None
else:
tool_type = None
if tool_type == "code_interpreter":
tools.append({"type": "code_interpreter"})
elif tool_type == "file_search":
tools.append({"type": "file_search"})
# Skip function tools - user must provide implementations
return tools
python/packages/core/pyproject.toml
+1
View File
@@ -54,6 +54,7 @@ all = [
"agent-framework-declarative",
"agent-framework-devui",
"agent-framework-durabletask",
"agent-framework-foundry",
"agent-framework-foundry-local",
"agent-framework-github-copilot; python_version >= '3.11'",
"agent-framework-lab",
python/packages/core/tests/azure/conftest.py
-62
View File
@@ -1,62 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from pytest import fixture
from agent_framework import Message
# region: Connector Settings fixtures
@fixture
def exclude_list(request: Any) -> list[str]:
"""Fixture that returns a list of environment variables to exclude."""
return request.param if hasattr(request, "param") else []
@fixture
def override_env_param_dict(request: Any) -> dict[str, str]:
"""Fixture that returns a dict of environment variables to override."""
return request.param if hasattr(request, "param") else {}
# These two fixtures are used for multiple things, also non-connector tests
@fixture()
def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
"""Fixture to set environment variables for AzureOpenAISettings."""
if exclude_list is None:
exclude_list = []
if override_env_param_dict is None:
override_env_param_dict = {}
env_vars = {
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.com",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_TEXT_DEPLOYMENT_NAME": "test_text_deployment",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
"AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME": "test_text_to_image_deployment",
"AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME": "test_audio_to_text_deployment",
"AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME": "test_text_to_audio_deployment",
"AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME": "test_realtime_deployment",
"AZURE_OPENAI_API_KEY": "test_api_key",
"AZURE_OPENAI_API_VERSION": "2023-03-15-preview",
"AZURE_OPENAI_BASE_URL": "https://test_text_deployment.test-base-url.com",
"AZURE_OPENAI_TOKEN_ENDPOINT": "https://test-token-endpoint.com",
}
env_vars.update(override_env_param_dict) # type: ignore
for key, value in env_vars.items():
if key in exclude_list:
monkeypatch.delenv(key, raising=False) # type: ignore
continue
monkeypatch.setenv(key, value) # type: ignore
return env_vars
@fixture(scope="function")
def chat_history() -> list[Message]:
return []
python/packages/core/tests/azure/test_azure_assistants_client.py
-725
View File
@@ -1,725 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from azure.identity import AzureCliCredential
from pydantic import Field
from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentSession,
ChatResponse,
ChatResponseUpdate,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework._settings import SecretString
from agent_framework.azure import AzureOpenAIAssistantsClient
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests.",
)
def create_test_azure_assistants_client(
mock_async_azure_openai: MagicMock,
deployment_name: str | None = None,
assistant_id: str | None = None,
assistant_name: str | None = None,
thread_id: str | None = None,
should_delete_assistant: bool = False,
) -> AzureOpenAIAssistantsClient:
"""Helper function to create AzureOpenAIAssistantsClient instances for testing."""
client = AzureOpenAIAssistantsClient(
deployment_name=deployment_name or "test_chat_deployment",
assistant_id=assistant_id,
assistant_name=assistant_name,
thread_id=thread_id,
api_key="test-api-key",
endpoint="https://test-endpoint.com",
async_client=mock_async_azure_openai,
)
# Set the _should_delete_assistant flag directly if needed
if should_delete_assistant:
object.__setattr__(client, "_should_delete_assistant", True)
return client
@pytest.fixture
def mock_async_azure_openai() -> MagicMock:
"""Mock AsyncAzureOpenAI client."""
mock_client = MagicMock()
# Mock beta.assistants
mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
mock_client.beta.assistants.delete = AsyncMock()
# Mock beta.threads
mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
mock_client.beta.threads.delete = AsyncMock()
# Mock beta.threads.runs
mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
mock_client.beta.threads.runs.retrieve = AsyncMock()
mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
# Mock beta.threads.messages
mock_client.beta.threads.messages.create = AsyncMock()
mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
return mock_client
def test_azure_assistants_client_init_with_client(mock_async_azure_openai: MagicMock) -> None:
"""Test AzureOpenAIAssistantsClient initialization with existing client."""
client = create_test_azure_assistants_client(
mock_async_azure_openai,
deployment_name="test_chat_deployment",
assistant_id="existing-assistant-id",
thread_id="test-thread-id",
)
assert client.client is mock_async_azure_openai
assert client.model_id == "test_chat_deployment"
assert client.assistant_id == "existing-assistant-id"
assert client.thread_id == "test-thread-id"
assert not client._should_delete_assistant # type: ignore
assert isinstance(client, SupportsChatGetResponse)
def test_azure_assistants_client_init_auto_create_client(
azure_openai_unit_test_env: dict[str, str],
mock_async_azure_openai: MagicMock,
) -> None:
"""Test AzureOpenAIAssistantsClient initialization with auto-created client."""
client = AzureOpenAIAssistantsClient(
deployment_name=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
assistant_name="TestAssistant",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
async_client=mock_async_azure_openai,
)
assert client.client is mock_async_azure_openai
assert client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert client.assistant_id is None
assert client.assistant_name == "TestAssistant"
assert not client._should_delete_assistant # type: ignore
def test_azure_assistants_client_init_validation_fail() -> None:
"""Test AzureOpenAIAssistantsClient initialization with validation failure."""
with pytest.raises(ValueError):
# Force failure by providing invalid deployment name type - this should cause validation to fail
AzureOpenAIAssistantsClient(deployment_name=123, api_key="valid-key") # type: ignore
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
def test_azure_assistants_client_init_missing_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with missing deployment name."""
with pytest.raises(ValueError):
AzureOpenAIAssistantsClient(api_key=azure_openai_unit_test_env.get("AZURE_OPENAI_API_KEY", "test-key"))
def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test AzureOpenAIAssistantsClient initialization with default headers."""
default_headers = {"X-Unit-Test": "test-guid"}
client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
assert client.model_id == "test_chat_deployment"
assert isinstance(client, SupportsChatGetResponse)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in client.client.default_headers
assert client.client.default_headers[key] == value
async def test_azure_assistants_client_get_assistant_id_or_create_existing_assistant(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "existing-assistant-id"
assert not client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_not_called()
async def test_azure_assistants_client_get_assistant_id_or_create_create_new(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when creating a new assistant."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
)
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "test-assistant-id"
assert client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_called_once()
async def test_azure_assistants_client_aclose_should_not_delete(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test close when assistant should not be deleted."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
)
await client.close() # type: ignore
# Verify assistant deletion was not called
mock_async_azure_openai.beta.assistants.delete.assert_not_called()
assert not client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_aclose_should_delete(mock_async_azure_openai: MagicMock) -> None:
"""Test close method calls cleanup."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
await client.close()
# Verify assistant deletion was called
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
assert not client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_async_context_manager(mock_async_azure_openai: MagicMock) -> None:
"""Test async context manager functionality."""
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
# Test context manager
async with client:
pass # Just test that we can enter and exit
# Verify cleanup was called on exit
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test serialization of AzureOpenAIAssistantsClient."""
default_headers = {"X-Unit-Test": "test-guid"}
# Test basic initialization and to_dict
client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
assistant_id="test-assistant-id",
assistant_name="TestAssistant",
thread_id="test-thread-id",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
dumped_settings = client.to_dict()
assert dumped_settings["model_id"] == "test_chat_deployment"
assert dumped_settings["assistant_id"] == "test-assistant-id"
assert dumped_settings["assistant_name"] == "TestAssistant"
assert dumped_settings["thread_id"] == "test-thread-id"
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
return f"The weather in {location} is sunny with a high of 25°C."
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response() -> None:
"""Test Azure Assistants Client response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(
Message(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(Message(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_get_response_tools() -> None:
"""Test Azure Assistants Client response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(
messages=messages,
options={"tools": [get_weather], "tool_choice": "auto"},
)
assert response is not None
assert isinstance(response, ChatResponse)
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming() -> None:
"""Test Azure Assistants Client streaming response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(
Message(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(Message(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_response(messages=messages, stream=True)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if content.type == "text" and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_streaming_tools() -> None:
"""Test Azure Assistants Client streaming response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_response(
messages=messages,
options={"tools": [get_weather], "tool_choice": "auto"},
stream=True,
)
full_message: str = ""
async for chunk in response:
assert chunk is not None
assert isinstance(chunk, ChatResponseUpdate)
for content in chunk.contents:
if content.type == "text" and content.text:
full_message += content.text
assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_with_existing_assistant() -> None:
"""Test Azure Assistants Client with existing assistant ID."""
# First create an assistant to use in the test
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [Message(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
# Now test using the existing assistant
async with AzureOpenAIAssistantsClient(
assistant_id=assistant_id, credential=AzureCliCredential()
) as azure_assistants_client:
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
assert azure_assistants_client.assistant_id == assistant_id
messages = [Message(role="user", text="What can you do?")]
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
assert response is not None
assert isinstance(response, ChatResponse)
assert len(response.text) > 0
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run():
"""Test Agent basic run functionality with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
assert "Hello World" in response.text
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run_streaming():
"""Test Agent basic streaming functionality with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run streaming query
full_message: str = ""
async for chunk in agent.run("Please respond with exactly: 'This is a streaming response test.'", stream=True):
assert chunk is not None
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
full_message += chunk.text
# Validate streaming response
assert len(full_message) > 0
assert "streaming response test" in full_message.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_session_persistence():
"""Test Agent session persistence across runs with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new session that will be reused
session = agent.create_session()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", session=session
)
assert isinstance(first_response, AgentResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", session=session
)
assert isinstance(second_response, AgentResponse)
assert "42" in second_response.text
# Verify session has been populated with conversation ID
assert session.service_session_id is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_existing_session_id():
"""Test Agent with existing session ID to continue conversations across agent instances."""
# First, create a conversation and capture the session ID
existing_session_id = None
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Start a conversation and get the session ID
session = agent.create_session()
response1 = await agent.run("What's the weather in Paris?", session=session)
# Validate first response
assert isinstance(response1, AgentResponse)
assert response1.text is not None
assert any(word in response1.text.lower() for word in ["weather", "paris"])
# The session ID is set after the first response
existing_session_id = session.service_session_id
assert existing_session_id is not None
# Now continue with the same session ID in a new agent instance
async with Agent(
client=AzureOpenAIAssistantsClient(thread_id=existing_session_id, credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_session_id)
# Ask about the previous conversation
response2 = await agent.run("What was the last city I asked about?", session=session)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentResponse)
assert response2.text is not None
# Should reference Paris from the previous conversation
assert "paris" in response2.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_code_interpreter():
"""Test Agent with code interpreter through AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[AzureOpenAIAssistantsClient.get_code_interpreter_tool()],
) as agent:
# Request code execution
response = await agent.run("Write Python code to calculate the factorial of 5 and show the result.")
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
# Factorial of 5 is 120
assert "120" in response.text or "factorial" in response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Assistants Client."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
def test_azure_assistants_client_entra_id_authentication() -> None:
"""Test credential authentication path with sync credential."""
mock_credential = MagicMock()
mock_provider = MagicMock(return_value="token-string")
with (
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
patch(
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
return_value=mock_provider,
) as mock_resolve,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": "https://cognitiveservices.azure.com/.default",
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_credential,
token_endpoint="https://cognitiveservices.azure.com/.default",
)
# Verify credential was resolved to a token provider
mock_resolve.assert_called_once_with(mock_credential, "https://cognitiveservices.azure.com/.default")
# Verify client was created with the token provider
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_provider
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_no_authentication_error() -> None:
"""Test authentication validation error when no auth provided."""
with patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings:
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
# Test missing authentication raises error
with pytest.raises(ValueError, match="api_key, credential, or a client"):
AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
# No authentication provided at all
)
def test_azure_assistants_client_callable_credential() -> None:
"""Test callable token provider as credential."""
mock_provider = MagicMock(return_value="my-token")
with (
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
patch(
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
return_value=mock_provider,
),
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": None,
"token_endpoint": "https://cognitiveservices.azure.com/.default",
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
endpoint="https://test-endpoint.openai.azure.com",
credential=mock_provider,
token_endpoint="https://cognitiveservices.azure.com/.default",
)
# Verify client was created with the token provider
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_ad_token_provider"] is mock_provider
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_base_url_configuration() -> None:
"""Test base_url client parameter path."""
with (
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": SecretString("test-api-key"),
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": None,
"base_url": "https://custom-base-url.com",
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment", api_key="test-api-key", base_url="https://custom-base-url.com"
)
# base_url path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["base_url"] == "https://custom-base-url.com"
assert "azure_endpoint" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
def test_azure_assistants_client_azure_endpoint_configuration() -> None:
"""Test azure_endpoint client parameter path."""
with (
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
):
mock_load_settings.return_value = {
"chat_deployment_name": "test-deployment",
"responses_deployment_name": None,
"api_key": SecretString("test-api-key"),
"token_endpoint": None,
"api_version": "2024-05-01-preview",
"endpoint": "https://test-endpoint.openai.azure.com",
"base_url": None,
}
client = AzureOpenAIAssistantsClient(
deployment_name="test-deployment",
api_key="test-api-key",
endpoint="https://test-endpoint.openai.azure.com",
)
# azure_endpoint path
mock_azure_client.assert_called_once()
call_args = mock_azure_client.call_args[1]
assert call_args["azure_endpoint"] == "https://test-endpoint.openai.azure.com"
assert "base_url" not in call_args
assert client is not None
assert isinstance(client, AzureOpenAIAssistantsClient)
python/packages/core/tests/azure/test_azure_chat_client.py
-1064
View File
File diff suppressed because it is too large Load Diff
python/packages/core/tests/azure/test_azure_embedding_client.py
-159
View File
@@ -1,159 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
from agent_framework.azure import AzureOpenAIEmbeddingClient
from agent_framework.openai import OpenAIEmbeddingOptions
def _make_openai_response(
embeddings: list[list[float]],
model: str = "text-embedding-3-small",
prompt_tokens: int = 5,
total_tokens: int = 5,
) -> CreateEmbeddingResponse:
"""Helper to create a mock OpenAI embeddings response."""
data = [OpenAIEmbedding(embedding=emb, index=i, object="embedding") for i, emb in enumerate(embeddings)]
return CreateEmbeddingResponse(
data=data,
model=model,
object="list",
usage=Usage(prompt_tokens=prompt_tokens, total_tokens=total_tokens),
)
@pytest.fixture
def azure_embedding_unit_test_env(monkeypatch: pytest.MonkeyPatch) -> None:
"""Clear ambient Azure OpenAI embedding env vars for deterministic unit tests."""
for key in (
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_TOKEN_ENDPOINT",
):
monkeypatch.delenv(key, raising=False)
def test_azure_construction_with_deployment_name(azure_embedding_unit_test_env: None) -> None:
client = AzureOpenAIEmbeddingClient(
deployment_name="text-embedding-3-small",
api_key="test-key",
endpoint="https://test.openai.azure.com/",
)
assert client.model_id == "text-embedding-3-small"
def test_azure_construction_with_existing_client(azure_embedding_unit_test_env: None) -> None:
mock_client = MagicMock()
client = AzureOpenAIEmbeddingClient(
deployment_name="my-deployment",
async_client=mock_client,
)
assert client.model_id == "my-deployment"
assert client.client is mock_client
def test_azure_construction_missing_deployment_name_raises(azure_embedding_unit_test_env: None) -> None:
with pytest.raises(ValueError, match="deployment name is required"):
AzureOpenAIEmbeddingClient(
api_key="test-key",
endpoint="https://test.openai.azure.com/",
)
def test_azure_construction_missing_credentials_raises(azure_embedding_unit_test_env: None) -> None:
with pytest.raises(ValueError, match="api_key, credential, or a client"):
AzureOpenAIEmbeddingClient(
deployment_name="test",
endpoint="https://test.openai.azure.com/",
)
async def test_azure_get_embeddings(azure_embedding_unit_test_env: None) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1, 0.2]],
)
mock_async_client = MagicMock()
mock_async_client.embeddings = MagicMock()
mock_async_client.embeddings.create = AsyncMock(return_value=mock_response)
client = AzureOpenAIEmbeddingClient(
deployment_name="text-embedding-3-small",
async_client=mock_async_client,
)
result = await client.get_embeddings(["hello"])
assert len(result) == 1
assert result[0].vector == [0.1, 0.2]
def test_azure_otel_provider_name(azure_embedding_unit_test_env: None) -> None:
mock_client = MagicMock()
client = AzureOpenAIEmbeddingClient(
deployment_name="test",
async_client=mock_client,
)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
not os.getenv("AZURE_OPENAI_ENDPOINT")
or (not os.getenv("AZURE_OPENAI_API_KEY") and not os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")),
reason="No Azure OpenAI credentials provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings() -> None:
"""End-to-end test of Azure OpenAI embedding generation."""
client = AzureOpenAIEmbeddingClient()
result = await client.get_embeddings(["hello world"])
assert len(result) == 1
assert isinstance(result[0].vector, list)
assert len(result[0].vector) > 0
assert all(isinstance(v, float) for v in result[0].vector)
assert result[0].model_id is not None
assert result.usage is not None
assert result.usage["input_token_count"] > 0
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings_multiple() -> None:
"""Test Azure OpenAI embedding generation for multiple inputs."""
client = AzureOpenAIEmbeddingClient()
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(e.vector) for e in result]
assert all(d == dims[0] for d in dims)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
"""Test Azure OpenAI embedding generation with custom dimensions."""
client = AzureOpenAIEmbeddingClient()
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 256
python/packages/core/tests/azure/test_azure_responses_client.py
-730
View File
@@ -1,730 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
import os
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import MagicMock
import pytest
from azure.identity import AzureCliCredential
from pydantic import BaseModel
from pytest import param
from agent_framework import (
Agent,
AgentResponse,
ChatResponse,
Content,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework.azure import AzureOpenAIResponsesClient
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests.",
)
logger = logging.getLogger(__name__)
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str
@tool(approval_mode="never_require")
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
return f"The weather in {location} is sunny and 72°F."
async def create_vector_store(
client: AzureOpenAIResponsesClient,
) -> tuple[str, Content]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."),
purpose="assistants",
)
vector_store = await client.client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
if result.last_error is not None:
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, Content.from_hosted_vector_store(vector_store_id=vector_store.id)
async def delete_vector_store(client: AzureOpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after tests."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"}) # type: ignore
def test_init_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
model_id = "test_model_id"
azure_responses_client = AzureOpenAIResponsesClient(deployment_name=model_id)
assert azure_responses_client.model_id == model_id
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test that model_id kwarg correctly sets the deployment name (issue #4299)."""
azure_responses_client = AzureOpenAIResponsesClient(model_id="gpt-4o")
assert azure_responses_client.model_id == "gpt-4o"
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg_does_not_override_deployment_name(
azure_openai_unit_test_env: dict[str, str],
) -> None:
"""Test that deployment_name takes precedence over model_id kwarg (issue #4299)."""
azure_responses_client = AzureOpenAIResponsesClient(deployment_name="my-deployment", model_id="gpt-4o")
assert azure_responses_client.model_id == "my-deployment"
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_model_id_kwarg_none(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test that model_id=None does not override the env-var deployment name."""
azure_responses_client = AzureOpenAIResponsesClient(model_id=None)
assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
def test_init_with_default_header(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
# Test successful initialization
azure_responses_client = AzureOpenAIResponsesClient(
default_headers=default_headers,
)
assert azure_responses_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(azure_responses_client, SupportsChatGetResponse)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in azure_responses_client.client.default_headers
assert azure_responses_client.client.default_headers[key] == value
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]], indirect=True)
def test_init_with_empty_model_id(azure_openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ValueError):
AzureOpenAIResponsesClient()
def test_init_with_project_client(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with an existing AIProjectClient."""
from unittest.mock import patch
from openai import AsyncOpenAI
# Create a mock AIProjectClient that returns a mock AsyncOpenAI client
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with patch(
"agent_framework.azure._responses_client.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_client=mock_project_client,
deployment_name="gpt-4o",
)
assert azure_responses_client.model_id == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_with_project_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with a project endpoint and credential."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
with patch(
"agent_framework.azure._responses_client.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_endpoint="https://test-project.services.ai.azure.com",
deployment_name="gpt-4o",
credential=AzureCliCredential(),
)
assert azure_responses_client.model_id == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_create_client_from_project_with_project_client() -> None:
"""Test _create_client_from_project with an existing project client."""
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=mock_project_client,
project_endpoint=None,
credential=None,
)
assert result is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
def test_create_client_from_project_with_endpoint() -> None:
"""Test _create_client_from_project with a project endpoint."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_credential = MagicMock()
with patch("agent_framework.azure._responses_client.AIProjectClient") as MockAIProjectClient:
mock_instance = MockAIProjectClient.return_value
mock_instance.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=mock_credential,
)
assert result is mock_openai_client
MockAIProjectClient.assert_called_once()
mock_instance.get_openai_client.assert_called_once()
def test_create_client_from_project_missing_endpoint() -> None:
"""Test _create_client_from_project raises error when endpoint is missing."""
with pytest.raises(ValueError, match="project endpoint is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint=None,
credential=MagicMock(),
)
def test_create_client_from_project_missing_credential() -> None:
"""Test _create_client_from_project raises error when credential is missing."""
with pytest.raises(ValueError, match="credential is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=None,
)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
settings = {
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"],
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
"default_headers": default_headers,
}
azure_responses_client = AzureOpenAIResponsesClient.from_dict(settings)
dumped_settings = azure_responses_client.to_dict()
assert dumped_settings["deployment_name"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert "api_key" not in dumped_settings
# Assert that the default header we added is present in the dumped_settings default headers
for key, value in default_headers.items():
assert key in dumped_settings["default_headers"]
assert dumped_settings["default_headers"][key] == value
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
assert "User-Agent" not in dumped_settings["default_headers"]
# region Integration Tests
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
# Simple ChatOptions - just verify they don't fail
param("temperature", 0.7, False, id="temperature"),
param("top_p", 0.9, False, id="top_p"),
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("frequency_penalty", 0.5, False, id="frequency_penalty"),
param("presence_penalty", 0.3, False, id="presence_penalty"),
param("stop", ["END"], False, id="stop"),
param("allow_multiple_tool_calls", True, False, id="allow_multiple_tool_calls"),
param("tool_choice", "none", True, id="tool_choice_none"),
# OpenAIResponsesOptions - just verify they don't fail
param("safety_identifier", "user-hash-abc123", False, id="safety_identifier"),
param("truncation", "auto", False, id="truncation"),
param("top_logprobs", 5, False, id="top_logprobs"),
param("prompt_cache_key", "test-cache-key", False, id="prompt_cache_key"),
param("max_tool_calls", 3, False, id="max_tool_calls"),
# Complex options requiring output validation
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param(
"tool_choice",
{"mode": "required", "required_function_name": "get_weather"},
True,
id="tool_choice_required",
),
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": [
"location",
"conditions",
"temperature_c",
"advisory",
],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
"""Parametrized test covering all ChatOptions and OpenAIResponsesOptions.
Tests both streaming and non-streaming modes for each option to ensure
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
# Prepare test message
if option_name == "tools" or option_name == "tool_choice":
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name == "response_format":
# Use prompt that works well with structured output
messages = [
Message(role="user", text="The weather in Seattle is sunny"),
Message(role="user", text="What is the weather in Seattle?"),
]
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name == "tool_choice":
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name == "tools" or option_name == "tool_choice":
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name == "response_format":
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_web_search() -> None:
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
for streaming in [False, True]:
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [AzureOpenAIResponsesClient.get_web_search_tool()],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [
AzureOpenAIResponsesClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})
],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
file_id, vector_store = await create_vector_store(azure_responses_client)
try:
# Test that the client will use the file search tool
response = await azure_responses_client.get_response(
messages=[
Message(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
options={
"tools": [
AzureOpenAIResponsesClient.get_file_search_tool(vector_store_ids=[vector_store.vector_store_id])
],
"tool_choice": "auto",
},
)
assert "sunny" in response.text.lower()
assert "75" in response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
file_id, vector_store = await create_vector_store(azure_responses_client)
# Test that the client will use the file search tool
try:
response_stream = azure_responses_client.get_response(
messages=[
Message(
role="user",
text="What is the weather today? Do a file search to find the answer.",
)
],
stream=True,
options={
"tools": [
AzureOpenAIResponsesClient.get_file_search_tool(vector_store_ids=[vector_store.vector_store_id])
],
"tool_choice": "auto",
},
)
full_response = await response_stream.get_final_response()
assert "sunny" in full_response.text.lower()
assert "75" in full_response.text
finally:
await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_mcp_tool() -> None:
"""Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
messages=[Message(role="user", text="How to create an Azure storage account using az cli?")],
options={
# this needs to be high enough to handle the full MCP tool response.
"max_tokens": 5000,
"tools": AzureOpenAIResponsesClient.get_mcp_tool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
),
},
)
assert isinstance(response, ChatResponse)
# MCP server may return empty response intermittently - skip test rather than fail
if not response.text:
pytest.skip("MCP server returned empty response - service-side issue")
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with code interpreter tool."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
messages=[
Message(
role="user",
text="Calculate the sum of numbers from 1 to 10 using Python code.",
)
],
options={
"tools": [AzureOpenAIResponsesClient.get_code_interpreter_tool()],
},
)
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My hobby is photography. Remember this.", session=session, store=True)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved session
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_responses_client_tool_rich_content_image() -> None:
"""Test that Azure OpenAI Responses client can handle tool results containing images."""
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@tool(approval_mode="never_require")
def get_test_image() -> Content:
"""Return a test image for analysis."""
return Content.from_data(data=image_bytes, media_type="image/jpeg")
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
messages = [
Message(
role="user",
text="Call the get_test_image tool and describe what you see.",
)
]
options: dict[str, Any] = {"tools": [get_test_image], "tool_choice": "auto"}
if streaming:
response = await client.get_response(messages=messages, stream=True, options=options).get_final_response()
else:
response = await client.get_response(messages=messages, options=options)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
# sample_image.jpg contains a photo of a house; the model should mention it.
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
# region Integration with Foundry V2
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/")
or os.getenv("AZURE_AI_MODEL_DEPLOYMENT_NAME", "") == "",
reason="No real AZURE_AI_PROJECT_ENDPOINT or AZURE_AI_MODEL_DEPLOYMENT_NAME provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_function_call_roundtrip_preserves_fidelity():
"""Test that function calls roundtrip correctly with full fidelity preserved.
This verifies the changes where:
1. raw_representation is preserved when parsing function calls
2. fc_id and status are included in additional_properties
3. When re-sending messages, the full object fidelity is preserved
"""
call_count = 0
@tool(name="get_weather", approval_mode="never_require")
async def get_weather_tool(location: str) -> str:
"""Get weather for a location."""
nonlocal call_count
call_count += 1
return f"Weather in {location} is sunny, 72F"
client = AzureOpenAIResponsesClient(
project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
deployment_name=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
credential=AzureCliCredential(),
)
async with Agent(
client=client,
name="WeatherAgent",
instructions="You help check weather. Use get_weather when asked about weather.",
tools=[get_weather_tool],
default_options={"store": False}, # Store messages locally to test fidelity across messages
) as agent:
session = agent.create_session()
# First request - should invoke the tool
response1 = await agent.run("What is the weather in Seattle?", session=session)
assert response1 is not None
assert response1.text is not None
assert call_count >= 1
# Verify the response contains expected content
response_text = response1.text.lower()
assert "seattle" in response_text or "sunny" in response_text or "72" in response_text
# Second request - should work correctly with the preserved conversation
response2 = await agent.run("And how about in Portland?", session=session)
assert response2 is not None
assert response2.text is not None
assert call_count >= 2
# endregion
python/packages/core/tests/azure/test_entra_id_authentication.py
-61
View File
@@ -1,61 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import MagicMock, patch
import pytest
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from agent_framework.azure._entra_id_authentication import (
resolve_credential_to_token_provider,
)
from agent_framework.exceptions import ChatClientInvalidAuthException
TOKEN_ENDPOINT = "https://cognitiveservices.azure.com/.default"
def test_resolve_sync_credential_returns_provider() -> None:
"""Test that a sync TokenCredential is resolved via azure.identity.get_bearer_token_provider."""
mock_credential = MagicMock(spec=TokenCredential)
mock_provider = MagicMock(return_value="token-string")
with patch("azure.identity.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
assert result is mock_provider
def test_resolve_async_credential_returns_provider() -> None:
"""Test that an AsyncTokenCredential is resolved via azure.identity.aio.get_bearer_token_provider."""
mock_credential = MagicMock(spec=AsyncTokenCredential)
mock_provider = MagicMock(return_value="token-string")
with patch("azure.identity.aio.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
assert result is mock_provider
def test_resolve_callable_provider_passthrough() -> None:
"""Test that a callable token provider is returned as-is, without needing token_endpoint."""
my_provider = lambda: "my-token" # noqa: E731
# Works with token_endpoint
assert resolve_credential_to_token_provider(my_provider, TOKEN_ENDPOINT) is my_provider
# Also works without token_endpoint
assert resolve_credential_to_token_provider(my_provider, None) is my_provider
assert resolve_credential_to_token_provider(my_provider, "") is my_provider
def test_resolve_missing_endpoint_raises() -> None:
"""Test that missing token endpoint raises ChatClientInvalidAuthException."""
mock_credential = MagicMock(spec=TokenCredential)
with pytest.raises(ChatClientInvalidAuthException, match="A token endpoint must be provided"):
resolve_credential_to_token_provider(mock_credential, "")
with pytest.raises(ChatClientInvalidAuthException, match="A token endpoint must be provided"):
resolve_credential_to_token_provider(mock_credential, None) # type: ignore[arg-type]
python/packages/core/tests/core/test_agents.py
+3 -3
View File
@@ -1950,13 +1950,13 @@ async def test_shared_local_storage_cross_provider_responses_history_does_not_le
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from agent_framework._sessions import InMemoryHistoryProvider
from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient, OpenAIChatCompletionClient
@tool(approval_mode="never_require")
def search_hotels(city: str) -> str:
return f"Found 3 hotels in {city}"
responses_client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
responses_client = OpenAIChatClient(model="test-model", api_key="test-key")
responses_agent = Agent(
client=responses_client,
tools=[search_hotels],
@@ -2024,7 +2024,7 @@ async def test_shared_local_storage_cross_provider_responses_history_does_not_le
responses_replay_call = next(item for item in responses_replay_input if item.get("type") == "function_call")
assert responses_replay_call["id"] == "fc_provider123"
chat_client = OpenAIChatClient(model_id="test-model", api_key="test-key")
chat_client = OpenAIChatCompletionClient(model="test-model", api_key="test-key")
chat_agent = Agent(client=chat_client)
chat_response = ChatCompletion(
python/packages/core/tests/core/test_clients.py
+14 -14
View File
@@ -66,10 +66,10 @@ def test_base_client_as_agent_uses_explicit_additional_properties(chat_client_ba
assert agent.additional_properties == {"team": "core"}
def test_openai_chat_client_get_response_docstring_surfaces_layered_runtime_docs() -> None:
from agent_framework.openai import OpenAIChatClient
def test_openai_chat_completion_client_get_response_docstring_surfaces_layered_runtime_docs() -> None:
from agent_framework.openai import OpenAIChatCompletionClient
docstring = inspect.getdoc(OpenAIChatClient.get_response)
docstring = inspect.getdoc(OpenAIChatCompletionClient.get_response)
assert docstring is not None
assert "Get a response from a chat client." in docstring
@@ -78,12 +78,12 @@ def test_openai_chat_client_get_response_docstring_surfaces_layered_runtime_docs
assert "function_middleware: Optional per-call function middleware." not in docstring
def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
from agent_framework.openai import OpenAIChatClient
def test_openai_chat_completion_client_get_response_is_defined_on_openai_class() -> None:
from agent_framework.openai import OpenAIChatCompletionClient
signature = inspect.signature(OpenAIChatClient.get_response)
signature = inspect.signature(OpenAIChatCompletionClient.get_response)
assert OpenAIChatClient.get_response.__qualname__ == "OpenAIChatClient.get_response"
assert OpenAIChatCompletionClient.get_response.__qualname__ == "OpenAIChatCompletionClient.get_response"
assert "middleware" in signature.parameters
@@ -349,15 +349,15 @@ def test_openai_responses_client_supports_all_tool_protocols():
assert isinstance(OpenAIResponsesClient, SupportsFileSearchTool)
def test_openai_chat_client_supports_web_search_only():
def test_openai_chat_completion_client_supports_web_search_only():
"""Test that OpenAIChatClient only supports web search tool."""
from agent_framework.openai import OpenAIChatClient
from agent_framework.openai import OpenAIChatCompletionClient
assert not isinstance(OpenAIChatClient, SupportsCodeInterpreterTool)
assert isinstance(OpenAIChatClient, SupportsWebSearchTool)
assert not isinstance(OpenAIChatClient, SupportsImageGenerationTool)
assert not isinstance(OpenAIChatClient, SupportsMCPTool)
assert not isinstance(OpenAIChatClient, SupportsFileSearchTool)
assert not isinstance(OpenAIChatCompletionClient, SupportsCodeInterpreterTool)
assert isinstance(OpenAIChatCompletionClient, SupportsWebSearchTool)
assert not isinstance(OpenAIChatCompletionClient, SupportsImageGenerationTool)
assert not isinstance(OpenAIChatCompletionClient, SupportsMCPTool)
assert not isinstance(OpenAIChatCompletionClient, SupportsFileSearchTool)
def test_openai_assistants_client_supports_code_interpreter_and_file_search():
python/packages/core/tests/core/test_foundry_namespace.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright (c) Microsoft. All rights reserved.
import pytest
from agent_framework_foundry import FoundryChatClient, FoundryMemoryProvider
from agent_framework_foundry_local import FoundryLocalClient
import agent_framework.azure as azure
import agent_framework.foundry as foundry
def test_foundry_namespace_exposes_cloud_and_local_symbols() -> None:
assert foundry.FoundryChatClient is FoundryChatClient
assert foundry.FoundryMemoryProvider is FoundryMemoryProvider
assert foundry.FoundryLocalClient is FoundryLocalClient
assert "FoundryChatClient" in dir(foundry)
assert "FoundryLocalClient" in dir(foundry)
def test_azure_namespace_no_longer_exposes_foundry_symbols() -> None:
assert "FoundryChatClient" not in dir(azure)
assert "FoundryLocalClient" not in dir(azure)
with pytest.raises(AttributeError, match="Module `azure` has no attribute FoundryChatClient\\."):
_ = azure.FoundryChatClient
python/packages/core/tests/core/test_mcp.py
+12 -7
View File
@@ -42,6 +42,14 @@ skip_if_mcp_integration_tests_disabled = pytest.mark.skipif(
)
def _mcp_result_to_text(result: str | list[Content]) -> str:
"""Normalize an MCP tool result to text for assertions."""
if isinstance(result, str):
return result
text = "\n".join(content.text for content in result if content.type == "text" and content.text)
return text or str(result)
# Helper function tests
def test_normalize_mcp_name():
"""Test MCP name normalization."""
@@ -1401,8 +1409,7 @@ async def test_streamable_http_integration():
assert hasattr(func, "name")
assert hasattr(func, "description")
result = await func.invoke(query="What is Agent Framework?")
assert isinstance(result, str)
result = _mcp_result_to_text(await func.invoke(query="What is Agent Framework?"))
assert len(result) > 0
@@ -1430,7 +1437,7 @@ async def test_mcp_connection_reset_integration():
# Get the first function and invoke it
func = tool.functions[0]
first_result = await func.invoke(query="What is Agent Framework?")
first_result = _mcp_result_to_text(await func.invoke(query="What is Agent Framework?"))
assert first_result is not None
assert len(first_result) > 0
@@ -1456,7 +1463,7 @@ async def test_mcp_connection_reset_integration():
tool.session.call_tool = call_tool_with_error
# Invoke the function again - this should trigger automatic reconnection on ClosedResourceError
second_result = await func.invoke(query="What is Agent Framework?")
second_result = _mcp_result_to_text(await func.invoke(query="What is Agent Framework?"))
assert second_result is not None
assert len(second_result) > 0
@@ -1469,10 +1476,8 @@ async def test_mcp_connection_reset_integration():
# Verify tools are still available after reconnection
assert len(tool.functions) > 0
# Both results should be valid strings (we don't compare content as it may vary)
assert isinstance(first_result, str)
# Both results should include text (we don't compare content as it may vary)
assert len(first_result) > 0
assert isinstance(second_result, str)
assert len(second_result) > 0
python/packages/core/tests/core/test_types.py
+13 -13
View File
@@ -1031,11 +1031,11 @@ def test_chat_tool_mode_from_dict():
def test_chat_options_init() -> None:
"""Test that ChatOptions can be created as a TypedDict."""
options: ChatOptions = {}
assert options.get("model_id") is None
assert options.get("model") is None
# With values
options_with_model: ChatOptions = {"model_id": "gpt-4o", "temperature": 0.7}
assert options_with_model.get("model_id") == "gpt-4o"
options_with_model: ChatOptions = {"model": "gpt-4o", "temperature": 0.7}
assert options_with_model.get("model") == "gpt-4o"
assert options_with_model.get("temperature") == 0.7
@@ -1069,18 +1069,18 @@ def test_chat_options_tool_choice_validation():
def test_chat_options_merge(tool_tool, ai_tool) -> None:
"""Test merge_chat_options utility function."""
options1: ChatOptions = {
"model_id": "gpt-4o",
"model": "gpt-4o",
"tools": [tool_tool],
"logit_bias": {"x": 1},
"metadata": {"a": "b"},
}
options2: ChatOptions = {"model_id": "gpt-4.1", "tools": [ai_tool]}
options2: ChatOptions = {"model": "gpt-4.1", "tools": [ai_tool]}
assert options1 != options2
# Merge options - override takes precedence for non-collection fields
options3 = merge_chat_options(options1, options2)
assert options3.get("model_id") == "gpt-4.1"
assert options3.get("model") == "gpt-4.1"
assert options3.get("tools") == [tool_tool, ai_tool] # tools are combined
assert options3.get("logit_bias") == {"x": 1} # base value preserved
assert options3.get("metadata") == {"a": "b"} # base value preserved
@@ -1089,7 +1089,7 @@ def test_chat_options_merge(tool_tool, ai_tool) -> None:
def test_chat_options_and_tool_choice_override() -> None:
"""Test that tool_choice from other takes precedence in ChatOptions merge."""
# Agent-level defaults to "auto"
agent_options: ChatOptions = {"model_id": "gpt-4o", "tool_choice": "auto"}
agent_options: ChatOptions = {"model": "gpt-4o", "tool_choice": "auto"}
# Run-level specifies "required"
run_options: ChatOptions = {"tool_choice": "required"}
@@ -1097,19 +1097,19 @@ def test_chat_options_and_tool_choice_override() -> None:
# Run-level should override agent-level
assert merged.get("tool_choice") == "required"
assert merged.get("model_id") == "gpt-4o" # Other fields preserved
assert merged.get("model") == "gpt-4o" # Other fields preserved
def test_chat_options_and_tool_choice_none_in_other_uses_self() -> None:
"""Test that when other.tool_choice is None, self.tool_choice is used."""
agent_options: ChatOptions = {"tool_choice": "auto"}
run_options: ChatOptions = {"model_id": "gpt-4.1"} # tool_choice is None
run_options: ChatOptions = {"model": "gpt-4.1"} # tool_choice is None
merged = merge_chat_options(agent_options, run_options)
# Should keep agent-level tool_choice since run-level is None
assert merged.get("tool_choice") == "auto"
assert merged.get("model_id") == "gpt-4.1"
assert merged.get("model") == "gpt-4.1"
def test_chat_options_and_tool_choice_with_tool_mode() -> None:
@@ -1845,7 +1845,7 @@ def test_chat_response_complex_serialization():
"output_token_count": 8,
"total_token_count": 13,
},
"model_id": "gpt-4", # Test alias handling
"model": "gpt-4", # Test alias handling
}
response = ChatResponse.from_dict(response_data)
@@ -1861,7 +1861,7 @@ def test_chat_response_complex_serialization():
assert isinstance(response_dict["messages"][0], dict)
assert isinstance(response_dict["finish_reason"], str) # FinishReason serializes to string
assert isinstance(response_dict["usage_details"], dict)
assert response_dict["model_id"] == "gpt-4" # Should serialize as model_id
assert response_dict["model"] == "gpt-4" # Should serialize as model_id
def test_chat_response_update_all_content_types():
@@ -2309,7 +2309,7 @@ def test_chat_response_deepcopy_deep_copies_additional_properties():
"total_token_count": 30,
},
"response_id": "resp-123",
"model_id": "gpt-4",
"model": "gpt-4",
},
id="chat_response",
),
python/packages/core/tests/openai/conftest.py
-51
View File
@@ -1,51 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
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 {}
@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-dummy-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 in exclude_list:
monkeypatch.delenv(key, raising=False) # type: ignore
continue
monkeypatch.setenv(key, value) # type: ignore
return env_vars
python/packages/core/tests/openai/test_assistant_provider.py
-814
View File
@@ -1,814 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from openai.types.beta.assistant import Assistant
from pydantic import BaseModel, Field
from agent_framework import Agent, normalize_tools, tool
from agent_framework.openai import OpenAIAssistantProvider, OpenAIAssistantsClient
from agent_framework.openai._shared import from_assistant_tools, to_assistant_tools
# region Test Helpers
def create_mock_assistant(
assistant_id: str = "asst_test123",
name: str = "TestAssistant",
model: str = "gpt-4",
instructions: str | None = "You are a helpful assistant.",
description: str | None = None,
tools: list[Any] | None = None,
) -> Assistant:
"""Create a mock Assistant object."""
mock = MagicMock(spec=Assistant)
mock.id = assistant_id
mock.name = name
mock.model = model
mock.instructions = instructions
mock.description = description
mock.tools = tools or []
return mock
def create_function_tool(name: str, description: str = "A test function") -> MagicMock:
"""Create a mock FunctionTool."""
mock = MagicMock()
mock.type = "function"
mock.function = MagicMock()
mock.function.name = name
mock.function.description = description
return mock
def create_code_interpreter_tool() -> MagicMock:
"""Create a mock CodeInterpreterTool."""
mock = MagicMock()
mock.type = "code_interpreter"
return mock
def create_file_search_tool() -> MagicMock:
"""Create a mock FileSearchTool."""
mock = MagicMock()
mock.type = "file_search"
return mock
@pytest.fixture
def mock_async_openai() -> MagicMock:
"""Mock AsyncOpenAI client."""
mock_client = MagicMock()
# Mock beta.assistants
mock_client.beta.assistants.create = AsyncMock(
return_value=create_mock_assistant(assistant_id="asst_created123", name="CreatedAssistant")
)
mock_client.beta.assistants.retrieve = AsyncMock(
return_value=create_mock_assistant(assistant_id="asst_retrieved123", name="RetrievedAssistant")
)
mock_client.beta.assistants.delete = AsyncMock()
# Mock close method
mock_client.close = AsyncMock()
return mock_client
# Test function for tool validation
def get_weather(location: Annotated[str, Field(description="The location")]) -> str:
"""Get the weather for a location."""
return f"Weather in {location}: sunny"
def search_database(query: Annotated[str, Field(description="Search query")]) -> str:
"""Search the database."""
return f"Results for: {query}"
# Pydantic model for structured output tests
class WeatherResponse(BaseModel):
location: str
temperature: float
conditions: str
# endregion
# region Initialization Tests
class TestOpenAIAssistantProviderInit:
"""Tests for provider initialization."""
def test_init_with_client(self, mock_async_openai: MagicMock) -> None:
"""Test initialization with existing AsyncOpenAI client."""
provider = OpenAIAssistantProvider(mock_async_openai)
assert provider._client is mock_async_openai # type: ignore[reportPrivateUsage]
assert provider._should_close_client is False # type: ignore[reportPrivateUsage]
def test_init_without_client_creates_one(self, openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization creates client from settings."""
provider = OpenAIAssistantProvider()
assert provider._client is not None # type: ignore[reportPrivateUsage]
assert provider._should_close_client is True # type: ignore[reportPrivateUsage]
def test_init_with_api_key(self) -> None:
"""Test initialization with explicit API key."""
provider = OpenAIAssistantProvider(api_key="sk-test-key")
assert provider._client is not None # type: ignore[reportPrivateUsage]
assert provider._should_close_client is True # type: ignore[reportPrivateUsage]
def test_init_fails_without_api_key(self) -> None:
"""Test initialization fails without API key when settings return None."""
from unittest.mock import patch
# Mock load_settings to return a dict with None for api_key
with patch("agent_framework.openai._assistant_provider.load_settings") as mock_load:
mock_load.return_value = {
"api_key": None,
"org_id": None,
"base_url": None,
"chat_model_id": None,
"responses_model_id": None,
}
with pytest.raises(ValueError) as exc_info:
OpenAIAssistantProvider()
assert "API key is required" in str(exc_info.value)
def test_init_with_org_id_and_base_url(self) -> None:
"""Test initialization with organization ID and base URL."""
provider = OpenAIAssistantProvider(
api_key="sk-test-key",
org_id="org-123",
base_url="https://custom.openai.com",
)
assert provider._client is not None # type: ignore[reportPrivateUsage]
class TestOpenAIAssistantProviderContextManager:
"""Tests for async context manager."""
async def test_context_manager_enter_exit(self, mock_async_openai: MagicMock) -> None:
"""Test async context manager entry and exit."""
provider = OpenAIAssistantProvider(mock_async_openai)
async with provider as p:
assert p is provider
async def test_context_manager_closes_owned_client(self, openai_unit_test_env: dict[str, str]) -> None:
"""Test that owned client is closed on exit."""
provider = OpenAIAssistantProvider()
client = provider._client # type: ignore[reportPrivateUsage]
assert client is not None
client.close = AsyncMock()
async with provider:
pass
client.close.assert_called_once()
async def test_context_manager_does_not_close_external_client(self, mock_async_openai: MagicMock) -> None:
"""Test that external client is not closed on exit."""
provider = OpenAIAssistantProvider(mock_async_openai)
async with provider:
pass
mock_async_openai.close.assert_not_called()
# endregion
# region create_agent Tests
class TestOpenAIAssistantProviderCreateAgent:
"""Tests for create_agent method."""
async def test_create_agent_basic(self, mock_async_openai: MagicMock) -> None:
"""Test basic assistant creation."""
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.create_agent(
name="TestAgent",
model="gpt-4",
instructions="You are helpful.",
)
assert isinstance(agent, Agent)
assert agent.name == "CreatedAssistant"
mock_async_openai.beta.assistants.create.assert_called_once()
# Verify create was called with correct parameters
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert call_kwargs["name"] == "TestAgent"
assert call_kwargs["model"] == "gpt-4"
assert call_kwargs["instructions"] == "You are helpful."
async def test_create_agent_with_description(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with description."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="TestAgent",
model="gpt-4",
description="A test agent description",
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert call_kwargs["description"] == "A test agent description"
async def test_create_agent_with_function_tools(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with function tools."""
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.create_agent(
name="WeatherAgent",
model="gpt-4",
tools=[get_weather],
)
assert isinstance(agent, Agent)
# Verify tools were passed to create
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert "tools" in call_kwargs
assert len(call_kwargs["tools"]) == 1
assert call_kwargs["tools"][0]["type"] == "function"
assert call_kwargs["tools"][0]["function"]["name"] == "get_weather"
async def test_create_agent_with_tool(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with FunctionTool."""
provider = OpenAIAssistantProvider(mock_async_openai)
@tool
def my_function(x: int) -> int:
"""Double a number."""
return x * 2
await provider.create_agent(
name="TestAgent",
model="gpt-4",
tools=[my_function],
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert call_kwargs["tools"][0]["function"]["name"] == "my_function"
async def test_create_agent_with_code_interpreter(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with code interpreter."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="CodeAgent",
model="gpt-4",
tools=[OpenAIAssistantsClient.get_code_interpreter_tool()],
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert {"type": "code_interpreter"} in call_kwargs["tools"]
async def test_create_agent_with_file_search(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with file search."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="SearchAgent",
model="gpt-4",
tools=[OpenAIAssistantsClient.get_file_search_tool()],
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert any(t["type"] == "file_search" for t in call_kwargs["tools"])
async def test_create_agent_with_file_search_max_results(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with file search and max_results."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="SearchAgent",
model="gpt-4",
tools=[OpenAIAssistantsClient.get_file_search_tool(max_num_results=10)],
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
file_search_tool = next(t for t in call_kwargs["tools"] if t["type"] == "file_search")
assert file_search_tool.get("file_search", {}).get("max_num_results") == 10
async def test_create_agent_with_mixed_tools(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with multiple tool types."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="MultiToolAgent",
model="gpt-4",
tools=[
get_weather,
OpenAIAssistantsClient.get_code_interpreter_tool(),
OpenAIAssistantsClient.get_file_search_tool(),
],
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert len(call_kwargs["tools"]) == 3
async def test_create_agent_with_metadata(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with metadata."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="TestAgent",
model="gpt-4",
metadata={"env": "test", "version": "1.0"},
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert call_kwargs["metadata"] == {"env": "test", "version": "1.0"}
async def test_create_agent_with_response_format_pydantic(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with Pydantic response format via default_options."""
provider = OpenAIAssistantProvider(mock_async_openai)
await provider.create_agent(
name="StructuredAgent",
model="gpt-4",
default_options={"response_format": WeatherResponse},
)
call_kwargs = mock_async_openai.beta.assistants.create.call_args.kwargs
assert call_kwargs["response_format"]["type"] == "json_schema"
assert call_kwargs["response_format"]["json_schema"]["name"] == "WeatherResponse"
async def test_create_agent_returns_chat_agent(self, mock_async_openai: MagicMock) -> None:
"""Test that create_agent returns a Agent instance."""
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.create_agent(
name="TestAgent",
model="gpt-4",
)
assert isinstance(agent, Agent)
# endregion
# region get_agent Tests
class TestOpenAIAssistantProviderGetAgent:
"""Tests for get_agent method."""
async def test_get_agent_basic(self, mock_async_openai: MagicMock) -> None:
"""Test retrieving an existing assistant."""
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.get_agent(assistant_id="asst_123")
assert isinstance(agent, Agent)
mock_async_openai.beta.assistants.retrieve.assert_called_once_with("asst_123")
async def test_get_agent_with_instructions_override(self, mock_async_openai: MagicMock) -> None:
"""Test retrieving assistant with instruction override."""
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.get_agent(
assistant_id="asst_123",
instructions="Custom instructions",
)
# Agent should be created successfully with the custom instructions
assert isinstance(agent, Agent)
assert agent.id == "asst_retrieved123"
async def test_get_agent_with_function_tools(self, mock_async_openai: MagicMock) -> None:
"""Test retrieving assistant with function tools provided."""
# Setup assistant with function tool
assistant = create_mock_assistant(tools=[create_function_tool("get_weather")])
mock_async_openai.beta.assistants.retrieve = AsyncMock(return_value=assistant)
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.get_agent(
assistant_id="asst_123",
tools=[get_weather],
)
assert isinstance(agent, Agent)
async def test_get_agent_validates_missing_function_tools(self, mock_async_openai: MagicMock) -> None:
"""Test that missing function tools raise ValueError."""
# Setup assistant with function tool
assistant = create_mock_assistant(tools=[create_function_tool("get_weather")])
mock_async_openai.beta.assistants.retrieve = AsyncMock(return_value=assistant)
provider = OpenAIAssistantProvider(mock_async_openai)
with pytest.raises(ValueError) as exc_info:
await provider.get_agent(assistant_id="asst_123")
assert "get_weather" in str(exc_info.value)
assert "no implementation was provided" in str(exc_info.value)
async def test_get_agent_validates_multiple_missing_function_tools(self, mock_async_openai: MagicMock) -> None:
"""Test validation with multiple missing function tools."""
assistant = create_mock_assistant(
tools=[create_function_tool("get_weather"), create_function_tool("search_database")]
)
mock_async_openai.beta.assistants.retrieve = AsyncMock(return_value=assistant)
provider = OpenAIAssistantProvider(mock_async_openai)
with pytest.raises(ValueError) as exc_info:
await provider.get_agent(assistant_id="asst_123")
error_msg = str(exc_info.value)
assert "get_weather" in error_msg or "search_database" in error_msg
async def test_get_agent_merges_hosted_tools(self, mock_async_openai: MagicMock) -> None:
"""Test that hosted tools are automatically included."""
assistant = create_mock_assistant(tools=[create_code_interpreter_tool(), create_file_search_tool()])
mock_async_openai.beta.assistants.retrieve = AsyncMock(return_value=assistant)
provider = OpenAIAssistantProvider(mock_async_openai)
agent = await provider.get_agent(assistant_id="asst_123")
# Hosted tools should be merged automatically
assert isinstance(agent, Agent)
# endregion
# region as_agent Tests
class TestOpenAIAssistantProviderAsAgent:
"""Tests for as_agent method."""
def test_as_agent_no_http_call(self, mock_async_openai: MagicMock) -> None:
"""Test that as_agent doesn't make HTTP calls."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant()
agent = provider.as_agent(assistant)
assert isinstance(agent, Agent)
# Verify no HTTP calls were made
mock_async_openai.beta.assistants.create.assert_not_called()
mock_async_openai.beta.assistants.retrieve.assert_not_called()
def test_as_agent_wraps_assistant(self, mock_async_openai: MagicMock) -> None:
"""Test wrapping an SDK Assistant object."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(
assistant_id="asst_wrap123",
name="WrappedAssistant",
instructions="Original instructions",
)
agent = provider.as_agent(assistant)
assert agent.id == "asst_wrap123"
assert agent.name == "WrappedAssistant"
# Instructions are passed to ChatOptions, not exposed as attribute
assert isinstance(agent, Agent)
def test_as_agent_with_instructions_override(self, mock_async_openai: MagicMock) -> None:
"""Test as_agent with instruction override."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(instructions="Original")
agent = provider.as_agent(assistant, instructions="Override")
# Agent should be created successfully with override instructions
assert isinstance(agent, Agent)
def test_as_agent_validates_function_tools(self, mock_async_openai: MagicMock) -> None:
"""Test that missing function tools raise ValueError."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(tools=[create_function_tool("get_weather")])
with pytest.raises(ValueError) as exc_info:
provider.as_agent(assistant)
assert "get_weather" in str(exc_info.value)
def test_as_agent_with_function_tools_provided(self, mock_async_openai: MagicMock) -> None:
"""Test as_agent with function tools provided."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(tools=[create_function_tool("get_weather")])
agent = provider.as_agent(assistant, tools=[get_weather])
assert isinstance(agent, Agent)
def test_as_agent_merges_hosted_tools(self, mock_async_openai: MagicMock) -> None:
"""Test that hosted tools are merged automatically."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(tools=[create_code_interpreter_tool()])
agent = provider.as_agent(assistant)
assert isinstance(agent, Agent)
def test_as_agent_hosted_tools_not_required(self, mock_async_openai: MagicMock) -> None:
"""Test that hosted tools don't require user implementations."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant = create_mock_assistant(tools=[create_code_interpreter_tool(), create_file_search_tool()])
# Should not raise - hosted tools don't need implementations
agent = provider.as_agent(assistant)
assert isinstance(agent, Agent)
# endregion
# region Tool Conversion Tests
class TestToolConversion:
"""Tests for tool conversion utilities (shared functions)."""
def test_to_assistant_tools_tool(self) -> None:
"""Test FunctionTool conversion to API format."""
@tool
def test_func(x: int) -> int:
"""Test function."""
return x
# Normalize tools first, then convert
normalized = normalize_tools([test_func])
api_tools = to_assistant_tools(normalized)
assert len(api_tools) == 1
assert api_tools[0]["type"] == "function"
assert api_tools[0]["function"]["name"] == "test_func"
def test_to_assistant_tools_callable(self) -> None:
"""Test raw callable conversion via normalize_tools."""
# normalize_tools converts callables to FunctionTool
normalized = normalize_tools([get_weather])
api_tools = to_assistant_tools(normalized)
assert len(api_tools) == 1
assert api_tools[0]["type"] == "function"
assert api_tools[0]["function"]["name"] == "get_weather"
def test_to_assistant_tools_code_interpreter(self) -> None:
"""Test code_interpreter tool dict conversion."""
api_tools = to_assistant_tools([OpenAIAssistantsClient.get_code_interpreter_tool()])
assert len(api_tools) == 1
assert api_tools[0] == {"type": "code_interpreter"}
def test_to_assistant_tools_file_search(self) -> None:
"""Test file_search tool dict conversion."""
api_tools = to_assistant_tools([OpenAIAssistantsClient.get_file_search_tool()])
assert len(api_tools) == 1
assert api_tools[0]["type"] == "file_search"
def test_to_assistant_tools_file_search_with_max_results(self) -> None:
"""Test file_search tool with max_results conversion."""
api_tools = to_assistant_tools([OpenAIAssistantsClient.get_file_search_tool(max_num_results=5)])
assert api_tools[0]["file_search"]["max_num_results"] == 5
def test_to_assistant_tools_dict(self) -> None:
"""Test raw dict tool passthrough."""
raw_tool = {"type": "function", "function": {"name": "custom", "description": "Custom tool"}}
api_tools = to_assistant_tools([raw_tool])
assert len(api_tools) == 1
assert api_tools[0] == raw_tool
def test_to_assistant_tools_empty(self) -> None:
"""Test conversion with no tools."""
api_tools = to_assistant_tools(None)
assert api_tools == []
def test_from_assistant_tools_code_interpreter(self) -> None:
"""Test converting code_interpreter tool from OpenAI format."""
assistant_tools = [create_code_interpreter_tool()]
tools = from_assistant_tools(assistant_tools)
assert len(tools) == 1
assert tools[0] == {"type": "code_interpreter"}
def test_from_assistant_tools_file_search(self) -> None:
"""Test converting file_search tool from OpenAI format."""
assistant_tools = [create_file_search_tool()]
tools = from_assistant_tools(assistant_tools)
assert len(tools) == 1
assert tools[0] == {"type": "file_search"}
def test_from_assistant_tools_function_skipped(self) -> None:
"""Test that function tools are skipped (no implementations)."""
assistant_tools = [create_function_tool("test_func")]
tools = from_assistant_tools(assistant_tools)
assert len(tools) == 0 # Function tools are skipped
def test_from_assistant_tools_empty(self) -> None:
"""Test conversion with no tools."""
tools = from_assistant_tools(None)
assert tools == []
# endregion
# region Tool Validation Tests
class TestToolValidation:
"""Tests for tool validation."""
def test_validate_missing_function_tool_raises(self, mock_async_openai: MagicMock) -> None:
"""Test that missing function tools raise ValueError."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_function_tool("my_function")]
with pytest.raises(ValueError) as exc_info:
provider._validate_function_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
assert "my_function" in str(exc_info.value)
def test_validate_all_tools_provided_passes(self, mock_async_openai: MagicMock) -> None:
"""Test that validation passes when all tools provided."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_function_tool("get_weather")]
# Should not raise
provider._validate_function_tools(assistant_tools, [get_weather]) # type: ignore[reportPrivateUsage]
def test_validate_hosted_tools_not_required(self, mock_async_openai: MagicMock) -> None:
"""Test that hosted tools don't require implementations."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_code_interpreter_tool(), create_file_search_tool()]
# Should not raise
provider._validate_function_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
def test_validate_with_tool(self, mock_async_openai: MagicMock) -> None:
"""Test validation with FunctionTool."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_function_tool("get_weather")]
wrapped = tool(get_weather)
# Should not raise
provider._validate_function_tools(assistant_tools, [wrapped]) # type: ignore[reportPrivateUsage]
def test_validate_partial_tools_raises(self, mock_async_openai: MagicMock) -> None:
"""Test that partial tool provision raises error."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [
create_function_tool("get_weather"),
create_function_tool("search_database"),
]
with pytest.raises(ValueError) as exc_info:
provider._validate_function_tools(assistant_tools, [get_weather]) # type: ignore[reportPrivateUsage]
assert "search_database" in str(exc_info.value)
# endregion
# region Tool Merging Tests
class TestToolMerging:
"""Tests for tool merging."""
def test_merge_code_interpreter(self, mock_async_openai: MagicMock) -> None:
"""Test merging code interpreter tool."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_code_interpreter_tool()]
merged = provider._merge_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
assert len(merged) == 1
assert merged[0] == {"type": "code_interpreter"}
def test_merge_file_search(self, mock_async_openai: MagicMock) -> None:
"""Test merging file search tool."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_file_search_tool()]
merged = provider._merge_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
assert len(merged) == 1
assert merged[0] == {"type": "file_search"}
def test_merge_with_user_tools(self, mock_async_openai: MagicMock) -> None:
"""Test merging hosted and user tools."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_code_interpreter_tool()]
merged = provider._merge_tools(assistant_tools, [get_weather]) # type: ignore[reportPrivateUsage]
assert len(merged) == 2
assert merged[0] == {"type": "code_interpreter"}
def test_merge_multiple_hosted_tools(self, mock_async_openai: MagicMock) -> None:
"""Test merging multiple hosted tools."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_code_interpreter_tool(), create_file_search_tool()]
merged = provider._merge_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
assert len(merged) == 2
def test_merge_single_user_tool(self, mock_async_openai: MagicMock) -> None:
"""Test merging with single user tool (not list)."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools: list[Any] = []
merged = provider._merge_tools(assistant_tools, get_weather) # type: ignore[reportPrivateUsage]
assert len(merged) == 1
# endregion
# region Integration Tests
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
reason="No real OPENAI_API_KEY provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_openai_integration_tests_disabled
class TestOpenAIAssistantProviderIntegration:
"""Integration tests requiring real OpenAI API."""
async def test_create_and_run_agent(self) -> None:
"""End-to-end test of creating and running an agent."""
provider = OpenAIAssistantProvider()
agent = await provider.create_agent(
name="IntegrationTestAgent",
model=os.environ.get("OPENAI_CHAT_MODEL_ID", "gpt-4"),
instructions="You are a helpful assistant. Respond briefly.",
)
try:
result = await agent.run("Say 'hello' and nothing else.")
result_text = str(result)
assert "hello" in result_text.lower()
finally:
# Clean up the assistant
await provider._client.beta.assistants.delete(agent.id) # type: ignore[reportPrivateUsage, union-attr]
async def test_create_agent_with_function_tools_integration(self) -> None:
"""Integration test with function tools."""
provider = OpenAIAssistantProvider()
@tool(approval_mode="never_require")
def get_current_time() -> str:
"""Get the current time."""
from datetime import datetime
return datetime.now().strftime("%H:%M")
agent = await provider.create_agent(
name="TimeAgent",
model=os.environ.get("OPENAI_CHAT_MODEL_ID", "gpt-4"),
instructions="You are a helpful assistant.",
tools=[get_current_time],
)
try:
result = await agent.run("What time is it? Use the get_current_time function.")
result_text = str(result)
# The response should contain time information
assert ":" in result_text or "time" in result_text.lower()
finally:
await provider._client.beta.assistants.delete(agent.id) # type: ignore[reportPrivateUsage, union-attr]
# endregion
python/packages/core/tests/openai/test_openai_assistants_client.py
-1833
View File
File diff suppressed because it is too large Load Diff
python/packages/core/tests/openai/test_openai_chat_client.py
-1607
View File
File diff suppressed because it is too large Load Diff
python/packages/core/tests/openai/test_openai_chat_client_base.py
-425
View File
@@ -1,425 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from copy import deepcopy
from datetime import datetime, timezone
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 ChatResponseUpdate, Message
from agent_framework.exceptions import ChatClientException
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(scope="function")
def chat_history() -> list[Message]:
return []
@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[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(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_messages_for_openai(chat_history), # type: ignore
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_chat_options(
mock_create: AsyncMock,
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(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_messages_for_openai(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[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(role="user", text="hello world"))
orig_chat_history = deepcopy(chat_history)
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_messages_for_openai(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[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(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[Message],
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(Message(role="user", text="hello world"))
openai_chat_completion = OpenAIChatClient()
async for msg in openai_chat_completion.get_response(
stream=True,
messages=chat_history,
):
assert isinstance(msg, ChatResponseUpdate)
assert msg.message_id is not None
assert msg.response_id is not None
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_messages_for_openai(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[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(role="user", text="hello world"))
openai_chat_completion = OpenAIChatClient()
with pytest.raises(ChatClientException):
await openai_chat_completion.get_response(
messages=chat_history,
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_additional_properties(
mock_create: AsyncMock,
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(Message(role="user", text="hello world"))
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(messages=chat_history, options={"reasoning_effort": "low"})
mock_create.assert_awaited_once_with(
model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
stream=False,
messages=openai_chat_completion._prepare_messages_for_openai(chat_history), # type: ignore
reasoning_effort="low",
)
# region Streaming
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_get_streaming(
mock_create: AsyncMock,
chat_history: list[Message],
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(Message(role="user", text="hello world"))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
async for msg in openai_chat_completion.get_response(
stream=True,
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_messages_for_openai(orig_chat_history), # type: ignore
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_get_streaming_singular(
mock_create: AsyncMock,
chat_history: list[Message],
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(Message(role="user", text="hello world"))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
async for msg in openai_chat_completion.get_response(
stream=True,
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_messages_for_openai(orig_chat_history), # type: ignore
)
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_get_streaming_structured_output_no_fcc(
mock_create: AsyncMock,
chat_history: list[Message],
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(Message(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_response(
stream=True,
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_get_streaming_no_fcc_in_response(
mock_create: AsyncMock,
chat_history: list[Message],
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(Message(role="user", text="hello world"))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
[
msg
async for msg in openai_chat_completion.get_response(
stream=True,
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_messages_for_openai(orig_chat_history), # type: ignore
)
# region UTC Timestamp Tests
def test_chat_response_created_at_uses_utc(openai_unit_test_env: dict[str, str]):
"""Test that ChatResponse.created_at uses UTC timestamp, not local time.
This is a regression test for the issue where created_at was using local time
but labeling it as UTC (with 'Z' suffix).
"""
# Use a specific Unix timestamp: 1733011890 = 2024-12-01T00:31:30Z (UTC)
# This ensures we test that the timestamp is actually converted to UTC
utc_timestamp = 1733011890
mock_response = ChatCompletion(
id="test_id",
choices=[
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
],
created=utc_timestamp,
model="test",
object="chat.completion",
)
client = OpenAIChatClient()
response = client._parse_response_from_openai(mock_response, {})
# Verify that created_at is correctly formatted as UTC
assert response.created_at is not None
assert response.created_at.endswith("Z"), "Timestamp should end with 'Z' for UTC"
# Parse the timestamp and verify it matches UTC time
expected_utc_time = datetime.fromtimestamp(utc_timestamp, tz=timezone.utc)
expected_formatted = expected_utc_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
assert response.created_at == expected_formatted, (
f"Expected UTC timestamp {expected_formatted}, got {response.created_at}"
)
def test_chat_response_update_created_at_uses_utc(openai_unit_test_env: dict[str, str]):
"""Test that ChatResponseUpdate.created_at uses UTC timestamp, not local time.
This is a regression test for the issue where created_at was using local time
but labeling it as UTC (with 'Z' suffix).
"""
# Use a specific Unix timestamp: 1733011890 = 2024-12-01T00:31:30Z (UTC)
utc_timestamp = 1733011890
mock_chunk = ChatCompletionChunk(
id="test_id",
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
created=utc_timestamp,
model="test",
object="chat.completion.chunk",
)
client = OpenAIChatClient()
response_update = client._parse_response_update_from_openai(mock_chunk)
# Verify that created_at is correctly formatted as UTC
assert response_update.created_at is not None
assert response_update.created_at.endswith("Z"), "Timestamp should end with 'Z' for UTC"
# Parse the timestamp and verify it matches UTC time
expected_utc_time = datetime.fromtimestamp(utc_timestamp, tz=timezone.utc)
expected_formatted = expected_utc_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
assert response_update.created_at == expected_formatted, (
f"Expected UTC timestamp {expected_formatted}, got {response_update.created_at}"
)
python/packages/core/tests/openai/test_openai_embedding_client.py
-243
View File
@@ -1,243 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
from agent_framework.openai import (
OpenAIEmbeddingClient,
OpenAIEmbeddingOptions,
)
def _make_openai_response(
embeddings: list[list[float]],
model: str = "text-embedding-3-small",
prompt_tokens: int = 5,
total_tokens: int = 5,
) -> CreateEmbeddingResponse:
"""Helper to create a mock OpenAI embeddings response."""
data = [OpenAIEmbedding(embedding=emb, index=i, object="embedding") for i, emb in enumerate(embeddings)]
return CreateEmbeddingResponse(
data=data,
model=model,
object="list",
usage=Usage(prompt_tokens=prompt_tokens, total_tokens=total_tokens),
)
@pytest.fixture
def openai_unit_test_env(monkeypatch: pytest.MonkeyPatch) -> None:
"""Set up environment variables for OpenAI embedding client."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL_ID", "text-embedding-3-small")
# --- OpenAI unit tests ---
def test_openai_construction_with_explicit_params() -> None:
client = OpenAIEmbeddingClient(
model_id="text-embedding-3-small",
api_key="test-key",
)
assert client.model_id == "text-embedding-3-small"
def test_openai_construction_from_env(openai_unit_test_env: None) -> None:
client = OpenAIEmbeddingClient()
assert client.model_id == "text-embedding-3-small"
def test_openai_construction_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(ValueError, match="API key is required"):
OpenAIEmbeddingClient(model_id="text-embedding-3-small")
def test_openai_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL_ID", raising=False)
with pytest.raises(ValueError, match="model ID is required"):
OpenAIEmbeddingClient(api_key="test-key")
async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
)
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock(return_value=mock_response)
result = await client.get_embeddings(["hello", "world"])
assert len(result) == 2
assert result[0].vector == [0.1, 0.2, 0.3]
assert result[1].vector == [0.4, 0.5, 0.6]
assert result[0].model_id == "text-embedding-3-small"
assert result[0].dimensions == 3
async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1]],
prompt_tokens=10,
total_tokens=10,
)
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock(return_value=mock_response)
result = await client.get_embeddings(["test"])
assert result.usage is not None
assert result.usage["input_token_count"] == 10
assert result.usage["total_token_count"] == 10
async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock(return_value=mock_response)
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["test"], options=options)
call_kwargs = client.client.embeddings.create.call_args[1]
assert call_kwargs["dimensions"] == 256
assert result.options is options
async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: None) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock(return_value=mock_response)
options: OpenAIEmbeddingOptions = {"encoding_format": "base64"}
await client.get_embeddings(["test"], options=options)
call_kwargs = client.client.embeddings.create.call_args[1]
assert call_kwargs["encoding_format"] == "base64"
async def test_openai_base64_decoding(openai_unit_test_env: None) -> None:
import base64
import struct
# Encode [0.1, 0.2, 0.3] as base64 little-endian floats
raw_floats = [0.1, 0.2, 0.3]
b64_str = base64.b64encode(struct.pack(f"<{len(raw_floats)}f", *raw_floats)).decode()
# Mock the embedding item to return a base64 string (as the API does with encoding_format=base64)
mock_item = MagicMock()
mock_item.embedding = b64_str
mock_item.index = 0
mock_response = MagicMock()
mock_response.data = [mock_item]
mock_response.model = "text-embedding-3-small"
mock_response.usage = MagicMock(prompt_tokens=3, total_tokens=3)
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock(return_value=mock_response)
options: OpenAIEmbeddingOptions = {"encoding_format": "base64"}
result = await client.get_embeddings(["test"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 3
assert result[0].dimensions == 3
for expected, actual in zip(raw_floats, result[0].vector):
assert abs(expected - actual) < 1e-6
async def test_openai_error_when_no_model_id() -> None:
client = OpenAIEmbeddingClient.__new__(OpenAIEmbeddingClient)
client.model_id = None
client.client = MagicMock()
client.additional_properties = {}
client.otel_provider_name = "openai"
with pytest.raises(ValueError, match="model_id is required"):
await client.get_embeddings(["test"])
async def test_openai_empty_values_returns_empty(openai_unit_test_env: None) -> None:
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
client.client.embeddings.create = AsyncMock()
result = await client.get_embeddings([])
assert len(result) == 0
assert result.usage is None
client.client.embeddings.create.assert_not_called()
# --- Integration tests ---
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
reason="No real OPENAI_API_KEY provided; skipping integration tests.",
)
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings() -> None:
"""End-to-end test of OpenAI embedding generation."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
result = await client.get_embeddings(["hello world"])
assert len(result) == 1
assert isinstance(result[0].vector, list)
assert len(result[0].vector) > 0
assert all(isinstance(v, float) for v in result[0].vector)
assert result[0].model_id is not None
assert result.usage is not None
assert result.usage["input_token_count"] > 0
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings_multiple() -> None:
"""Test embedding generation for multiple inputs."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(e.vector) for e in result]
assert all(d == dims[0] for d in dims)
@skip_if_openai_integration_tests_disabled
@pytest.mark.flaky
@pytest.mark.integration
async def test_integration_openai_get_embeddings_with_dimensions() -> None:
"""Test embedding generation with custom dimensions."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 256
python/packages/core/tests/openai/test_openai_responses_client.py
-3776
View File
File diff suppressed because it is too large Load Diff