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/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: