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

[BREAKING] Python: fix OpenAI Azure routing and provider samples (#4925)

Browse Source 
* Python: fix OpenAI Azure routing and provider samples

Prefer OpenAI when OPENAI_API_KEY is present unless Azure is explicitly requested. Clarify constructor docs, keep deprecated Azure wrappers compatible with stricter settings validation, and refresh the provider samples and tests to use the current client patterns.

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

* fix bandit

* Python: align OpenAI embedding Azure routing

Extend the shared OpenAI-vs-Azure routing and credential behavior to the embedding client, add Azure embedding regression coverage, and refresh the embedding samples to use the generic client path.

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

* Python: fix embedding client pyright check

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

* Python: thin OpenAI embedding wrapper

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

* Python: document embedding overload routing

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

* Python: fix callable OpenAI key routing

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

* Python: fix Azure credential routing tests

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

* Python: address OpenAI review feedback

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

* Python: narrow Azure routing markers

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

* Python: refine OpenAI model fallback order

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

* Python: narrow Azure deployment docs

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

* Python: remove embedding routing wording

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

* Python: run embedding Azure integration tests

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

* changed variable name

* Python: expand OpenAI package README

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

* clarified readme

* Python: fix Azure OpenAI integration setup

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

* Python: correct Azure integration env mapping

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

* updated code to fix int tests

* test updates

* test fix

* fix test setup

* updates to tests and setup

* remove openai assistants int tests

* improvements in int tests

* fix env var

* fix env vars

* fix azure responses test

* trigger actions

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-27 14:33:39 +01:00
committed by GitHub
Unverified
parent 3611be82cf
commit cc0cfaaac8
103 changed files with 5451 additions and 4216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/openai/AGENTS.md
+4
View File
@@ -27,6 +27,10 @@ agent_framework_openai/
All clients follow the Raw + Full-Featured pattern (e.g., `RawOpenAIChatClient` + `OpenAIChatClient`).
The generic OpenAI clients support both OpenAI and Azure OpenAI routing. Precedence is:
explicit Azure inputs (`credential`, `azure_endpoint`, `api_version`) → OpenAI API key
(`OPENAI_API_KEY`) → Azure environment fallback (`AZURE_OPENAI_*`).
## Dependencies
- `agent-framework-core` — core abstractions
python/packages/openai/README.md
+93 -4
View File
@@ -1,17 +1,106 @@
# agent-framework-openai
OpenAI integration for Microsoft Agent Framework. Provides chat clients for the OpenAI Responses API and Chat Completions API.
OpenAI integration for Microsoft Agent Framework.
This package provides:
- `OpenAIChatClient` for the OpenAI Responses API
- `OpenAIChatCompletionClient` for the Chat Completions API
- `OpenAIEmbeddingClient` for embeddings
## Installation
```bash
pip install agent-framework-openai
pip install agent-framework-openai --pre
```
## Usage
## Which chat client should I use?
Use `OpenAIChatClient` for new work unless you specifically need the Chat Completions API.
- `OpenAIChatClient` uses the Responses API and is the preferred general-purpose chat client.
- `OpenAIChatCompletionClient` uses the Chat Completions API and is mainly for compatibility with
existing Chat Completions-based integrations.
The deprecated `OpenAIResponsesClient` alias points to `OpenAIChatClient`.
## Environment variables
### OpenAI
These variables are used when the client is configured for OpenAI:
| Variable | Purpose |
| --- | --- |
| `OPENAI_API_KEY` | OpenAI API key |
| `OPENAI_ORG_ID` | OpenAI organization ID |
| `OPENAI_BASE_URL` | Custom OpenAI-compatible base URL |
| `OPENAI_MODEL` | Generic fallback model |
| `OPENAI_RESPONSES_MODEL` | Preferred model for `OpenAIChatClient` |
| `OPENAI_CHAT_MODEL` | Preferred model for `OpenAIChatCompletionClient` |
| `OPENAI_EMBEDDING_MODEL` | Preferred model for `OpenAIEmbeddingClient` |
Model lookup order:
- `OpenAIChatClient`: `OPENAI_RESPONSES_MODEL` -> `OPENAI_MODEL`
- `OpenAIChatCompletionClient`: `OPENAI_CHAT_MODEL` -> `OPENAI_MODEL`
- `OpenAIEmbeddingClient`: `OPENAI_EMBEDDING_MODEL` -> `OPENAI_MODEL`
### Azure OpenAI
These variables are used when the client is configured for Azure OpenAI:
| Variable | Purpose |
| --- | --- |
| `AZURE_OPENAI_ENDPOINT` | Azure OpenAI resource endpoint |
| `AZURE_OPENAI_BASE_URL` | Full Azure OpenAI base URL (`.../openai/v1`) |
| `AZURE_OPENAI_API_KEY` | Azure OpenAI API key |
| `AZURE_OPENAI_API_VERSION` | Azure OpenAI API version |
| `AZURE_OPENAI_DEPLOYMENT_NAME` | Generic fallback deployment |
| `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME` | Preferred deployment for `OpenAIChatClient` |
| `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME` | Preferred deployment for `OpenAIChatCompletionClient` |
| `AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME` | Preferred deployment for `OpenAIEmbeddingClient` |
Deployment lookup order:
- `OpenAIChatClient`: `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME` -> `AZURE_OPENAI_DEPLOYMENT_NAME`
- `OpenAIChatCompletionClient`: `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME` -> `AZURE_OPENAI_DEPLOYMENT_NAME`
- `OpenAIEmbeddingClient`: `AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME` -> `AZURE_OPENAI_DEPLOYMENT_NAME`
When both OpenAI and Azure environment variables are present, the generic clients prefer OpenAI
when `OPENAI_API_KEY` is configured. To use Azure explicitly, pass `azure_endpoint` or
`credential`.
## OpenAI example
```python
from agent_framework.openai import OpenAIChatClient
client = OpenAIChatClient(model_id="gpt-4o")
client = OpenAIChatClient(model="gpt-4.1")
```
## Azure OpenAI example
```python
from azure.identity.aio import AzureCliCredential
from agent_framework.openai import OpenAIChatClient
client = OpenAIChatClient(
model="my-responses-deployment",
azure_endpoint="https://my-resource.openai.azure.com",
credential=AzureCliCredential(),
)
```
## ChatClient vs ChatCompletionClient
Use `OpenAIChatClient` when you want the Responses API as your default chat surface.
Use `OpenAIChatCompletionClient` when you specifically need the Chat Completions API:
```python
from agent_framework.openai import OpenAIChatCompletionClient
client = OpenAIChatCompletionClient(model="gpt-4o-mini")
```
python/packages/openai/agent_framework_openai/_chat_client.py
+236 -164
View File
@@ -14,7 +14,6 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from copy import copy
from datetime import datetime, timezone
from itertools import chain
from typing import (
@@ -28,12 +27,11 @@ from typing import (
cast,
overload,
)
from urllib.parse import urljoin, urlparse
from agent_framework._clients import BaseChatClient
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from agent_framework._settings import SecretString
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._tools import (
SHELL_TOOL_KIND_VALUE,
FunctionInvocationConfiguration,
@@ -87,8 +85,7 @@ from pydantic import BaseModel
from ._exceptions import OpenAIContentFilterException
from ._shared import (
DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION,
get_api_key,
AzureTokenProvider,
load_openai_service_settings,
maybe_append_azure_endpoint_guidance,
)
@@ -107,14 +104,15 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
)
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
logger = logging.getLogger("agent_framework.openai")
DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION = "preview"
OPENAI_SHELL_ENVIRONMENT_KEY = "openai.responses.shell.environment"
OPENAI_SHELL_OUTPUT_TYPE_KEY = "openai.responses.shell.output_type"
OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY = "openai.responses.local_shell.call_item_id"
@@ -139,7 +137,7 @@ class ReasoningOptions(TypedDict, total=False):
See: https://platform.openai.com/docs/guides/reasoning
"""
effort: Literal["low", "medium", "high"]
effort: Literal["none", "low", "medium", "high", "xhigh"]
"""The effort level for reasoning. Higher effort means more reasoning tokens."""
summary: Literal["auto", "concise", "detailed"]
@@ -272,8 +270,8 @@ class RawOpenAIChatClient( # type: ignore[misc]
@overload
def __init__(
self,
*,
model: str | None = None,
*,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
@@ -282,31 +280,77 @@ class RawOpenAIChatClient( # type: ignore[misc]
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None: ...
) -> None:
"""Initialize a raw OpenAI Chat client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_RESPONSES_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
*,
model: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
*,
azure_endpoint: str,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None: ...
) -> None:
"""Initialize a raw OpenAI Chat client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the Responses default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted,
but ``credential`` is the preferred Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
def __init__(
self,
*,
model: str | None = None,
*,
model_id: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str | None = None,
@@ -318,29 +362,53 @@ class RawOpenAIChatClient( # type: ignore[misc]
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a raw OpenAI Responses client.
"""Initialize a raw OpenAI Chat client.
Keyword Args:
model: OpenAI model name.
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_RESPONSES_MODEL`` and then ``OPENAI_MODEL`` for OpenAI,
or ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME`` for Azure.
model_id: Deprecated alias for ``model``.
api_key: OpenAI API key, SecretString, or callable returning a key.
org_id: OpenAI organization ID.
base_url: Custom API base URL.
azure_endpoint: Azure OpenAI endpoint. When provided, the client uses
``AsyncAzureOpenAI`` instead of ``AsyncOpenAI``. The value should be the
resource endpoint and should not end with ``/openai/v1``. For Azure OpenAI
key auth, either pass the resource endpoint without that suffix to
``azure_endpoint`` or pass the full ``.../openai/v1`` URL to ``base_url``.
Can also be set via ``AZURE_OPENAI_ENDPOINT`` when no ``OPENAI_BASE_URL``
is configured.
api_version: Azure OpenAI API version. Can also be set via
``AZURE_OPENAI_API_VERSION``.
api_key: API key override. For OpenAI this maps to ``OPENAI_API_KEY``.
For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI and resolved from
``OPENAI_ORG_ID`` when not provided.
base_url: Base URL override. For OpenAI this maps to ``OPENAI_BASE_URL``.
For Azure this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use once Azure routing is selected. When
not provided explicitly, Azure routing falls back to
``AZURE_OPENAI_API_VERSION`` and then the Responses default.
default_headers: Additional HTTP headers.
async_client: Pre-configured AsyncOpenAI client (skips client creation).
instruction_role: Role for instruction messages (e.g. ``"system"``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
kwargs: Additional keyword arguments forwarded to ``BaseChatClient``.
Notes:
Environment resolution and routing precedence are:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_RESPONSES_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing
reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
"""
if model_id is not None and model is None:
import warnings
@@ -348,98 +416,39 @@ class RawOpenAIChatClient( # type: ignore[misc]
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
model = model_id
openai_settings: dict[str, Any] = {}
use_azure_client = isinstance(async_client, AsyncAzureOpenAI)
if not async_client:
resolved_settings, use_azure_client = load_openai_service_settings(
model=model,
api_key=api_key,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
azure_model_env_vars=("AZURE_OPENAI_DEPLOYMENT_NAME",),
default_azure_api_version=DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION,
)
openai_settings = dict(resolved_settings)
settings, client, use_azure_client = load_openai_service_settings(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
endpoint=azure_endpoint,
api_version=api_version,
default_azure_api_version=DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION,
default_headers=default_headers,
client=async_client,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
openai_model_fields=("responses_model", "model"),
azure_deployment_fields=("responses_deployment_name", "deployment_name"),
responses_mode=True,
)
api_key_value = openai_settings.get("api_key")
if not api_key_value:
raise ValueError(
"OpenAI API key is required. Set via the 'api_key' parameter or the "
"'OPENAI_API_KEY' or 'AZURE_OPENAI_API_KEY' environment variables."
)
resolved_model = openai_settings.get("model") or model
if not resolved_model:
raise ValueError(
"OpenAI model is required. Set via the 'model' parameter or the "
"'OPENAI_MODEL' or 'AZURE_OPENAI_DEPLOYMENT_NAME' environment variables."
)
model = resolved_model
resolved_api_key = get_api_key(api_key_value)
# Merge APP_INFO into the headers
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)
client_args: dict[str, Any] = {"api_key": resolved_api_key, "default_headers": merged_headers}
if use_azure_client:
endpoint_value = openai_settings.get("azure_endpoint")
if (
not openai_settings.get("base_url")
and endpoint_value
and (hostname := urlparse(str(endpoint_value)).hostname)
and hostname.endswith(".openai.azure.com")
):
openai_settings["base_url"] = urljoin(str(endpoint_value), "/openai/v1/")
client_args.pop("api_key")
if resolved_api_version := openai_settings.get("api_version"):
client_args["api_version"] = resolved_api_version
if resolved_base_url := openai_settings.get("base_url"):
client_args["base_url"] = resolved_base_url
elif resolved_azure_endpoint := openai_settings.get("azure_endpoint"):
client_args["azure_endpoint"] = resolved_azure_endpoint
if callable(resolved_api_key):
client_args["azure_ad_token_provider"] = resolved_api_key
else:
client_args["api_key"] = resolved_api_key
client_args["azure_deployment"] = resolved_model
async_client = AsyncAzureOpenAI(**client_args)
else:
if resolved_org_id := openai_settings.get("org_id"):
client_args["organization"] = resolved_org_id
if resolved_base_url := openai_settings.get("base_url"):
client_args["base_url"] = resolved_base_url
async_client = AsyncOpenAI(**client_args)
self.client = async_client
self.model: str | None = model.strip() if model else None
self.client = client
self.model: str = settings.get("model") or settings.get("deployment_name") or ""
# Store configuration for serialization
resolved_base_url = openai_settings.get("base_url") or base_url
resolved_azure_endpoint = openai_settings.get("azure_endpoint") or azure_endpoint
resolved_api_version = openai_settings.get("api_version") or api_version
self.org_id = openai_settings.get("org_id") or org_id
self.base_url = str(resolved_base_url) if resolved_base_url else None
self.azure_endpoint = str(resolved_azure_endpoint) if resolved_azure_endpoint else None
self.api_version = str(resolved_api_version) if use_azure_client and resolved_api_version else None
self.org_id = settings.get("org_id")
self.base_url = settings.get("base_url")
self.azure_endpoint = settings.get("endpoint")
self.api_version = settings.get("api_version")
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
if instruction_role is not None:
self.instruction_role = instruction_role
self.instruction_role = instruction_role
if use_azure_client:
self.OTEL_PROVIDER_NAME = "azure.ai.openai" # type: ignore[misc]
@@ -2452,8 +2461,8 @@ class OpenAIChatClient( # type: ignore[misc]
@overload
def __init__(
self,
*,
model: str | None = None,
*,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
@@ -2462,38 +2471,84 @@ class OpenAIChatClient( # type: ignore[misc]
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None: ...
) -> None:
"""Initialize an OpenAI Responses client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_RESPONSES_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
middleware: Optional middleware to apply to the client.
function_invocation_configuration: Optional function invocation configuration override.
"""
...
@overload
def __init__(
self,
*,
model: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str,
*,
azure_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None: ...
) -> None:
"""Initialize an OpenAI Responses client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the Responses default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted, but ``credential`` is the preferred
Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
middleware: Optional middleware to apply to the client.
function_invocation_configuration: Optional function invocation configuration override.
"""
...
def __init__(
self,
*,
model: str | None = None,
*,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str | None = None,
@@ -2503,43 +2558,59 @@ class OpenAIChatClient( # type: ignore[misc]
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an OpenAI Responses client.
Keyword Args:
model: OpenAI model name, see https://platform.openai.com/docs/models.
Can also be set via environment variable OPENAI_MODEL.
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.
azure_endpoint: Azure OpenAI endpoint. When provided, the client uses
``AsyncAzureOpenAI``. The value should be the Azure resource endpoint and
should not end with ``/openai/v1``. For Azure OpenAI key auth, either pass
the resource endpoint without that suffix to ``azure_endpoint`` or pass the
full ``.../openai/v1`` URL to ``base_url`` instead. Can also be discovered
from ``AZURE_OPENAI_ENDPOINT`` when no OpenAI base URL is configured.
api_version: Azure OpenAI API version. Can also be set via
``AZURE_OPENAI_API_VERSION``.
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".
env_file_path: Use the environment settings file as a fallback
to environment variables.
env_file_encoding: The encoding of the environment settings file.
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_RESPONSES_MODEL`` and then ``OPENAI_MODEL`` for OpenAI
routing, or ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME`` for Azure routing.
api_key: API key override. For OpenAI routing this maps to ``OPENAI_API_KEY``.
For Azure routing this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI routing and resolved from
``OPENAI_ORG_ID`` when not provided.
base_url: Base URL override. For OpenAI routing this maps to ``OPENAI_BASE_URL``.
For Azure routing this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure routing
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use once Azure routing is selected. When
not provided explicitly, Azure routing falls back to
``AZURE_OPENAI_API_VERSION`` and then the Responses default.
default_headers: Default HTTP headers that are merged into each request.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role to use for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
middleware: Optional middleware to apply to the client.
function_invocation_configuration: Optional function invocation configuration override.
kwargs: Other keyword parameters.
Notes:
Environment resolution and routing precedence are:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_RESPONSES_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing
reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
Examples:
.. code-block:: python
@@ -2571,6 +2642,7 @@ class OpenAIChatClient( # type: ignore[misc]
super().__init__(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
python/packages/openai/agent_framework_openai/_chat_completion_client.py
+326 -197
View File
@@ -13,16 +13,15 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from copy import copy
from datetime import datetime, timezone
from itertools import chain
from typing import Any, ClassVar, Generic, Literal, cast, overload
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload
from agent_framework._clients import BaseChatClient
from agent_framework._docstrings import apply_layered_docstring
from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from agent_framework._settings import SecretString
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
@@ -59,8 +58,7 @@ from pydantic import BaseModel
from ._exceptions import OpenAIContentFilterException
from ._shared import (
DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION,
get_api_key,
AzureTokenProvider,
load_openai_service_settings,
maybe_append_azure_endpoint_guidance,
)
@@ -78,8 +76,16 @@ if sys.version_info >= (3, 11):
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
logger = logging.getLogger("agent_framework.openai")
DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION = "2024-12-01-preview"
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
@@ -179,8 +185,8 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
@overload
def __init__(
self,
*,
model: str | None = None,
*,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
@@ -189,31 +195,77 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None: ...
) -> None:
"""Initialize a raw OpenAI Chat completion client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
*,
model: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str,
*,
azure_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None: ...
) -> None:
"""Initialize a raw OpenAI Chat completion client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the Chat Completions default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted, but ``credential`` is the preferred
Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
def __init__(
self,
*,
model: str | None = None,
*,
model_id: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str | None = None,
@@ -228,26 +280,50 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
"""Initialize a raw OpenAI Chat completion client.
Keyword Args:
model: OpenAI model name.
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL`` for OpenAI routing,
or ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME`` for Azure routing.
model_id: Deprecated alias for ``model``.
api_key: OpenAI API key, SecretString, or callable returning a key.
org_id: OpenAI organization ID.
base_url: Custom API base URL.
azure_endpoint: Azure OpenAI endpoint. When provided, the client uses
``AsyncAzureOpenAI`` instead of ``AsyncOpenAI``. The value should be the
resource endpoint and should not end with ``/openai/v1``. For Azure OpenAI
key auth, either pass the resource endpoint without that suffix to
``azure_endpoint`` or pass the full ``.../openai/v1`` URL to ``base_url``.
Can also be set via ``AZURE_OPENAI_ENDPOINT`` when no ``OPENAI_BASE_URL``
is configured.
api_version: Azure OpenAI API version. Can also be set via
``AZURE_OPENAI_API_VERSION``.
api_key: API key override. For OpenAI routing this maps to ``OPENAI_API_KEY``.
For Azure routing this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI routing and resolved from
``OPENAI_ORG_ID`` when not provided.
base_url: Base URL override. For OpenAI routing this maps to ``OPENAI_BASE_URL``.
For Azure routing this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure routing
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use once Azure routing is selected. When
not provided explicitly, Azure routing falls back to
``AZURE_OPENAI_API_VERSION`` and then the Chat Completions default.
default_headers: Additional HTTP headers.
async_client: Pre-configured AsyncOpenAI client (skips client creation).
instruction_role: Role for instruction messages (e.g. ``"system"``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
kwargs: Additional keyword arguments forwarded to ``BaseChatClient``.
Notes:
Environment resolution and routing precedence are:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_CHAT_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing
reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
"""
if model_id is not None and model is None:
import warnings
@@ -255,89 +331,38 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
model = model_id
openai_settings: dict[str, Any] = {}
use_azure_client = isinstance(async_client, AsyncAzureOpenAI)
if not async_client:
resolved_settings, use_azure_client = load_openai_service_settings(
model=model,
api_key=api_key,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
azure_model_env_vars=("AZURE_OPENAI_DEPLOYMENT_NAME",),
default_azure_api_version=DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION,
)
openai_settings = dict(resolved_settings)
settings, client, use_azure_client = load_openai_service_settings(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
endpoint=azure_endpoint,
api_version=api_version,
default_azure_api_version=DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION,
default_headers=default_headers,
client=async_client,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
openai_model_fields=("chat_model", "model"),
azure_deployment_fields=("chat_deployment_name", "deployment_name"),
)
api_key_value = openai_settings.get("api_key")
if not api_key_value:
raise ValueError(
"OpenAI API key is required. Set via the 'api_key' parameter or the "
"'OPENAI_API_KEY' or 'AZURE_OPENAI_API_KEY' environment variables."
)
resolved_model = openai_settings.get("model") or model
if not resolved_model:
raise ValueError(
"OpenAI model is required. Set via the 'model' parameter or the "
"'OPENAI_MODEL' or 'AZURE_OPENAI_DEPLOYMENT_NAME' environment variables."
)
model = resolved_model
resolved_api_key = get_api_key(api_key_value)
# Merge APP_INFO into the headers
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)
client_args: dict[str, Any] = {"api_key": resolved_api_key, "default_headers": merged_headers}
if use_azure_client:
client_args.pop("api_key")
if resolved_api_version := openai_settings.get("api_version"):
client_args["api_version"] = resolved_api_version
if resolved_base_url := openai_settings.get("base_url"):
client_args["base_url"] = resolved_base_url
elif resolved_azure_endpoint := openai_settings.get("azure_endpoint"):
client_args["azure_endpoint"] = resolved_azure_endpoint
if callable(resolved_api_key):
client_args["azure_ad_token_provider"] = resolved_api_key
else:
client_args["api_key"] = resolved_api_key
client_args["azure_deployment"] = resolved_model
async_client = AsyncAzureOpenAI(**client_args)
else:
if resolved_org_id := openai_settings.get("org_id"):
client_args["organization"] = resolved_org_id
if resolved_base_url := openai_settings.get("base_url"):
client_args["base_url"] = resolved_base_url
async_client = AsyncOpenAI(**client_args)
self.client = async_client
self.model: str | None = model.strip() if model else None
self.client = client
self.model: str = settings.get("model") or settings.get("deployment_name") or ""
# Store configuration for serialization
resolved_base_url = openai_settings.get("base_url") or base_url
resolved_azure_endpoint = openai_settings.get("azure_endpoint") or azure_endpoint
resolved_api_version = openai_settings.get("api_version") or api_version
self.org_id = openai_settings.get("org_id") or org_id
self.base_url = str(resolved_base_url) if resolved_base_url else None
self.azure_endpoint = str(resolved_azure_endpoint) if resolved_azure_endpoint else None
self.api_version = str(resolved_api_version) if use_azure_client and resolved_api_version else None
self.org_id = settings.get("org_id")
self.base_url = settings.get("base_url")
self.azure_endpoint = settings.get("endpoint")
self.api_version = settings.get("api_version")
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
if instruction_role is not None:
self.instruction_role = instruction_role
self.instruction_role = instruction_role
if use_azure_client:
self.OTEL_PROVIDER_NAME = "azure.ai.openai" # type: ignore[misc]
@@ -977,6 +1002,202 @@ class OpenAIChatCompletionClient( # type: ignore[misc]
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@overload
def __init__(
self,
model: 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,
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None:
"""Initialize an OpenAI Chat completion client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
instruction_role: Role for instruction messages (for example ``"system"``).
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
middleware: Optional sequence of ChatAndFunctionMiddlewareTypes to apply to requests.
function_invocation_configuration: Optional configuration for function invocation support.
"""
...
@overload
def __init__(
self,
model: str | None = None,
*,
azure_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
instruction_role: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
) -> None:
"""Initialize an OpenAI Chat completion client.
Keyword Args:
model: Model identifier to use for the request. When not provided, the constructor
reads ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the Chat Completions default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted, but ``credential`` is the preferred
Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role for instruction messages (for example ``"system"``).
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
middleware: Optional sequence of ChatAndFunctionMiddlewareTypes to apply to requests.
function_invocation_configuration: Optional configuration for function invocation support.
"""
...
def __init__(
self,
model: str | None = None,
*,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | 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,
azure_endpoint: str | None = None,
api_version: 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: Model identifier to use for the request. When not provided, the constructor
reads ``OPENAI_CHAT_MODEL`` and then ``OPENAI_MODEL`` for OpenAI routing,
or ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME`` for Azure routing.
api_key: API key override. For OpenAI routing this maps to ``OPENAI_API_KEY``.
For Azure routing this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI routing and resolved from
``OPENAI_ORG_ID`` when not provided.
default_headers: Default HTTP headers that are merged into each request.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI and bypasses env lookup.
instruction_role: Role to use for instruction messages (for example ``"system"``).
base_url: Base URL override. For OpenAI routing this maps to ``OPENAI_BASE_URL``.
For Azure routing this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure routing
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use once Azure routing is selected. When
not provided explicitly, Azure routing falls back to
``AZURE_OPENAI_API_VERSION`` and then the Chat Completions default.
middleware: Optional sequence of ChatAndFunctionMiddlewareTypes to apply to requests.
function_invocation_configuration: Optional configuration for function invocation support.
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
Notes:
Environment resolution and routing precedence are:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI routing reads ``OPENAI_API_KEY``, ``OPENAI_CHAT_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure routing
reads ``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_CHAT_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIChatCompletionClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_MODEL=<model name>
client = OpenAIChatCompletionClient()
# Or passing parameters directly
client = OpenAIChatCompletionClient(model="<model name>", api_key="sk-...")
# Or loading from a .env file
client = OpenAIChatCompletionClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.openai import OpenAIChatCompletionOptions
class MyOptions(OpenAIChatCompletionOptions, total=False):
my_custom_option: str
client: OpenAIChatCompletionClient[MyOptions] = OpenAIChatCompletionClient(model="<model name>")
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
super().__init__(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
api_version=api_version,
default_headers=default_headers,
async_client=async_client,
instruction_role=instruction_role,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
@overload
def get_response(
self,
@@ -1045,98 +1266,6 @@ class OpenAIChatCompletionClient( # type: ignore[misc]
**kwargs,
)
def __init__(
self,
*,
model: 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,
azure_endpoint: str | None = None,
api_version: 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: OpenAI model name, see https://platform.openai.com/docs/models.
Can also be set via environment variable OPENAI_MODEL.
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.
azure_endpoint: Azure OpenAI endpoint. When provided, the client uses
``AsyncAzureOpenAI``. The value should be the Azure resource endpoint and
should not end with ``/openai/v1``. For Azure OpenAI key auth, either pass
the resource endpoint without that suffix to ``azure_endpoint`` or pass the
full ``.../openai/v1`` URL to ``base_url`` instead. Can also be discovered
from ``AZURE_OPENAI_ENDPOINT`` when no OpenAI base URL is configured.
api_version: Azure OpenAI API version. Can also be set via
``AZURE_OPENAI_API_VERSION``.
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 OpenAIChatCompletionClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_MODEL=<model name>
client = OpenAIChatCompletionClient()
# Or passing parameters directly
client = OpenAIChatCompletionClient(model="<model name>", api_key="sk-...")
# Or loading from a .env file
client = OpenAIChatCompletionClient(env_file_path="path/to/.env")
# Using custom ChatOptions with type safety:
from typing import TypedDict
from agent_framework.openai import OpenAIChatCompletionOptions
class MyOptions(OpenAIChatCompletionOptions, total=False):
my_custom_option: str
client: OpenAIChatCompletionClient[MyOptions] = OpenAIChatCompletionClient(model="<model name>")
response = await client.get_response("Hello", options={"my_custom_option": "value"})
"""
super().__init__(
model=model,
api_key=api_key,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
api_version=api_version,
default_headers=default_headers,
async_client=async_client,
instruction_role=instruction_role,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
)
def _apply_openai_chat_completion_client_docstrings() -> None:
"""Align OpenAI chat completion client docstrings with the raw implementation."""
python/packages/openai/agent_framework_openai/_embedding_client.py
+310 -108
View File
@@ -6,23 +6,31 @@ import base64
import struct
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from copy import copy
from typing import Any, ClassVar, Generic, Literal, TypedDict
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypedDict, overload
from agent_framework._clients import BaseEmbeddingClient
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from agent_framework._settings import SecretString
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails
from agent_framework.observability import EmbeddingTelemetryLayer
from openai import AsyncOpenAI
from openai import AsyncAzureOpenAI, AsyncOpenAI
from ._shared import OpenAISettings, get_api_key
from ._shared import AzureTokenProvider, load_openai_service_settings
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 TYPE_CHECKING:
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
DEFAULT_AZURE_OPENAI_EMBEDDING_API_VERSION = "2024-10-21"
class OpenAIEmbeddingOptions(EmbeddingGenerationOptions, total=False):
"""OpenAI-specific embedding options.
@@ -61,11 +69,11 @@ class RawOpenAIEmbeddingClient(
INJECTABLE: ClassVar[set[str]] = {"client"}
@overload
def __init__(
self,
*,
model: str | None = None,
model_id: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
base_url: str | None = None,
@@ -73,21 +81,130 @@ class RawOpenAIEmbeddingClient(
async_client: AsyncOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize a raw OpenAI embedding client.
Keyword Args:
model: Embedding model identifier. When not provided, the constructor reads
``OPENAI_EMBEDDING_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
*,
model: str | None = None,
azure_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize a raw OpenAI embedding client.
Keyword Args:
model: Embedding deployment name. When not provided, the constructor reads
``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the embedding default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted, but ``credential`` is the preferred
Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI.
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
def __init__(
self,
*,
model: str | None = None,
model_id: str | None = None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
org_id: str | None = None,
base_url: str | None = None,
azure_endpoint: str | None = None,
api_version: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a raw OpenAI embedding client.
Keyword Args:
model: OpenAI embedding model name.
model: Embedding model or Azure OpenAI deployment name. When not provided, the
constructor reads ``OPENAI_EMBEDDING_MODEL`` and then ``OPENAI_MODEL``
for OpenAI. For Azure it first checks ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME``
and then ``AZURE_OPENAI_DEPLOYMENT_NAME``.
model_id: Deprecated alias for ``model``.
api_key: OpenAI API key, SecretString, or callable returning a key.
org_id: OpenAI organization ID.
base_url: Custom API base URL.
api_key: API key override. For OpenAI this maps to ``OPENAI_API_KEY``.
For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key auth.
A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI and resolved from
``OPENAI_ORG_ID`` when not provided.
base_url: Base URL override. For OpenAI this maps to ``OPENAI_BASE_URL``.
For Azure this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use for Azure requests. When not provided explicitly,
Azure falls back to
``AZURE_OPENAI_API_VERSION`` and then the embedding default.
default_headers: Additional HTTP headers.
async_client: Pre-configured AsyncOpenAI client (skips client creation).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI.
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
kwargs: Additional keyword arguments forwarded to ``BaseEmbeddingClient``.
Notes:
Environment resolution precedence is:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI reads ``OPENAI_API_KEY``, ``OPENAI_EMBEDDING_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure reads
``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
"""
if model_id is not None and model is None:
import warnings
@@ -95,59 +212,40 @@ class RawOpenAIEmbeddingClient(
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
model = model_id
if not async_client:
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
org_id=org_id,
base_url=base_url,
embedding_model=model,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
settings, client, use_azure_client = load_openai_service_settings(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
endpoint=azure_endpoint,
api_version=api_version,
default_azure_api_version=DEFAULT_AZURE_OPENAI_EMBEDDING_API_VERSION,
default_headers=default_headers,
client=async_client,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
openai_model_fields=("embedding_model", "model"),
azure_deployment_fields=("embedding_deployment_name", "deployment_name"),
)
api_key_value = openai_settings.get("api_key")
resolved_model = openai_settings.get("embedding_model") or model
# Only create a client when we have enough configuration.
# Subclasses that manage their own client pass no args here
if api_key_value:
if not resolved_model:
raise ValueError(
"OpenAI embedding model is required. "
"Set via 'model' parameter or 'OPENAI_EMBEDDING_MODEL' environment variable."
)
model = resolved_model
resolved_api_key = get_api_key(api_key_value)
# Merge APP_INFO into the headers
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)
client_args: dict[str, Any] = {"api_key": resolved_api_key, "default_headers": merged_headers}
if resolved_org_id := openai_settings.get("org_id"):
client_args["organization"] = resolved_org_id
if resolved_base_url := openai_settings.get("base_url"):
client_args["base_url"] = resolved_base_url
async_client = AsyncOpenAI(**client_args)
self.client = async_client
self.model: str | None = model.strip() if model else None
self.client = client
resolved_model = settings.get("model") or settings.get("deployment_name")
self.model: str | None = resolved_model.strip() if isinstance(resolved_model, str) and resolved_model else None
# Store configuration for serialization
self.org_id = org_id
self.base_url = str(base_url) if base_url else None
self.org_id = settings.get("org_id")
self.base_url = settings.get("base_url")
self.azure_endpoint = settings.get("endpoint")
self.api_version = settings.get("api_version")
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
if use_azure_client:
self.OTEL_PROVIDER_NAME = "azure.ai.openai" # type: ignore[misc]
super().__init__(**kwargs)
@@ -225,45 +323,11 @@ class OpenAIEmbeddingClient(
RawOpenAIEmbeddingClient[OpenAIEmbeddingOptionsT],
Generic[OpenAIEmbeddingOptionsT],
):
"""OpenAI embedding client with telemetry support.
Keyword Args:
model: The embedding model (e.g. "text-embedding-3-small").
Can also be set via environment variable OPENAI_EMBEDDING_MODEL.
model_id: Deprecated alias for ``model``.
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=text-embedding-3-small
client = OpenAIEmbeddingClient()
# Or passing parameters directly
client = OpenAIEmbeddingClient(
model="text-embedding-3-small",
api_key="sk-...",
)
# Generate embeddings
result = await client.get_embeddings(["Hello, world!"])
print(result[0].vector)
"""
"""OpenAI embedding client with telemetry support."""
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
@overload
def __init__(
self,
*,
@@ -277,27 +341,165 @@ class OpenAIEmbeddingClient(
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OpenAI embedding client."""
"""Initialize an OpenAI embedding client.
Keyword Args:
model: Embedding model identifier. When not provided, the constructor reads
``OPENAI_EMBEDDING_MODEL`` and then ``OPENAI_MODEL``.
api_key: API key. When not provided explicitly, the constructor reads
``OPENAI_API_KEY``. A callable API key is also supported.
org_id: OpenAI organization ID. When not provided explicitly, the constructor reads
``OPENAI_ORG_ID``.
default_headers: Additional HTTP headers.
async_client: Pre-configured OpenAI client.
base_url: Base URL override. When not provided explicitly, the constructor reads
``OPENAI_BASE_URL``.
otel_provider_name: Optional telemetry provider name override.
env_file_path: Optional ``.env`` file that is checked before the process environment
for ``OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
*,
model: str | None = None,
azure_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
api_version: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
base_url: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | 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.
Keyword Args:
model: Embedding deployment name. When not provided, the constructor reads
``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME`` and then
``AZURE_OPENAI_DEPLOYMENT_NAME``.
azure_endpoint: Azure resource endpoint. When not provided explicitly, the constructor
reads ``AZURE_OPENAI_ENDPOINT``.
credential: Azure credential or token provider for Entra auth.
api_version: Azure API version. When not provided explicitly, the constructor reads
``AZURE_OPENAI_API_VERSION`` and then uses the embedding default.
api_key: API key. For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key
auth. A callable token provider is also accepted, but ``credential`` is the preferred
Azure auth surface.
base_url: Base URL override. When not provided explicitly, the constructor reads
``AZURE_OPENAI_BASE_URL``. Use this instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI.
otel_provider_name: Optional telemetry provider name override.
env_file_path: Optional ``.env`` file that is checked before process environment
variables for ``AZURE_OPENAI_*`` values.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
def __init__(
self,
*,
model: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | AsyncOpenAI | None = None,
base_url: str | None = None,
azure_endpoint: str | None = None,
api_version: 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.
Keyword Args:
model: Embedding model or Azure OpenAI deployment name. When not provided, the
constructor reads ``OPENAI_EMBEDDING_MODEL`` and then ``OPENAI_MODEL``
for OpenAI. For Azure it first checks ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME``
and then ``AZURE_OPENAI_DEPLOYMENT_NAME``.
api_key: API key override. For OpenAI this maps to ``OPENAI_API_KEY``.
For Azure this can be used instead of ``AZURE_OPENAI_API_KEY`` for key auth.
A callable token provider is also accepted for backwards compatibility,
but ``credential`` is the preferred Azure auth surface.
credential: Azure credential or token provider for Azure OpenAI auth. Passing this
is an explicit Azure signal, even when ``OPENAI_API_KEY`` is also configured.
Credential objects require the optional ``azure-identity`` package.
org_id: OpenAI organization ID. Used only for OpenAI and resolved from
``OPENAI_ORG_ID`` when not provided.
default_headers: Additional HTTP headers.
async_client: Pre-configured client. Passing ``AsyncAzureOpenAI`` keeps the client on
Azure; passing ``AsyncOpenAI`` keeps the client on OpenAI.
base_url: Base URL override. For OpenAI this maps to ``OPENAI_BASE_URL``.
For Azure this may be used instead of ``azure_endpoint`` when you want
to pass the full ``.../openai/v1`` base URL directly.
azure_endpoint: Azure resource endpoint. When not provided explicitly, Azure
falls back to ``AZURE_OPENAI_ENDPOINT``.
api_version: Azure API version to use for Azure requests. When not provided explicitly,
Azure falls back to
``AZURE_OPENAI_API_VERSION`` and then the embedding default.
otel_provider_name: Override the OpenTelemetry provider name.
env_file_path: Optional ``.env`` file that is checked before process environment
variables. The same file is used for both ``OPENAI_*`` and ``AZURE_OPENAI_*``
lookups.
env_file_encoding: Encoding for the ``.env`` file.
Notes:
Environment resolution precedence is:
1. Explicit Azure inputs (``azure_endpoint`` or ``credential``)
2. Explicit OpenAI API key or ``OPENAI_API_KEY``
3. Azure environment fallback
OpenAI reads ``OPENAI_API_KEY``, ``OPENAI_EMBEDDING_MODEL``,
``OPENAI_MODEL``, ``OPENAI_ORG_ID``, and ``OPENAI_BASE_URL``. Azure reads
``AZURE_OPENAI_ENDPOINT``, ``AZURE_OPENAI_BASE_URL``,
``AZURE_OPENAI_API_KEY``, ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME``,
``AZURE_OPENAI_DEPLOYMENT_NAME``, and ``AZURE_OPENAI_API_VERSION``.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIEmbeddingClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_EMBEDDING_MODEL=text-embedding-3-small
client = OpenAIEmbeddingClient()
# Or passing OpenAI parameters directly
client = OpenAIEmbeddingClient(
model="text-embedding-3-small",
api_key="sk-...",
)
# Or using Azure OpenAI with an Azure credential
client = OpenAIEmbeddingClient(
model="text-embedding-3-small",
azure_endpoint="https://example-resource.openai.azure.com/",
credential=my_azure_credential,
)
"""
super().__init__(
model=model,
api_key=api_key,
credential=credential,
org_id=org_id,
base_url=base_url,
azure_endpoint=azure_endpoint,
api_version=api_version,
default_headers=default_headers,
async_client=async_client,
otel_provider_name=otel_provider_name,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if otel_provider_name is not None:
self.OTEL_PROVIDER_NAME = otel_provider_name # type: ignore[misc]
# Validate that the client was created successfully (from explicit args or env vars)
if self.client is None:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
if not self.model:
raise ValueError(
"OpenAI embedding model is required. "
"Set via 'model' parameter or 'OPENAI_EMBEDDING_MODEL' environment variable."
)
python/packages/openai/agent_framework_openai/_shared.py
+217 -98
View File
@@ -3,19 +3,18 @@
from __future__ import annotations
import logging
import os
import sys
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from copy import copy
from typing import Any, ClassVar, Union, cast
from typing import TYPE_CHECKING, Any, ClassVar, Literal, Union, cast
import openai
from agent_framework._serialization import SerializationMixin
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
from agent_framework._tools import FunctionTool
from dotenv import dotenv_values
from openai import AsyncOpenAI, AsyncStream, _legacy_response # type: ignore
from agent_framework.exceptions import SettingNotFoundError
from openai import AsyncAzureOpenAI, AsyncOpenAI, AsyncStream, _legacy_response # type: ignore
from openai.types import Completion
from openai.types.audio import Transcription
from openai.types.chat import ChatCompletion, ChatCompletionChunk
@@ -24,10 +23,21 @@ from openai.types.responses.response import Response
from openai.types.responses.response_stream_event import ResponseStreamEvent
from packaging.version import parse
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 azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
logger: logging.Logger = logging.getLogger("agent_framework.openai")
DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION = "2024-10-21"
DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION = "preview"
AZURE_OPENAI_TOKEN_SCOPE = "https://cognitiveservices.azure.com/.default" # noqa: S105 # nosec B105
RESPONSE_TYPE = Union[
@@ -43,12 +53,7 @@ RESPONSE_TYPE = Union[
_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
AzureTokenProvider = Callable[[], str | Awaitable[str]]
def _check_openai_version_for_callable_api_key() -> None:
@@ -92,6 +97,10 @@ class OpenAISettings(TypedDict, total=False):
Can be set via environment variable OPENAI_MODEL.
embedding_model: The OpenAI embedding model to use, for example, text-embedding-3-small.
Can be set via environment variable OPENAI_EMBEDDING_MODEL.
chat_model: The OpenAI chat-completions model to prefer before OPENAI_MODEL.
Can be set via environment variable OPENAI_CHAT_MODEL.
responses_model: The OpenAI responses model to prefer before OPENAI_MODEL.
Can be set via environment variable OPENAI_RESPONSES_MODEL.
Examples:
.. code-block:: python
@@ -110,122 +119,232 @@ class OpenAISettings(TypedDict, total=False):
settings = load_settings(OpenAISettings, env_prefix="OPENAI_", env_file_path="path/to/.env")
"""
api_key: SecretString | Callable[[], str | Awaitable[str]] | None
api_key: SecretString | None
base_url: str | None
org_id: str | None
model: str | None
embedding_model: str | None
azure_endpoint: str | None
chat_model: str | None
responses_model: str | None
class AzureOpenAISettings(TypedDict, total=False):
"""Azure OpenAI environment settings."""
endpoint: str | None
base_url: str | None
api_key: SecretString | None
deployment_name: str | None
embedding_deployment_name: str | None
chat_deployment_name: str | None
responses_deployment_name: str | None
api_version: str | None
def _load_dotenv_values(*, env_file_path: str | None, env_file_encoding: str | None) -> dict[str, str]:
"""Load dotenv values for non-standard environment variable aliases."""
if env_file_path is None or not os.path.exists(env_file_path):
return {}
OpenAIModelSettingName = Literal["model", "embedding_model", "chat_model", "responses_model"]
AzureDeploymentSettingName = Literal[
"deployment_name", "embedding_deployment_name", "chat_deployment_name", "responses_deployment_name"
]
raw_dotenv_values = dotenv_values(dotenv_path=env_file_path, encoding=env_file_encoding or "utf-8")
return {key: value for key, value in raw_dotenv_values.items() if value is not None}
OPENAI_MODEL_ENV_VARS: dict[OpenAIModelSettingName, str] = {
"model": "OPENAI_MODEL",
"embedding_model": "OPENAI_EMBEDDING_MODEL",
"chat_model": "OPENAI_CHAT_MODEL",
"responses_model": "OPENAI_RESPONSES_MODEL",
}
AZURE_DEPLOYMENT_ENV_VARS: dict[AzureDeploymentSettingName, str] = {
"deployment_name": "AZURE_OPENAI_DEPLOYMENT_NAME",
"embedding_deployment_name": "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"chat_deployment_name": "AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"responses_deployment_name": "AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME",
}
def _get_setting_from_alias(
name: str,
*,
dotenv_values_by_name: Mapping[str, str],
def _resolve_named_setting(
settings: Mapping[str, Any],
fields: Sequence[OpenAIModelSettingName | AzureDeploymentSettingName],
) -> str | None:
"""Resolve a setting from an explicit env-var alias."""
if dotenv_value := dotenv_values_by_name.get(name):
return dotenv_value
return os.getenv(name)
"""Return the first populated value from ``fields``."""
for field in fields:
value = settings.get(field)
if isinstance(value, str) and value:
return value
return None
def _join_env_names(env_names: Sequence[str]) -> str:
"""Format env var names for user-facing error messages."""
return ", ".join(f"'{env_name}'" for env_name in env_names)
def load_openai_service_settings(
*,
model: str | None,
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None,
credential: AzureCredentialTypes | AzureTokenProvider | None,
org_id: str | None,
base_url: str | None,
azure_endpoint: str | None,
endpoint: str | None,
api_version: str | None,
default_azure_api_version: str,
default_headers: Mapping[str, str] | None = None,
client: AsyncOpenAI | None = None,
env_file_path: str | None,
env_file_encoding: str | None,
azure_model_env_vars: Sequence[str],
default_azure_api_version: str,
) -> tuple[OpenAISettings, bool]:
openai_model_fields: Sequence[OpenAIModelSettingName] = ("model",),
azure_deployment_fields: Sequence[AzureDeploymentSettingName] = ("deployment_name",),
responses_mode: bool = False,
) -> tuple[dict[str, Any], AsyncOpenAI, bool]:
"""Load OpenAI settings, including Azure OpenAI aliases.
The generic OpenAI clients primarily read from ``OPENAI_*`` variables. When an
``AZURE_OPENAI_ENDPOINT`` (or ``AZURE_OPENAI_BASE_URL``) is available and no
explicit OpenAI base URL is configured, this helper switches to Azure-specific
environment variables for endpoint, API key, model deployment, and API version.
The generic OpenAI clients primarily read from ``OPENAI_*`` variables. Azure-specific
environment variables are used only when an explicit Azure signal is present
(``endpoint`` or ``credential``) or when no explicit
OpenAI API key is available.
"""
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
org_id=org_id,
base_url=base_url,
model=model,
azure_endpoint=azure_endpoint,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
# Merge APP_INFO into the headers
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)
dotenv_values_by_name = _load_dotenv_values(
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
resolved_azure_endpoint = azure_endpoint
resolved_azure_base_url: str | None = None
if not openai_settings.get("base_url"):
if resolved_azure_endpoint is None:
resolved_azure_endpoint = _get_setting_from_alias(
"AZURE_OPENAI_ENDPOINT",
dotenv_values_by_name=dotenv_values_by_name,
)
if resolved_azure_endpoint is None:
resolved_azure_base_url = _get_setting_from_alias(
"AZURE_OPENAI_BASE_URL",
dotenv_values_by_name=dotenv_values_by_name,
)
if resolved_azure_base_url is not None:
openai_settings["base_url"] = resolved_azure_base_url
use_azure_client = resolved_azure_endpoint is not None or resolved_azure_base_url is not None
if resolved_azure_endpoint is not None:
openai_settings["azure_endpoint"] = resolved_azure_endpoint
if use_azure_client:
if api_key is None:
resolved_azure_api_key = _get_setting_from_alias(
"AZURE_OPENAI_API_KEY",
dotenv_values_by_name=dotenv_values_by_name,
)
if resolved_azure_api_key is not None:
openai_settings["api_key"] = SecretString(resolved_azure_api_key)
if model is None:
for env_var_name in azure_model_env_vars:
resolved_model = _get_setting_from_alias(
env_var_name,
dotenv_values_by_name=dotenv_values_by_name,
api_key_callable = api_key if callable(api_key) else None
api_key_str = api_key if not callable(api_key) else None
azure_client = isinstance(client, AsyncAzureOpenAI)
use_azure = azure_client or endpoint is not None or credential is not None
checked_openai = False
if not use_azure:
openai_settings_kwargs: dict[str, Any] = {
"api_key": api_key_str,
"org_id": org_id,
"base_url": base_url,
"env_file_path": env_file_path,
"env_file_encoding": env_file_encoding,
}
if model is not None:
openai_settings_kwargs[openai_model_fields[0]] = model
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
**openai_settings_kwargs,
)
if resolved_model := _resolve_named_setting(openai_settings, openai_model_fields):
openai_settings["model"] = resolved_model
if client:
return openai_settings, client, False # type: ignore[return-value]
if openai_settings.get("api_key") is not None or api_key_callable is not None:
resolved_model = _resolve_named_setting(openai_settings, openai_model_fields)
if not resolved_model:
raise SettingNotFoundError(
"Model must be specified via the 'model' parameter or the "
f"{_join_env_names([OPENAI_MODEL_ENV_VARS[field] for field in openai_model_fields])} "
"environment variable."
)
if resolved_model is not None:
openai_settings["model"] = resolved_model
break
if api_version is not None:
openai_settings["api_version"] = api_version
else:
resolved_api_version = _get_setting_from_alias(
"AZURE_OPENAI_API_VERSION",
dotenv_values_by_name=dotenv_values_by_name,
client_args: dict[str, Any] = {
"api_key": api_key_callable
if api_key_callable is not None
else openai_settings["api_key"].get_secret_value(), # type: ignore[reportOptionalMemberAccess, union-attr]
"organization": openai_settings.get("org_id"),
"default_headers": merged_headers,
}
if base_url := openai_settings.get("base_url"):
client_args["base_url"] = base_url
return openai_settings, AsyncOpenAI(**client_args), False # type: ignore[return-value]
checked_openai = True
azure_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
required_fields=None if client else [("base_url", "endpoint")],
api_key=api_key_str,
endpoint=endpoint,
base_url=base_url,
api_version=api_version or default_azure_api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if model is not None:
azure_settings[azure_deployment_fields[0]] = model
client_args = {}
resolved_azure_deployment = _resolve_named_setting(azure_settings, azure_deployment_fields)
if resolved_azure_deployment is None and client:
azure_deployment = getattr(client, "_azure_deployment", None)
if isinstance(azure_deployment, str) and azure_deployment:
resolved_azure_deployment = azure_deployment
if resolved_azure_deployment:
azure_settings["deployment_name"] = resolved_azure_deployment
client_args["azure_deployment"] = resolved_azure_deployment
else:
deployment_env_guidance = _join_env_names([
AZURE_DEPLOYMENT_ENV_VARS[field] for field in azure_deployment_fields
])
has_azure_configuration = (
client is not None
or azure_settings.get("endpoint") is not None
or azure_settings.get("base_url") is not None
)
if checked_openai and not has_azure_configuration:
raise SettingNotFoundError(
"OpenAI credentials are required. Provide the 'api_key' parameter or set 'OPENAI_API_KEY'. "
"To use Azure OpenAI instead, pass 'azure_endpoint' or set 'AZURE_OPENAI_ENDPOINT' or "
"'AZURE_OPENAI_BASE_URL'."
)
openai_settings["api_version"] = resolved_api_version or default_azure_api_version
raise SettingNotFoundError(
"Azure OpenAI client requires a deployment name, which can be provided via the 'model' parameter, "
f"or the {deployment_env_guidance} environment variable."
)
if client:
return azure_settings, client, True # type: ignore[return-value]
client_args["default_headers"] = merged_headers
if endpoint := azure_settings.get("endpoint"):
if responses_mode:
client_args["base_url"] = f"{endpoint.rstrip('/')}/openai/v1/"
else:
client_args["azure_endpoint"] = endpoint
if base_url := azure_settings.get("base_url"):
client_args["base_url"] = base_url
if api_key := azure_settings.get("api_key"):
client_args["api_key"] = api_key.get_secret_value()
if api_key_callable:
client_args["api_key"] = api_key_callable
if api_version := azure_settings.get("api_version"):
client_args["api_version"] = api_version
if credential:
client_args["azure_ad_token_provider"] = _resolve_azure_credential_to_token_provider(credential)
if "api_key" not in client_args and "azure_ad_token_provider" not in client_args:
raise SettingNotFoundError(
"Azure OpenAI client requires either an API key or an Azure AD token provider."
" This can be provided either as a callable api_key or via the credential parameter."
)
return azure_settings, AsyncAzureOpenAI(**client_args), True # type: ignore[return-value]
return openai_settings, use_azure_client
def _resolve_azure_credential_to_token_provider(
credential: AzureCredentialTypes | AzureTokenProvider,
) -> AzureTokenProvider:
"""Resolve an Azure credential or token provider for Azure OpenAI auth."""
if callable(credential):
return credential
try:
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity import get_bearer_token_provider
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
"Azure credential auth requires the 'azure-identity' package. Install it with: pip install azure-identity"
) from exc
if isinstance(credential, AsyncTokenCredential):
return get_async_bearer_token_provider(credential, AZURE_OPENAI_TOKEN_SCOPE)
if isinstance(credential, TokenCredential):
return get_bearer_token_provider(credential, AZURE_OPENAI_TOKEN_SCOPE) # type: ignore[arg-type]
raise ValueError(
"The 'credential' parameter must be an Azure TokenCredential, AsyncTokenCredential, or a "
"callable token provider."
)
def maybe_append_azure_endpoint_guidance(message: str, *, azure_endpoint: str | None) -> str:
python/packages/openai/tests/openai/conftest.py
+13
View File
@@ -43,6 +43,8 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
"OPENAI_ORG_ID",
"OPENAI_MODEL",
"OPENAI_EMBEDDING_MODEL",
"OPENAI_CHAT_MODEL",
"OPENAI_RESPONSES_MODEL",
"OPENAI_TEXT_MODEL_ID",
"OPENAI_TEXT_TO_IMAGE_MODEL_ID",
"OPENAI_AUDIO_TO_TEXT_MODEL_ID",
@@ -53,6 +55,9 @@ def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): #
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_DEPLOYMENT_NAME",
"AZURE_OPENAI_API_VERSION",
],
@@ -97,6 +102,8 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
"OPENAI_ORG_ID",
"OPENAI_MODEL",
"OPENAI_EMBEDDING_MODEL",
"OPENAI_CHAT_MODEL",
"OPENAI_RESPONSES_MODEL",
"OPENAI_TEXT_MODEL_ID",
"OPENAI_TEXT_TO_IMAGE_MODEL_ID",
"OPENAI_AUDIO_TO_TEXT_MODEL_ID",
@@ -107,6 +114,9 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_BASE_URL",
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME",
"AZURE_OPENAI_DEPLOYMENT_NAME",
"AZURE_OPENAI_API_VERSION",
],
@@ -114,6 +124,9 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
env_vars = {
"AZURE_OPENAI_ENDPOINT": "https://test-endpoint.openai.azure.com",
"AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
"AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_responses_deployment",
"AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
"AZURE_OPENAI_DEPLOYMENT_NAME": "test_deployment",
"AZURE_OPENAI_API_KEY": "test_api_key",
"AZURE_OPENAI_API_VERSION": "2024-12-01-preview",
python/packages/openai/tests/openai/test_assistant_provider.py
-62
View File
@@ -1,6 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock
@@ -750,64 +749,3 @@ class TestToolMerging:
# 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_MODEL", "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_MODEL", "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/openai/tests/openai/test_openai_chat_client.py
+175 -173
View File
@@ -28,6 +28,7 @@ from agent_framework._sessions import (
from agent_framework.exceptions import (
ChatClientException,
ChatClientInvalidRequestException,
SettingNotFoundError,
)
from openai import BadRequestError
from openai.types.responses.response_reasoning_item import Summary
@@ -109,6 +110,14 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
assert isinstance(openai_responses_client, SupportsChatGetResponse)
def test_init_prefers_openai_responses_model(monkeypatch, openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_RESPONSES_MODEL", "test_responses_model_id")
openai_responses_client = OpenAIChatClient()
assert openai_responses_client.model == "test_responses_model_id"
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
@@ -143,7 +152,7 @@ def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
@pytest.mark.parametrize("exclude_list", [["OPENAI_MODEL"]], indirect=True)
def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatClient()
@@ -151,7 +160,7 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
model_id = "test_model_id"
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatClient(
model=model_id,
)
@@ -203,34 +212,56 @@ async def test_get_response_with_invalid_input() -> None:
async def test_get_response_with_all_parameters() -> None:
"""Test get_response with all possible parameters to cover parameter handling logic."""
"""Test request preparation with a comprehensive parameter set."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Test with comprehensive parameter set - should fail due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="Test message")],
options={
"include": ["message.output_text.logprobs"],
"instructions": "You are a helpful assistant",
"max_tokens": 100,
"parallel_tool_calls": True,
"model": "gpt-4",
"previous_response_id": "prev-123",
"reasoning": {"chain_of_thought": "enabled"},
"service_tier": "auto",
"response_format": OutputStruct,
"seed": 42,
"store": True,
"temperature": 0.7,
"tool_choice": "auto",
"tools": [get_weather],
"top_p": 0.9,
"user": "test-user",
"truncation": "auto",
"timeout": 30.0,
"additional_properties": {"custom": "value"},
},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="Test message")],
options={
"include": ["message.output_text.logprobs"],
"instructions": "You are a helpful assistant",
"max_tokens": 100,
"parallel_tool_calls": True,
"model": "gpt-4",
"previous_response_id": "prev-123",
"reasoning": {"chain_of_thought": "enabled"},
"service_tier": "auto",
"response_format": OutputStruct,
"seed": 42,
"store": True,
"temperature": 0.7,
"tool_choice": "auto",
"tools": [get_weather],
"top_p": 0.9,
"user": "test-user",
"truncation": "auto",
"timeout": 30.0,
"additional_properties": {"custom": "value"},
},
)
assert run_options["include"] == ["message.output_text.logprobs"]
assert run_options["max_output_tokens"] == 100
assert run_options["parallel_tool_calls"] is True
assert run_options["model"] == "gpt-4"
assert run_options["previous_response_id"] == "prev-123"
assert run_options["reasoning"] == {"chain_of_thought": "enabled"}
assert run_options["service_tier"] == "auto"
assert run_options["text_format"] is OutputStruct
assert run_options["store"] is True
assert run_options["temperature"] == 0.7
assert run_options["tool_choice"] == "auto"
assert run_options["top_p"] == 0.9
assert run_options["user"] == "test-user"
assert run_options["truncation"] == "auto"
assert run_options["timeout"] == 30.0
assert run_options["additional_properties"] == {"custom": "value"}
assert len(run_options["tools"]) == 1
assert run_options["tools"][0]["type"] == "function"
assert run_options["tools"][0]["name"] == "get_weather"
assert run_options["input"][0]["role"] == "system"
assert run_options["input"][0]["content"][0]["text"] == "You are a helpful assistant"
assert run_options["input"][1]["role"] == "user"
assert run_options["input"][1]["content"][0]["text"] == "Test message"
@pytest.mark.asyncio
@@ -248,12 +279,13 @@ async def test_web_search_tool_with_location() -> None:
}
)
# Should raise an authentication error due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="What's the weather?")],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="What's the weather?")],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
assert run_options["tools"] == [web_search_tool]
assert run_options["tool_choice"] == "auto"
async def test_code_interpreter_tool_variations() -> None:
@@ -263,20 +295,22 @@ async def test_code_interpreter_tool_variations() -> None:
# Test code interpreter using static method
code_tool = OpenAIChatClient.get_code_interpreter_tool()
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
assert run_options["tools"] == [code_tool]
# Test code interpreter with files using static method
code_tool_with_files = OpenAIChatClient.get_code_interpreter_tool(file_ids=["file1", "file2"])
with pytest.raises(ChatClientException):
await client.get_response(
messages=[Message(role="user", text="Process these files")],
options={"tools": [code_tool_with_files]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="Process these files")],
options={"tools": [code_tool_with_files]},
)
assert run_options["tools"] == [code_tool_with_files]
async def test_content_filter_exception() -> None:
@@ -300,23 +334,23 @@ async def test_content_filter_exception() -> None:
@pytest.mark.asyncio
async def test_hosted_file_search_tool_validation() -> None:
"""Test get_response HostedFileSearchTool validation."""
"""Test HostedFileSearchTool validation and request preparation."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Test file search tool with vector store IDs
file_search_tool = OpenAIChatClient.get_file_search_tool(vector_store_ids=["vs_123"])
# Test using file search tool - may raise various exceptions depending on API response
with pytest.raises((ValueError, ChatClientInvalidRequestException, ChatClientException)):
await client.get_response(
messages=[Message("user", ["Test"])],
options={"tools": [file_search_tool]},
)
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Test"])],
options={"tools": [file_search_tool]},
)
assert run_options["tools"] == [file_search_tool]
async def test_chat_message_parsing_with_function_calls() -> None:
"""Test get_response message preparation with function call and result content types in conversation flow."""
"""Test message preparation with function call and function result content."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
# Create messages with function call and result content
@@ -335,9 +369,27 @@ async def test_chat_message_parsing_with_function_calls() -> None:
Message(role="tool", contents=[function_result]),
]
# This should exercise the message parsing logic - will fail due to invalid API key
with pytest.raises(ChatClientException):
await client.get_response(messages=messages)
prepared_messages = client._prepare_messages_for_openai(messages)
assert prepared_messages == [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "Call a function"}],
},
{
"call_id": "test-call-id",
"id": "fc_test-fc-id",
"type": "function_call",
"name": "test_function",
"arguments": '{"param": "value"}',
},
{
"call_id": "test-call-id",
"type": "function_call_output",
"output": "Function executed successfully",
},
]
async def test_response_format_parse_path() -> None:
@@ -3043,8 +3095,6 @@ def test_with_callable_api_key() -> None:
"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"),
@@ -3057,7 +3107,6 @@ def test_with_callable_api_key() -> None:
# OpenAIChatOptions - 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
@@ -3113,70 +3162,56 @@ async def test_integration_options(
they don't cause failures. Options marked with needs_validation also
check that the feature actually works correctly.
"""
openai_responses_client = OpenAIChatClient()
client = OpenAIChatClient()
# Need at least 2 iterations for tool_choice tests: one to get function call, one to get final response
openai_responses_client.function_invocation_configuration["max_iterations"] = 2
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
# Prepare test message
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(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.startswith("tool_choice"):
options["tools"] = [get_weather]
# Test streaming mode
response = await client.get_response(stream=True, messages=messages, options=options).get_final_response()
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.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
# 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.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(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.startswith("tool_choice"):
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_stream = openai_responses_client.get_response(
stream=True,
messages=messages,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await openai_responses_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.startswith("tools") or option_name.startswith("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.startswith("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()
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.timeout(300)
@@ -3186,53 +3221,24 @@ async def test_integration_options(
async def test_integration_web_search() -> None:
client = OpenAIChatClient(model="gpt-5")
for streaming in [False, True]:
# Use static method for web search tool
web_search_tool = OpenAIChatClient.get_web_search_tool()
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": [web_search_tool],
},
}
if streaming:
response = await client.get_response(stream=True, **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
web_search_tool_with_location = OpenAIChatClient.get_web_search_tool(
user_location={"country": "US", "city": "Seattle"},
)
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [web_search_tool_with_location],
},
}
if streaming:
response = await client.get_response(stream=True, **content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
# Test that the client will use the web search tool with location
web_search_tool_with_location = OpenAIChatClient.get_web_search_tool(
user_location={"country": "US", "city": "Seattle"},
)
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [web_search_tool_with_location],
},
}
response = await client.get_response(stream=True, **content).get_final_response()
assert response.text is not None
@pytest.mark.skip(
@@ -3351,7 +3357,6 @@ async def test_integration_tool_rich_content_image() -> None:
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
@pytest.mark.timeout(300)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_openai_integration_tests_disabled
@@ -3363,14 +3368,11 @@ async def test_integration_agent_replays_local_tool_history_without_stale_fc_id(
async def search_hotels(city: Annotated[str, "The city to search for hotels in"]) -> str:
return f"The only hotel option in {city} is {hotel_code}."
client = OpenAIChatClient()
# override with model that does not do reasoning by default
client = OpenAIChatClient(model="gpt-5.4")
client.function_invocation_configuration["max_iterations"] = 2
agent = Agent(
client=client,
tools=[search_hotels],
default_options={"store": False},
)
agent = Agent(client=client, tools=[search_hotels], default_options={"store": False})
session = agent.create_session()
first_response = await agent.run(
python/packages/openai/tests/openai/test_openai_chat_client_azure.py
+162 -95
View File
@@ -4,12 +4,16 @@ from __future__ import annotations
import json
import os
from functools import wraps
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import Agent, AgentResponse, ChatResponse, Content, Message, SupportsChatGetResponse, tool
from azure.identity.aio import AzureCliCredential, get_bearer_token_provider
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from pydantic import BaseModel
from pytest import param
@@ -20,11 +24,40 @@ pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == "",
or (
os.getenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", "") == ""
and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or responses deployment provided; skipping integration tests.",
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION") or "preview"
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
@@ -32,18 +65,6 @@ class OutputStruct(BaseModel):
weather: str | None = None
def _create_azure_openai_chat_client(
*,
api_key: Any = None,
) -> OpenAIChatClient:
return OpenAIChatClient(
model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
api_key=api_key or os.environ["AZURE_OPENAI_API_KEY"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
)
async def create_vector_store(client: OpenAIChatClient) -> tuple[str, Content]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
@@ -79,30 +100,117 @@ async def get_weather(location: str) -> str:
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_openai_chat_client()
client = OpenAIChatClient(credential=AzureCliCredential())
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client, SupportsChatGetResponse)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
assert client.api_version == azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"]
assert client.azure_endpoint.startswith(azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"])
def test_init_auto_detects_azure_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_openai_chat_client()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient()
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_api_version_alone_does_not_override_openai_api_key(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient(api_version="2024-10-21")
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
client = OpenAIChatClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.api_version == "preview"
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_responses_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_RESPONSES_MODEL", "test_responses_model")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_RESPONSES_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatClient()
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIChatClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatClient(credential=AzureCliCredential())
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
assert client.api_version is not None
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
@@ -123,8 +231,6 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
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"),
@@ -136,7 +242,6 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
param("tool_choice", "none", True, id="tool_choice_none"),
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"),
param("tools", [get_weather], True, id="tools_function"),
@@ -174,15 +279,14 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
),
],
)
@_with_azure_openai_debug()
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
@@ -233,64 +337,34 @@ async def test_integration_options(
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_web_search() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
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": [OpenAIChatClient.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 isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [OpenAIChatClient.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
response = await client.get_response(
messages=[
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
options={
"tools": [OpenAIChatClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})],
},
stream=True,
).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
file_id, vector_store = await create_vector_store(client)
try:
response = await client.get_response(
@@ -310,11 +384,10 @@ async def test_integration_client_file_search() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search_streaming() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
file_id, vector_store = await create_vector_store(client)
try:
response_stream = client.get_response(
@@ -336,11 +409,10 @@ async def test_integration_client_file_search_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_mcp_tool() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="How to create an Azure storage account using az cli?")],
options={
@@ -361,11 +433,10 @@ async def test_integration_client_agent_hosted_mcp_tool() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_code_interpreter_tool() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="Calculate the sum of numbers from 1 to 10 using Python code.")],
@@ -381,14 +452,13 @@ async def test_integration_client_agent_hosted_code_interpreter_tool() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_existing_session() -> None:
async with AzureCliCredential() as credential:
preserved_session = None
async with Agent(
client=_create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
session = first_agent.create_session()
@@ -403,9 +473,7 @@ async def test_integration_client_agent_existing_session() -> None:
if preserved_session:
async with Agent(
client=_create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
@@ -418,6 +486,7 @@ async def test_integration_client_agent_existing_session() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_tool_rich_content_image() -> None:
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@@ -428,9 +497,7 @@ async def test_azure_openai_chat_client_tool_rich_content_image() -> None:
return Content.from_data(data=image_bytes, media_type="image/jpeg")
async with AzureCliCredential() as credential:
client = _create_azure_openai_chat_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatClient(credential=credential)
client.function_invocation_configuration["max_iterations"] = 2
for streaming in [False, True]:
python/packages/openai/tests/openai/test_openai_chat_completion_client.py
+64 -66
View File
@@ -13,7 +13,7 @@ from agent_framework import (
SupportsChatGetResponse,
tool,
)
from agent_framework.exceptions import ChatClientException
from agent_framework.exceptions import ChatClientException, SettingNotFoundError
from openai import BadRequestError
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_message import ChatCompletionMessage
@@ -37,6 +37,14 @@ def test_init(openai_unit_test_env: dict[str, str]) -> None:
assert isinstance(open_ai_chat_completion, SupportsChatGetResponse)
def test_init_prefers_openai_chat_model(monkeypatch, openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_CHAT_MODEL", "test_chat_model_id")
open_ai_chat_completion = OpenAIChatCompletionClient()
assert open_ai_chat_completion.model == "test_chat_model_id"
def test_init_validation_fail() -> None:
# Test successful initialization
with pytest.raises(ValueError):
@@ -93,7 +101,7 @@ def test_init_base_url_from_settings_env() -> None:
@pytest.mark.parametrize("exclude_list", [["OPENAI_MODEL"]], indirect=True)
def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatCompletionClient()
@@ -101,7 +109,7 @@ def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
model_id = "test_model_id"
with pytest.raises(ValueError):
with pytest.raises(SettingNotFoundError):
OpenAIChatCompletionClient(
model=model_id,
)
@@ -1480,71 +1488,61 @@ async def test_integration_options(
# 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
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(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.startswith("tool_choice"):
options["tools"] = [get_weather]
# Test streaming mode
response = await client.get_response(
messages=messages,
stream=True,
options=options,
).get_final_response()
assert response is not None
assert isinstance(response, ChatResponse)
assert response.messages is not None
if not option_name.startswith("tool_choice") and (
(isinstance(option_value, str) and option_value != "required")
or (isinstance(option_value, dict) and option_value.get("mode") != "required")
):
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.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
# 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.startswith("response_format"):
# Use prompt that works well with structured output
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(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.startswith("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.messages is not None
if not option_name.startswith("tool_choice") and (
(isinstance(option_value, str) and option_value != "required")
or (isinstance(option_value, dict) and option_value.get("mode") != "required")
):
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.startswith("tools") or option_name.startswith("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.startswith("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()
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
python/packages/openai/tests/openai/test_openai_chat_completion_client_azure.py
+129 -57
View File
@@ -3,7 +3,9 @@
from __future__ import annotations
import os
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import (
@@ -16,7 +18,9 @@ from agent_framework import (
SupportsChatGetResponse,
tool,
)
from azure.identity.aio import AzureCliCredential, get_bearer_token_provider
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from agent_framework_openai import OpenAIChatCompletionClient
@@ -25,21 +29,37 @@ pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == "",
or (
os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "") == "" and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or chat deployment provided; skipping integration tests.",
)
def _create_azure_chat_completion_client(
*,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
) -> OpenAIChatCompletionClient:
return OpenAIChatCompletionClient(
model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
api_key=api_key or os.environ["AZURE_OPENAI_API_KEY"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "<unset>")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
@tool(approval_mode="never_require")
@@ -60,9 +80,9 @@ async def get_weather(location: str) -> str:
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_chat_completion_client()
client = OpenAIChatCompletionClient(azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"))
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client, SupportsChatGetResponse)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
@@ -73,18 +93,86 @@ def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) ->
def test_init_auto_detects_azure_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIChatCompletionClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_VERSION", "preview")
client = _create_azure_chat_completion_client()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatCompletionClient()
assert client.model == "gpt-5"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
client = OpenAIChatCompletionClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
client = OpenAIChatCompletionClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert client.api_version == "2024-10-21"
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_chat_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_CHAT_MODEL", "test_chat_model")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatCompletionClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_CHAT_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIChatCompletionClient()
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIChatCompletionClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
@@ -102,11 +190,10 @@ def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_response() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
assert isinstance(client, SupportsChatGetResponse)
messages = [
@@ -134,11 +221,10 @@ async def test_azure_openai_chat_completion_client_response() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_response_tools() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = await client.get_response(
messages=[Message(role="user", text="who are Emily and David?")],
@@ -153,11 +239,10 @@ async def test_azure_openai_chat_completion_client_response_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_streaming() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = client.get_response(
messages=[
@@ -190,11 +275,10 @@ async def test_azure_openai_chat_completion_client_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_streaming_tools() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
)
client = OpenAIChatCompletionClient(credential=credential)
response = client.get_response(
messages=[Message(role="user", text="who are Emily and David?")],
@@ -215,13 +299,12 @@ async def test_azure_openai_chat_completion_client_streaming_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_basic_run() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
) as agent,
):
response = await agent.run("Please respond with exactly: 'This is a response test.'")
@@ -234,20 +317,14 @@ async def test_azure_openai_chat_completion_client_agent_basic_run() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_basic_run_streaming() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
) as agent,
Agent(client=OpenAIChatCompletionClient(credential=credential)) as agent,
):
full_text = ""
async for chunk in agent.run(
"Please respond with exactly: 'This is a streaming response test.'",
stream=True,
):
async for chunk in agent.run("Please respond with exactly: 'This is a streaming response test.'", stream=True):
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
full_text += chunk.text
@@ -258,13 +335,12 @@ async def test_azure_openai_chat_completion_client_agent_basic_run_streaming() -
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_session_persistence() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as agent,
):
@@ -281,14 +357,13 @@ async def test_azure_openai_chat_completion_client_agent_session_persistence() -
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_completion_client_agent_existing_session() -> None:
async with AzureCliCredential() as credential:
preserved_session = None
async with Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
session = first_agent.create_session()
@@ -299,9 +374,7 @@ async def test_azure_openai_chat_completion_client_agent_existing_session() -> N
if preserved_session:
async with Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
second_response = await second_agent.run("What is my name?", session=preserved_session)
@@ -314,13 +387,12 @@ async def test_azure_openai_chat_completion_client_agent_existing_session() -> N
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_chat_completion_client_agent_level_tool_persistence() -> None:
async with (
AzureCliCredential() as credential,
Agent(
client=_create_azure_chat_completion_client(
api_key=get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
),
client=OpenAIChatCompletionClient(credential=credential),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather],
) as agent,
python/packages/openai/tests/openai/test_openai_embedding_client.py
+30 -21
View File
@@ -6,6 +6,7 @@ import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework.exceptions import SettingNotFoundError
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
@@ -32,13 +33,6 @@ def _make_openai_response(
)
@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", "text-embedding-3-small")
# --- OpenAI unit tests ---
@@ -50,24 +44,39 @@ def test_openai_construction_with_explicit_params() -> None:
assert client.model == "text-embedding-3-small"
def test_openai_construction_from_env(openai_unit_test_env: None) -> None:
def test_openai_construction_from_env(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == openai_unit_test_env["OPENAI_EMBEDDING_MODEL"]
def test_with_callable_api_key() -> None:
"""Test OpenAIEmbeddingClient initialization with callable API key."""
async def get_api_key() -> str:
return "test-api-key-123"
client = OpenAIEmbeddingClient(model="text-embedding-3-small", api_key=get_api_key)
assert client.model == "text-embedding-3-small"
assert client.client is not None
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"):
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
def test_openai_construction_without_openai_or_azure_config_raises_clear_error(
openai_unit_test_env: dict[str, str],
) -> None:
with pytest.raises(SettingNotFoundError):
OpenAIEmbeddingClient(model="text-embedding-3-small")
def test_openai_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
with pytest.raises(ValueError, match="embedding model is required"):
OpenAIEmbeddingClient(api_key="test-key")
@pytest.mark.parametrize("exclude_list", [["OPENAI_EMBEDDING_MODEL"]], indirect=True)
def test_openai_construction_falls_back_to_openai_model(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == openai_unit_test_env["OPENAI_MODEL"]
async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
async def test_openai_get_embeddings(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
)
@@ -85,7 +94,7 @@ async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
assert result[0].dimensions == 3
async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
async def test_openai_get_embeddings_usage(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(
embeddings=[[0.1]],
prompt_tokens=10,
@@ -103,7 +112,7 @@ async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
assert result.usage["total_token_count"] == 10
async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None) -> None:
async def test_openai_options_passthrough_dimensions(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
@@ -118,7 +127,7 @@ async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None)
assert result.options is options
async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: None) -> None:
async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: dict[str, str]) -> None:
mock_response = _make_openai_response(embeddings=[[0.1]])
client = OpenAIEmbeddingClient()
client.client = MagicMock()
@@ -132,7 +141,7 @@ async def test_openai_options_passthrough_encoding_format(openai_unit_test_env:
assert call_kwargs["encoding_format"] == "base64"
async def test_openai_base64_decoding(openai_unit_test_env: None) -> None:
async def test_openai_base64_decoding(openai_unit_test_env: dict[str, str]) -> None:
import base64
import struct
@@ -176,7 +185,7 @@ async def test_openai_error_when_no_model_id() -> None:
await client.get_embeddings(["test"])
async def test_openai_empty_values_returns_empty(openai_unit_test_env: None) -> None:
async def test_openai_empty_values_returns_empty(openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
client.client = MagicMock()
client.client.embeddings = MagicMock()
python/packages/openai/tests/openai/test_openai_embedding_client_azure.py
+250
View File
@@ -0,0 +1,250 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
from functools import wraps
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials_async import AsyncTokenCredential
from azure.identity.aio import AzureCliCredential
from openai import AsyncAzureOpenAI
from agent_framework_openai import OpenAIEmbeddingClient, OpenAIEmbeddingOptions
pytestmark = pytest.mark.azure
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.openai.azure.com")
or (
os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", "") == ""
and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No real Azure OpenAI endpoint or embedding deployment provided; skipping integration tests.",
)
def _with_azure_openai_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
model = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") or os.getenv(
"AZURE_OPENAI_DEPLOYMENT_NAME", "<unset>"
)
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "<unset>")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "<unset>")
debug_message = f"Azure OpenAI debug: endpoint={endpoint}, model={model}, api_version={api_version}"
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
def _get_azure_embedding_deployment_name() -> str:
return os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") or os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"]
def _create_azure_embedding_client(
*,
api_key: str | None = None,
credential: AsyncTokenCredential | None = None,
) -> OpenAIEmbeddingClient:
resolved_api_key = (
api_key if api_key is not None else None if credential is not None else os.environ["AZURE_OPENAI_API_KEY"]
)
return OpenAIEmbeddingClient(
model=_get_azure_embedding_deployment_name(),
api_key=resolved_api_key,
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=credential,
)
def test_init_with_azure_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_embedding_client()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
assert client.api_version == azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"]
def test_init_auto_detects_azure_embedding_env(azure_openai_unit_test_env: dict[str, str]) -> None:
client = OpenAIEmbeddingClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_falls_back_to_generic_azure_deployment_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
client = OpenAIEmbeddingClient()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
def test_init_does_not_fall_back_to_openai_embedding_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIEmbeddingClient()
def test_init_does_not_fall_back_to_openai_model_for_azure_env(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("AZURE_OPENAI_DEPLOYMENT_NAME", raising=False)
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
monkeypatch.setenv("OPENAI_MODEL", "gpt-5")
with pytest.raises(SettingNotFoundError, match="Azure OpenAI client requires a deployment name"):
OpenAIEmbeddingClient()
def test_openai_api_key_wins_over_azure_env(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient()
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_api_version_alone_does_not_override_openai_api_key(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient(api_version="2024-10-21")
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
def test_explicit_credential_wins_over_openai_api_key(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
client = OpenAIEmbeddingClient(credential=lambda: "token")
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint == azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"]
def test_init_with_credential_wraps_async_token_credential(
monkeypatch, azure_openai_unit_test_env: dict[str, str]
) -> None:
class TestAsyncTokenCredential(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
credential = TestAsyncTokenCredential()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
client = OpenAIEmbeddingClient(credential=credential)
assert isinstance(client.client, AsyncAzureOpenAI)
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
mock_provider.assert_called_once_with(credential, "https://cognitiveservices.azure.com/.default")
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_VERSION"]], indirect=True)
def test_init_uses_default_azure_api_version(azure_openai_unit_test_env: dict[str, str]) -> None:
client = _create_azure_embedding_client()
assert client.model == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
assert client.api_version == "2024-10-21"
def test_openai_base_url_wins_over_azure_aliases(monkeypatch, azure_openai_unit_test_env: dict[str, str]) -> None:
monkeypatch.setenv("OPENAI_API_KEY", "test-dummy-key")
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
monkeypatch.setenv("OPENAI_BASE_URL", "https://custom-openai-endpoint.com/v1")
client = OpenAIEmbeddingClient()
assert client.model == "text-embedding-3-small"
assert not isinstance(client.client, AsyncAzureOpenAI)
assert client.azure_endpoint is None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
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 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
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings_multiple() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(embedding.vector) for embedding in result]
assert all(dimension == dims[0] for dimension in dims)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_get_embeddings_with_dimensions() -> None:
async with AzureCliCredential() as credential:
client = _create_azure_embedding_client(credential=credential)
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/openai/tests/openai/test_openai_shared.py
+54
View File
@@ -0,0 +1,54 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from agent_framework_openai._shared import AZURE_OPENAI_TOKEN_SCOPE, _resolve_azure_credential_to_token_provider
class _AsyncTokenCredentialStub(AsyncTokenCredential):
async def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
class _TokenCredentialStub(TokenCredential):
def get_token(self, *scopes: str, **kwargs: object):
raise NotImplementedError
def test_resolve_azure_async_credential_wraps_provider() -> None:
credential = _AsyncTokenCredentialStub()
token_provider = MagicMock()
with patch("azure.identity.aio.get_bearer_token_provider", return_value=token_provider) as mock_provider:
resolved = _resolve_azure_credential_to_token_provider(credential)
assert resolved is token_provider
mock_provider.assert_called_once_with(credential, AZURE_OPENAI_TOKEN_SCOPE)
def test_resolve_azure_sync_credential_wraps_provider() -> None:
credential = _TokenCredentialStub()
token_provider = MagicMock()
with patch("azure.identity.get_bearer_token_provider", return_value=token_provider) as mock_provider:
resolved = _resolve_azure_credential_to_token_provider(credential)
assert resolved is token_provider
mock_provider.assert_called_once_with(credential, AZURE_OPENAI_TOKEN_SCOPE)
def test_resolve_azure_callable_token_provider_passthrough() -> None:
token_provider = MagicMock()
assert _resolve_azure_credential_to_token_provider(token_provider) is token_provider
def test_resolve_azure_invalid_credential_raises() -> None:
with pytest.raises(ValueError, match="credential"):
_resolve_azure_credential_to_token_provider(object()) # type: ignore[arg-type]