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

* 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 13:33:39 +00:00
committed by GitHub
co-authored by Copilot
parent 3611be82cf
commit cc0cfaaac8
103 changed files with 5451 additions and 4216 deletions
+4 -4
View File
@@ -1,6 +1,6 @@
# Azure AI
AZURE_AI_PROJECT_ENDPOINT=""
AZURE_AI_MODEL_DEPLOYMENT_NAME=""
FOUNDRY_PROJECT_ENDPOINT=""
FOUNDRY_MODEL=""
# Bing connection for web search (optional, used by samples with web search)
BING_CONNECTION_ID=""
# Azure AI Search (optional, used by AzureAISearchContextProvider samples)
@@ -13,8 +13,8 @@ AZURE_SEARCH_KNOWLEDGE_BASE_NAME=""
# (different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs OpenAI endpoint for model calls)
# OpenAI
OPENAI_API_KEY=""
OPENAI_CHAT_MODEL_ID=""
OPENAI_RESPONSES_MODEL_ID=""
OPENAI_CHAT_MODEL=""
OPENAI_RESPONSES_MODEL=""
# Azure OpenAI
AZURE_OPENAI_ENDPOINT=""
AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=""
+2 -2
View File
@@ -108,10 +108,10 @@ Content of `.env` or `openai.env`:
```env
OPENAI_API_KEY=""
OPENAI_CHAT_MODEL_ID="gpt-4o-mini"
OPENAI_MODEL="gpt-4o-mini"
```
You will then configure the ChatClient class with the keyword argument `env_file_path`:
You will then configure the ChatClient class with the keyword argument `env_file_path` (alternatively you can use `load_dotenv` in your code):
```python
from agent_framework.openai import OpenAIChatClient
+15 -5
View File
@@ -47,7 +47,7 @@ Set as environment variables, or create a .env file at your project root:
```bash
OPENAI_API_KEY=sk-...
OPENAI_CHAT_MODEL_ID=...
OPENAI_MODEL=...
...
AZURE_OPENAI_API_KEY=...
AZURE_OPENAI_ENDPOINT=...
@@ -57,15 +57,25 @@ FOUNDRY_PROJECT_ENDPOINT=...
FOUNDRY_MODEL=...
```
For the generic OpenAI clients (`OpenAIChatClient` and `OpenAIChatCompletionClient`), configuration
resolves in this order:
1. Explicit Azure inputs such as `credential` or `azure_endpoint`
2. `OPENAI_API_KEY` / explicit OpenAI API-key parameters
3. Azure environment fallback such as `AZURE_OPENAI_ENDPOINT` and `AZURE_OPENAI_API_KEY`
This means mixed shells default to OpenAI when `OPENAI_API_KEY` is present. To force Azure routing,
pass an explicit Azure input such as `credential=AzureCliCredential()`.
You can also override environment variables by explicitly passing configuration parameters to the chat client constructor:
```python
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.openai import OpenAIChatClient
client = AzureOpenAIChatClient(
client = OpenAIChatClient(
api_key='',
endpoint='',
deployment_name='',
azure_endpoint='',
model='',
api_version='',
)
```
@@ -13,6 +13,7 @@ import json
import logging
import sys
from collections.abc import Mapping, Sequence
from contextlib import contextmanager
from copy import copy
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, cast
from urllib.parse import urljoin, urlparse
@@ -109,6 +110,12 @@ def _apply_azure_defaults(
settings["token_endpoint"] = default_token_endpoint
@contextmanager
def _prefer_single_azure_endpoint_env(*, endpoint: str | None, base_url: str | None) -> Any:
"""Preserve the legacy call shape without mutating process-wide environment state."""
yield
# endregion
@@ -315,6 +322,8 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
"or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
)
endpoint_value = azure_openai_settings.get("endpoint")
client_base_url = azure_openai_settings.get("base_url")
if not async_client:
# Create the Azure OpenAI client directly
merged_headers = dict(copy(default_headers)) if default_headers else {}
@@ -332,9 +341,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
client_endpoint = azure_openai_settings.get("endpoint")
client_base_url = azure_openai_settings.get("base_url")
if not client_endpoint and not client_base_url:
if not endpoint_value and not client_base_url:
raise ValueError("Please provide an endpoint or a base_url")
client_args: dict[str, Any] = {"default_headers": merged_headers}
@@ -346,8 +353,8 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
client_args["api_key"] = api_key_secret.get_secret_value()
if client_base_url:
client_args["base_url"] = str(client_base_url)
if client_endpoint and not client_base_url:
client_args["azure_endpoint"] = str(client_endpoint)
if endpoint_value and not client_base_url:
client_args["azure_endpoint"] = str(endpoint_value)
if responses_deployment_name:
client_args["azure_deployment"] = responses_deployment_name
if "websocket_base_url" in kwargs:
@@ -360,16 +367,19 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = responses_deployment_name
super().__init__(
async_client=async_client,
model=responses_deployment_name,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
with _prefer_single_azure_endpoint_env(endpoint=endpoint_value, base_url=client_base_url):
super().__init__(
async_client=async_client,
model=responses_deployment_name,
azure_endpoint=str(endpoint_value) if endpoint_value else None,
base_url=str(client_base_url) if client_base_url else None,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
@staticmethod
def _create_client_from_project(
@@ -530,6 +540,8 @@ class AzureOpenAIChatClient( # type: ignore[misc]
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
)
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not async_client:
# Create the Azure OpenAI client directly
merged_headers = dict(copy(default_headers)) if default_headers else {}
@@ -547,8 +559,6 @@ class AzureOpenAIChatClient( # type: ignore[misc]
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not endpoint_value and not base_url_value:
raise ValueError("Please provide an endpoint or a base_url")
@@ -573,16 +583,19 @@ class AzureOpenAIChatClient( # type: ignore[misc]
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = chat_deployment_name
super().__init__(
async_client=async_client,
model=chat_deployment_name,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
additional_properties=additional_properties,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
)
with _prefer_single_azure_endpoint_env(endpoint=endpoint_value, base_url=base_url_value):
super().__init__(
async_client=async_client,
model=chat_deployment_name,
azure_endpoint=str(endpoint_value) if endpoint_value else None,
base_url=str(base_url_value) if base_url_value else None,
api_version=azure_openai_settings.get("api_version"),
instruction_role=instruction_role,
default_headers=default_headers,
additional_properties=additional_properties,
middleware=middleware, # type: ignore[arg-type]
function_invocation_configuration=function_invocation_configuration,
)
@override
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
@@ -842,6 +855,8 @@ class AzureOpenAIEmbeddingClient(
"or 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME' environment variable."
)
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not async_client:
# Create the Azure OpenAI client directly
merged_headers = dict(copy(default_headers)) if default_headers else {}
@@ -859,8 +874,6 @@ class AzureOpenAIEmbeddingClient(
if not api_key_secret and not ad_token_provider:
raise ValueError("Please provide either api_key, credential, or a client.")
endpoint_value = azure_openai_settings.get("endpoint")
base_url_value = azure_openai_settings.get("base_url")
if not endpoint_value and not base_url_value:
raise ValueError("Please provide an endpoint or a base_url")
@@ -885,11 +898,15 @@ class AzureOpenAIEmbeddingClient(
self.api_version = azure_openai_settings.get("api_version") or ""
self.deployment_name = embedding_deployment_name
super().__init__(
async_client=async_client,
model=embedding_deployment_name,
default_headers=default_headers,
)
with _prefer_single_azure_endpoint_env(endpoint=endpoint_value, base_url=base_url_value):
super().__init__(
async_client=async_client,
model=embedding_deployment_name,
azure_endpoint=str(endpoint_value) if endpoint_value else None,
base_url=str(base_url_value) if base_url_value else None,
api_version=azure_openai_settings.get("api_version"),
default_headers=default_headers,
)
if otel_provider_name is not None:
self.OTEL_PROVIDER_NAME = otel_provider_name # type: ignore[misc]
@@ -2,6 +2,8 @@
import json
import os
from functools import wraps
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import openai
@@ -33,6 +35,8 @@ from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
from openai.types.chat.chat_completion_message import ChatCompletionMessage
pytestmark = pytest.mark.filterwarnings("ignore:AzureOpenAIChatClient is deprecated\\..*:DeprecationWarning")
# region Service Setup
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
@@ -41,6 +45,32 @@ skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
)
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
def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
# Test successful initialization
azure_chat_client = AzureOpenAIChatClient()
@@ -820,6 +850,7 @@ def get_weather(location: str) -> str:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_response() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
@@ -851,6 +882,7 @@ async def test_azure_openai_chat_client_response() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_response_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
@@ -873,6 +905,7 @@ async def test_azure_openai_chat_client_response_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
@@ -909,6 +942,7 @@ async def test_azure_openai_chat_client_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
@@ -937,6 +971,7 @@ async def test_azure_openai_chat_client_streaming_tools() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_agent_basic_run():
"""Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient."""
async with Agent(
@@ -954,6 +989,7 @@ async def test_azure_openai_chat_client_agent_basic_run():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_agent_basic_run_streaming():
"""Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient."""
async with Agent(
@@ -976,6 +1012,7 @@ async def test_azure_openai_chat_client_agent_basic_run_streaming():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_agent_session_persistence():
"""Test Azure OpenAI chat client agent session persistence across runs with AzureOpenAIChatClient."""
async with Agent(
@@ -1002,6 +1039,7 @@ async def test_azure_openai_chat_client_agent_session_persistence():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_chat_client_agent_existing_session():
"""Test Azure OpenAI chat client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
@@ -1038,6 +1076,7 @@ async def test_azure_openai_chat_client_agent_existing_session():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Chat Client."""
@@ -3,15 +3,20 @@
from __future__ import annotations
import os
from functools import wraps
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework.azure import AzureOpenAIEmbeddingClient
from agent_framework_openai import OpenAIEmbeddingOptions
from agent_framework.openai import OpenAIEmbeddingOptions
from azure.identity.aio import AzureCliCredential
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
from openai.types.create_embedding_response import Usage
pytestmark = pytest.mark.filterwarnings("ignore:AzureOpenAIEmbeddingClient is deprecated\\..*:DeprecationWarning")
def _make_openai_response(
embeddings: list[list[float]],
@@ -106,20 +111,72 @@ def test_azure_otel_provider_name(azure_embedding_unit_test_env: None) -> None:
skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
not os.getenv("AZURE_OPENAI_ENDPOINT")
or (not os.getenv("AZURE_OPENAI_API_KEY") and not os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")),
reason="No Azure OpenAI credentials provided; skipping integration tests.",
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com")
or (
os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", "") == ""
and os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "") == ""
),
reason="No 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_openai_embedding_client(
*,
api_key: str | None = None,
credential: AzureCliCredential | None = None,
) -> AzureOpenAIEmbeddingClient:
resolved_api_key = (
api_key if api_key is not None else None if credential is not None else os.getenv("AZURE_OPENAI_API_KEY")
)
return AzureOpenAIEmbeddingClient(
deployment_name=_get_azure_embedding_deployment_name(),
api_key=resolved_api_key,
endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=credential,
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_azure_openai_get_embeddings() -> None:
"""End-to-end test of Azure OpenAI embedding generation."""
client = AzureOpenAIEmbeddingClient()
async with AzureCliCredential() as credential:
client = _create_azure_openai_embedding_client(credential=credential)
result = await client.get_embeddings(["hello world"])
result = await client.get_embeddings(["hello world"])
assert len(result) == 1
assert isinstance(result[0].vector, list)
@@ -133,11 +190,13 @@ async def test_integration_azure_openai_get_embeddings() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_azure_openai_get_embeddings_multiple() -> None:
"""Test Azure OpenAI embedding generation for multiple inputs."""
client = AzureOpenAIEmbeddingClient()
async with AzureCliCredential() as credential:
client = _create_azure_openai_embedding_client(credential=credential)
result = await client.get_embeddings(["hello", "world", "test"])
result = await client.get_embeddings(["hello", "world", "test"])
assert len(result) == 3
dims = [len(e.vector) for e in result]
@@ -147,12 +206,14 @@ async def test_integration_azure_openai_get_embeddings_multiple() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_openai_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
"""Test Azure OpenAI embedding generation with custom dimensions."""
client = AzureOpenAIEmbeddingClient()
async with AzureCliCredential() as credential:
client = _create_azure_openai_embedding_client(credential=credential)
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
options: OpenAIEmbeddingOptions = {"dimensions": 256}
result = await client.get_embeddings(["hello world"], options=options)
assert len(result) == 1
assert len(result[0].vector) == 256
@@ -3,9 +3,9 @@
import json
import logging
import os
from functools import wraps
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import MagicMock
import pytest
from agent_framework import (
@@ -22,11 +22,40 @@ from azure.identity import AzureCliCredential
from pydantic import BaseModel
from pytest import param
pytestmark = pytest.mark.filterwarnings("ignore:AzureOpenAIResponsesClient is deprecated\\..*:DeprecationWarning")
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests.",
)
def _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", "<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
logger = logging.getLogger(__name__)
@@ -141,119 +170,6 @@ def test_init_with_empty_model_id(azure_openai_unit_test_env: dict[str, str]) ->
AzureOpenAIResponsesClient()
def test_init_with_project_client(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with an existing AIProjectClient."""
from unittest.mock import patch
from openai import AsyncOpenAI
# Create a mock AIProjectClient that returns a mock AsyncOpenAI client
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_client=mock_project_client,
deployment_name="gpt-4o",
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_with_project_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with a project endpoint and credential."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_endpoint="https://test-project.services.ai.azure.com",
deployment_name="gpt-4o",
credential=AzureCliCredential(),
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_create_client_from_project_with_project_client() -> None:
"""Test _create_client_from_project with an existing project client."""
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=mock_project_client,
project_endpoint=None,
credential=None,
)
assert result is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
def test_create_client_from_project_with_endpoint() -> None:
"""Test _create_client_from_project with a project endpoint."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_credential = MagicMock()
with patch("agent_framework_azure_ai._deprecated_azure_openai.AIProjectClient") as MockAIProjectClient:
mock_instance = MockAIProjectClient.return_value
mock_instance.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=mock_credential,
)
assert result is mock_openai_client
MockAIProjectClient.assert_called_once()
mock_instance.get_openai_client.assert_called_once()
def test_create_client_from_project_missing_endpoint() -> None:
"""Test _create_client_from_project raises error when endpoint is missing."""
with pytest.raises(ValueError, match="project endpoint is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint=None,
credential=MagicMock(),
)
def test_create_client_from_project_missing_credential() -> None:
"""Test _create_client_from_project raises error when credential is missing."""
with pytest.raises(ValueError, match="credential is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=None,
)
def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
default_headers = {"X-Unit-Test": "test-guid"}
@@ -285,8 +201,6 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> 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"),
@@ -299,7 +213,6 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
# OpenAIResponsesOptions - just verify they don't fail
param("safety_identifier", "user-hash-abc123", False, id="safety_identifier"),
param("truncation", "auto", False, id="truncation"),
param("top_logprobs", 5, False, id="top_logprobs"),
param("prompt_cache_key", "test-cache-key", False, id="prompt_cache_key"),
param("max_tool_calls", 3, False, id="max_tool_calls"),
# Complex options requiring output validation
@@ -343,6 +256,7 @@ def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
),
],
)
@_with_azure_openai_debug()
async def test_integration_options(
option_name: str,
option_value: Any,
@@ -358,127 +272,84 @@ 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 == "tools" or option_name == "tool_choice":
# Use weather-related prompt for tool tests
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name == "response_format":
# Use prompt that works well with structured output
messages = [
Message(role="user", text="The weather in Seattle is sunny"),
Message(role="user", text="What is the weather in Seattle?"),
]
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name == "tool_choice":
options["tools"] = [get_weather]
# 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.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name == "tools" or option_name == "tool_choice":
# 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 == "response_format":
# Use prompt that works well with structured output
messages = [
Message(role="user", text="The weather in Seattle is sunny"),
Message(role="user", text="What is the weather in Seattle?"),
]
else:
# Generic prompt for simple options
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
# Add tools if testing tool_choice to avoid errors
if option_name == "tool_choice":
options["tools"] = [get_weather]
if streaming:
# Test streaming mode
response_stream = client.get_response(
messages=messages,
stream=True,
options=options,
)
response = await response_stream.get_final_response()
else:
# Test non-streaming mode
response = await client.get_response(
messages=messages,
options=options,
)
assert response is not None
assert isinstance(response, ChatResponse)
assert response.text is not None, f"No text in response for option '{option_name}'"
assert len(response.text) > 0, f"Empty response for option '{option_name}'"
# Validate based on option type
if needs_validation:
if option_name == "tools" or option_name == "tool_choice":
# Should have called the weather function
text = response.text.lower()
assert "sunny" in text or "seattle" in text, f"Tool not invoked for {option_name}"
elif option_name == "response_format":
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
if option_value == OutputStruct:
# Should have structured output
assert response.value is not None, "No structured output"
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
# Runtime JSON schema
assert response.value is None, "No structured output, can't parse any json."
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
assert "seattle" in response_value["location"].lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_web_search() -> None:
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
response = await client.get_response(
messages=[
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
options={
"tools": [
AzureOpenAIResponsesClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})
]
},
stream=True,
).get_final_response()
for streaming in [False, True]:
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {
"tool_choice": "auto",
"tools": [AzureOpenAIResponsesClient.get_web_search_tool()],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response is not None
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
# Test that the client will use the web search tool with location
content = {
"messages": [
Message(
role="user",
text="What is the current weather? Do not ask for my current location.",
)
],
"options": {
"tool_choice": "auto",
"tools": [
AzureOpenAIResponsesClient.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})
],
},
"stream": streaming,
}
if streaming:
response = await client.get_response(**content).get_final_response()
else:
response = await client.get_response(**content)
assert response.text is not None
assert response.text is not None
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search() -> None:
"""Test Azure responses client with file search tool."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
@@ -509,6 +380,7 @@ async def test_integration_client_file_search() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_file_search_streaming() -> None:
"""Test Azure responses client with file search tool and streaming."""
azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
@@ -541,6 +413,7 @@ async def test_integration_client_file_search_streaming() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_mcp_tool() -> None:
"""Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
@@ -566,6 +439,7 @@ async def test_integration_client_agent_hosted_mcp_tool() -> None:
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_hosted_code_interpreter_tool():
"""Test Azure Responses Client agent with code interpreter tool."""
client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
@@ -591,6 +465,7 @@ async def test_integration_client_agent_hosted_code_interpreter_tool():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_integration_client_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
@@ -627,6 +502,7 @@ async def test_integration_client_agent_existing_session():
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_integration_tests_disabled
@_with_azure_openai_debug()
async def test_azure_openai_responses_client_tool_rich_content_image() -> None:
"""Test that Azure OpenAI Responses client can handle tool results containing images."""
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
@@ -660,70 +536,3 @@ async def test_azure_openai_responses_client_tool_rich_content_image() -> None:
assert len(response.text) > 0
# sample_image.jpg contains a photo of a house; the model should mention it.
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
# region Integration with Foundry V2
skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("AZURE_AI_PROJECT_ENDPOINT", "") in ("", "https://test-project.cognitiveservices.azure.com/")
or os.getenv("AZURE_AI_MODEL", "") == "",
reason="No real AZURE_AI_PROJECT_ENDPOINT or AZURE_AI_MODEL provided; skipping integration tests.",
)
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_function_call_roundtrip_preserves_fidelity():
"""Test that function calls roundtrip correctly with full fidelity preserved.
This verifies the changes where:
1. raw_representation is preserved when parsing function calls
2. fc_id and status are included in additional_properties
3. When re-sending messages, the full object fidelity is preserved
"""
call_count = 0
@tool(name="get_weather", approval_mode="never_require")
async def get_weather_tool(location: str) -> str:
"""Get weather for a location."""
nonlocal call_count
call_count += 1
return f"Weather in {location} is sunny, 72F"
client = AzureOpenAIResponsesClient(
project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
deployment_name=os.environ["AZURE_AI_MODEL"],
credential=AzureCliCredential(),
)
async with Agent(
client=client,
name="WeatherAgent",
instructions="You help check weather. Use get_weather when asked about weather.",
tools=[get_weather_tool],
default_options={"store": False}, # Store messages locally to test fidelity across messages
) as agent:
session = agent.create_session()
# First request - should invoke the tool
response1 = await agent.run("What is the weather in Seattle?", session=session)
assert response1 is not None
assert response1.text is not None
assert call_count >= 1
# Verify the response contains expected content
response_text = response1.text.lower()
assert "seattle" in response_text or "sunny" in response_text or "72" in response_text
# Second request - should work correctly with the preserved conversation
response2 = await agent.run("And how about in Portland?", session=session)
assert response2 is not None
assert response2.text is not None
assert call_count >= 2
# endregion
@@ -0,0 +1,131 @@
# Copyright (c) Microsoft. All rights reserved.
import warnings
from unittest.mock import MagicMock
import pytest
from agent_framework import SupportsChatGetResponse
warnings.filterwarnings(
"ignore",
message=r"RawAzureAIClient is deprecated\..*",
category=DeprecationWarning,
)
from agent_framework.azure import AzureOpenAIResponsesClient # noqa: E402
from azure.identity import AzureCliCredential # noqa: E402
pytestmark = pytest.mark.filterwarnings("ignore:AzureOpenAIResponsesClient is deprecated\\..*:DeprecationWarning")
def test_init_with_project_client(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with an existing AIProjectClient."""
from unittest.mock import patch
from openai import AsyncOpenAI
# Create a mock AIProjectClient that returns a mock AsyncOpenAI client
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_client=mock_project_client,
deployment_name="gpt-4o",
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_init_with_project_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test initialization with a project endpoint and credential."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_openai_client.default_headers = {}
with patch(
"agent_framework_azure_ai._deprecated_azure_openai.AzureOpenAIResponsesClient._create_client_from_project",
return_value=mock_openai_client,
):
azure_responses_client = AzureOpenAIResponsesClient(
project_endpoint="https://test-project.services.ai.azure.com",
deployment_name="gpt-4o",
credential=AzureCliCredential(),
)
assert azure_responses_client.model == "gpt-4o"
assert azure_responses_client.client is mock_openai_client
assert isinstance(azure_responses_client, SupportsChatGetResponse)
def test_create_client_from_project_with_project_client() -> None:
"""Test _create_client_from_project with an existing project client."""
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=mock_project_client,
project_endpoint=None,
credential=None,
)
assert result is mock_openai_client
mock_project_client.get_openai_client.assert_called_once()
def test_create_client_from_project_with_endpoint() -> None:
"""Test _create_client_from_project with a project endpoint."""
from unittest.mock import patch
from openai import AsyncOpenAI
mock_openai_client = MagicMock(spec=AsyncOpenAI)
mock_credential = MagicMock()
with patch("agent_framework_azure_ai._deprecated_azure_openai.AIProjectClient") as MockAIProjectClient:
mock_instance = MockAIProjectClient.return_value
mock_instance.get_openai_client.return_value = mock_openai_client
result = AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=mock_credential,
)
assert result is mock_openai_client
MockAIProjectClient.assert_called_once()
mock_instance.get_openai_client.assert_called_once()
def test_create_client_from_project_missing_endpoint() -> None:
"""Test _create_client_from_project raises error when endpoint is missing."""
with pytest.raises(ValueError, match="project endpoint is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint=None,
credential=MagicMock(),
)
def test_create_client_from_project_missing_credential() -> None:
"""Test _create_client_from_project raises error when credential is missing."""
with pytest.raises(ValueError, match="credential is required"):
AzureOpenAIResponsesClient._create_client_from_project(
project_client=None,
project_endpoint="https://test-project.services.ai.azure.com",
credential=None,
)
@@ -3,6 +3,7 @@
import json
import os
import sys
import warnings
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated, Any
@@ -41,8 +42,21 @@ from openai.types.responses.response import Response as OpenAIResponse
from pydantic import BaseModel, ConfigDict, Field
from pytest import fixture
from agent_framework_azure_ai import AzureAIClient, AzureAISettings
from agent_framework_azure_ai._shared import from_azure_ai_tools
from agent_framework_azure_ai import AzureAIClient, AzureAISettings # noqa: E402
from agent_framework_azure_ai._shared import from_azure_ai_tools # noqa: E402
warnings.filterwarnings(
"ignore",
message=r"RawAzureAIClient is deprecated\..*",
category=DeprecationWarning,
)
warnings.filterwarnings(
"ignore",
message=r"AzureAIClient is deprecated\..*",
category=DeprecationWarning,
)
pytestmark = pytest.mark.filterwarnings("ignore:AzureAIClient is deprecated\\..*:DeprecationWarning")
@pytest.fixture
+2 -2
View File
@@ -29,8 +29,8 @@ Set as environment variables, or create a .env file at your project root:
```bash
OPENAI_API_KEY=sk-...
OPENAI_CHAT_MODEL_ID=...
OPENAI_RESPONSES_MODEL_ID=...
OPENAI_CHAT_MODEL=...
OPENAI_RESPONSES_MODEL=...
...
AZURE_OPENAI_API_KEY=...
AZURE_OPENAI_ENDPOINT=...
File diff suppressed because one or more lines are too long
+1 -1
View File
@@ -33,7 +33,7 @@ Then edit `.env` and add your API keys:
```bash
# For OpenAI (minimum required)
OPENAI_API_KEY="your-api-key-here"
OPENAI_CHAT_MODEL_ID="gpt-4o-mini"
OPENAI_CHAT_MODEL="gpt-4o-mini"
# Or for Azure OpenAI
AZURE_OPENAI_ENDPOINT="your-endpoint"
@@ -243,7 +243,7 @@ services:
environment:
# OpenAI
- OPENAI_API_KEY=\${OPENAI_API_KEY}
- OPENAI_CHAT_MODEL_ID=\${OPENAI_CHAT_MODEL_ID:-gpt-4o-mini}
- OPENAI_CHAT_MODEL=\${OPENAI_CHAT_MODEL:-gpt-4o-mini}
# Or Azure OpenAI
- AZURE_OPENAI_API_KEY=\${AZURE_OPENAI_API_KEY}
- AZURE_OPENAI_ENDPOINT=\${AZURE_OPENAI_ENDPOINT}
@@ -802,7 +802,7 @@ az acr build --registry myregistry \\
--target-port 8080 \\
--ingress 'external' \\
--registry-server myregistry.azurecr.io \\
--env-vars OPENAI_API_KEY=secretref:openai-key OPENAI_CHAT_MODEL_ID=gpt-4o-mini`}
--env-vars OPENAI_API_KEY=secretref:openai-key OPENAI_CHAT_MODEL=gpt-4o-mini`}
</pre>
</div>
@@ -2,10 +2,9 @@
import importlib.metadata
from ._foundry_agent import FoundryAgent, RawFoundryAgent
from ._foundry_agent_client import RawFoundryAgentChatClient
from ._foundry_chat_client import FoundryChatClient, FoundryChatOptions, RawFoundryChatClient
from ._foundry_memory_provider import FoundryMemoryProvider
from ._agent import FoundryAgent, RawFoundryAgent, RawFoundryAgentChatClient
from ._chat_client import FoundryChatClient, FoundryChatOptions, RawFoundryChatClient
from ._memory_provider import FoundryMemoryProvider
try:
__version__ = importlib.metadata.version(__name__)
@@ -1,30 +1,37 @@
# Copyright (c) Microsoft. All rights reserved.
"""Microsoft Foundry Agent client for connecting to pre-configured agents in Foundry.
"""Microsoft Foundry Agent for connecting to pre-configured agents in Foundry.
This module provides ``RawFoundryAgentClient`` and ``FoundryAgentClient`` for
communicating with PromptAgents and HostedAgents via the Responses API.
This module provides ``RawFoundryAgent`` and ``FoundryAgent`` Agent subclasses
that connect to existing PromptAgents or HostedAgents in Foundry. Use
``FoundryAgent`` for the recommended experience with full middleware and telemetry.
"""
from __future__ import annotations
import logging
import sys
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._settings import load_settings
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer, FunctionTool
from agent_framework._types import Message
from agent_framework.observability import ChatTelemetryLayer
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentMiddlewareLayer,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
Message,
RawAgent,
load_settings,
)
from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
from ._entra_id_authentication import AzureCredentialTypes
logger: logging.Logger = logging.getLogger(__name__)
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -33,22 +40,25 @@ else:
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
from typing_extensions import override # type: ignore[import] # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework import Agent, BaseContextProvider
from agent_framework._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
from agent_framework import (
Agent,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
MiddlewareTypes,
ToolTypes,
)
from agent_framework._tools import ToolTypes
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
AzureTokenProvider = Callable[[], str | Awaitable[str]]
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundryAgentSettings(TypedDict, total=False):
@@ -203,8 +213,6 @@ class RawFoundryAgentChatClient( # type: ignore[misc]
**kwargs: Any,
) -> Agent[FoundryAgentOptionsT]:
"""Create a FoundryAgent that reuses this client's Foundry configuration."""
from ._foundry_agent import FoundryAgent
function_tools = cast(
FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None,
tools,
@@ -359,9 +367,7 @@ class _FoundryAgentChatClient( # type: ignore[misc]
allow_preview: bool | 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:
@@ -393,3 +399,236 @@ class _FoundryAgentChatClient( # type: ignore[misc]
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
class RawFoundryAgent( # type: ignore[misc]
RawAgent[FoundryAgentOptionsT],
):
"""Raw Microsoft Foundry Agent without agent-level middleware or telemetry.
Connects to an existing PromptAgent or HostedAgent in Foundry.
For full middleware and telemetry support, use :class:`FoundryAgent`.
Examples:
.. code-block:: python
from agent_framework.foundry import RawFoundryAgent
from azure.identity import AzureCliCredential
agent = RawFoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
)
result = await agent.run("Hello!")
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
agent_name: The name of the Foundry agent to connect to.
Can also be set via environment variable FOUNDRY_AGENT_NAME.
agent_version: The version of the agent (required for PromptAgents, optional for HostedAgents).
Can also be set via environment variable FOUNDRY_AGENT_VERSION.
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers for injecting dynamic context.
client_type: Custom client class to use (must be a subclass of ``RawFoundryAgentChatClient``).
Defaults to ``_FoundryAgentChatClient`` (full client middleware).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments passed to the Agent base class.
"""
# Create the client
actual_client_type = client_type or _FoundryAgentChatClient
if not issubclass(actual_client_type, RawFoundryAgentChatClient):
raise TypeError(
f"client_type must be a subclass of RawFoundryAgentChatClient, got {actual_client_type.__name__}"
)
client = actual_client_type(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
super().__init__(
client=client, # type: ignore[arg-type]
tools=tools, # type: ignore[arg-type]
context_providers=context_providers,
**kwargs,
)
async def configure_azure_monitor(
self,
enable_sensitive_data: bool = False,
**kwargs: Any,
) -> None:
"""Setup observability with Azure Monitor (Microsoft Foundry integration).
This method configures Azure Monitor for telemetry collection using the
connection string from the Foundry project client (accessed via the internal client).
Args:
enable_sensitive_data: Enable sensitive data logging (prompts, responses).
Should only be enabled in development/test environments. Default is False.
**kwargs: Additional arguments passed to configure_azure_monitor().
Raises:
ImportError: If azure-monitor-opentelemetry-exporter is not installed.
"""
from azure.core.exceptions import ResourceNotFoundError
client = self.client
if not isinstance(client, RawFoundryAgentChatClient):
raise TypeError("configure_azure_monitor requires a RawFoundryAgentChatClient-based client.")
try:
conn_string = await client.project_client.telemetry.get_application_insights_connection_string()
except ResourceNotFoundError:
logger.warning(
"No Application Insights connection string found for the Foundry project. "
"Please ensure Application Insights is configured in your project, "
"or call configure_otel_providers() manually with custom exporters."
)
return
try:
from azure.monitor.opentelemetry import configure_azure_monitor # type: ignore[import]
except ImportError as exc:
raise ImportError(
"azure-monitor-opentelemetry is required for Azure Monitor integration. "
"Install it with: pip install azure-monitor-opentelemetry"
) from exc
from agent_framework.observability import create_metric_views, create_resource, enable_instrumentation
if "resource" not in kwargs:
kwargs["resource"] = create_resource()
configure_azure_monitor(
connection_string=conn_string,
views=create_metric_views(),
**kwargs,
)
enable_instrumentation(enable_sensitive_data=enable_sensitive_data)
class FoundryAgent( # type: ignore[misc]
AgentMiddlewareLayer,
AgentTelemetryLayer,
RawFoundryAgent[FoundryAgentOptionsT],
):
"""Microsoft Foundry Agent with full middleware and telemetry support.
Connects to an existing PromptAgent or HostedAgent in Foundry.
This is the recommended class for production use.
Examples:
.. code-block:: python
from agent_framework.foundry import FoundryAgent
from azure.identity import AzureCliCredential
# Connect to a PromptAgent
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
tools=[my_function_tool],
)
result = await agent.run("Hello!")
# Connect to a HostedAgent (no version needed)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-hosted-agent",
credential=AzureCliCredential(),
)
# Custom client (e.g., raw client without client middleware)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-agent",
credential=AzureCliCredential(),
client_type=RawFoundryAgentChatClient,
)
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent with full middleware and telemetry.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
agent_name: The name of the Foundry agent to connect to.
agent_version: The version of the agent (for PromptAgents).
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers.
middleware: Optional agent-level middleware.
client_type: Custom client class (must subclass ``RawFoundryAgentChatClient``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments.
"""
super().__init__(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
tools=tools,
context_providers=context_providers,
middleware=middleware,
client_type=client_type,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
**kwargs,
)
@@ -4,14 +4,17 @@ from __future__ import annotations
import logging
import sys
from collections.abc import Sequence
from collections.abc import Awaitable, Callable, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._settings import load_settings
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
from agent_framework._types import Content
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatMiddlewareLayer,
Content,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
load_settings,
)
from agent_framework.observability import ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
@@ -25,9 +28,8 @@ from azure.ai.projects.models import (
)
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.ai.projects.models import MCPTool as FoundryMCPTool
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import resolve_file_ids
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -43,15 +45,13 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
)
from agent_framework import ChatAndFunctionMiddlewareTypes
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
AzureTokenProvider = Callable[[], str | Awaitable[str]]
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundrySettings(TypedDict, total=False):
"""Settings for Microsoft FoundryChatClient resolved from args and environment.
@@ -67,6 +67,33 @@ class FoundrySettings(TypedDict, total=False):
project_endpoint: str | None
def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
"""Resolve file IDs from strings or hosted-file Content objects."""
if not file_ids:
return None
resolved: list[str] = []
for item in file_ids:
if isinstance(item, str):
if not item:
raise ValueError("file_ids must not contain empty strings.")
resolved.append(item)
elif isinstance(item, Content):
if item.type != "hosted_file":
raise ValueError(
f"Unsupported Content type {item.type!r} for code interpreter file_ids. "
"Only Content.from_hosted_file() is supported."
)
if item.file_id is None:
raise ValueError(
"Content.from_hosted_file() item is missing a file_id. "
"Ensure the Content object has a valid file_id before using it in file_ids."
)
resolved.append(item.file_id)
return resolved if resolved else None
FoundryChatOptionsT = TypeVar(
"FoundryChatOptionsT",
bound=TypedDict, # type: ignore[valid-type]
@@ -492,9 +519,7 @@ class FoundryChatClient( # type: ignore[misc]
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: 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:
@@ -1,67 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from typing import Union
from agent_framework.exceptions import ChatClientInvalidAuthException
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
logger: logging.Logger = logging.getLogger(__name__)
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]
"""A callable that returns a bearer token string, either synchronously or asynchronously."""
AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
"""Union of Azure credential types.
Accepts:
- ``TokenCredential`` — synchronous Azure credential (e.g. ``DefaultAzureCredential()``)
- ``AsyncTokenCredential`` — asynchronous Azure credential (e.g. ``azure.identity.aio.DefaultAzureCredential()``)
"""
def resolve_credential_to_token_provider(
credential: AzureCredentialTypes | AzureTokenProvider,
token_endpoint: str | None,
) -> AzureTokenProvider:
"""Convert an Azure credential or token provider into an ``ad_token_provider`` callable.
If the credential is already a callable token provider, it is returned as-is
(``token_endpoint`` is not required in this case).
If it is a ``TokenCredential`` or ``AsyncTokenCredential``, it is wrapped using
``azure.identity.get_bearer_token_provider`` (sync or async variant) which
handles token caching and automatic refresh.
Args:
credential: An Azure credential or token provider callable.
token_endpoint: The token scope/endpoint
(e.g. ``"https://cognitiveservices.azure.com/.default"``).
Required when ``credential`` is a ``TokenCredential`` or ``AsyncTokenCredential``.
Returns:
A callable that returns a bearer token string (sync or async).
Raises:
ServiceInvalidAuthError: If the token endpoint is empty when needed for credential wrapping.
"""
# Already a token provider callable (not a credential object) — use directly
if callable(credential) and not isinstance(credential, (TokenCredential, AsyncTokenCredential)):
return credential
if not token_endpoint:
raise ChatClientInvalidAuthException(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
if isinstance(credential, AsyncTokenCredential):
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
return get_async_bearer_token_provider(credential, token_endpoint)
from azure.identity import get_bearer_token_provider
return get_bearer_token_provider(credential, token_endpoint) # type: ignore[arg-type]
@@ -1,287 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Microsoft Foundry Agent for connecting to pre-configured agents in Foundry.
This module provides ``RawFoundryAgent`` and ``FoundryAgent`` — Agent subclasses
that connect to existing PromptAgents or HostedAgents in Foundry. Use
``FoundryAgent`` for the recommended experience with full middleware and telemetry.
"""
from __future__ import annotations
import logging
import sys
from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING, Any
from agent_framework import (
AgentMiddlewareLayer,
BaseContextProvider,
RawAgent,
)
from agent_framework.observability import AgentTelemetryLayer
from azure.ai.projects.aio import AIProjectClient
from ._entra_id_authentication import AzureCredentialTypes
from ._foundry_agent_client import (
RawFoundryAgentChatClient,
_FoundryAgentChatClient, # pyright: ignore[reportPrivateUsage]
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import MiddlewareTypes
from agent_framework._tools import FunctionTool
from agent_framework_openai._chat_client import OpenAIChatOptions
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
FoundryAgentOptionsT = TypeVar(
"FoundryAgentOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIChatOptions",
covariant=True,
)
class RawFoundryAgent( # type: ignore[misc]
RawAgent[FoundryAgentOptionsT],
):
"""Raw Microsoft Foundry Agent without agent-level middleware or telemetry.
Connects to an existing PromptAgent or HostedAgent in Foundry.
For full middleware and telemetry support, use :class:`FoundryAgent`.
Examples:
.. code-block:: python
from agent_framework.foundry import RawFoundryAgent
from azure.identity import AzureCliCredential
agent = RawFoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
)
result = await agent.run("Hello!")
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
agent_name: The name of the Foundry agent to connect to.
Can also be set via environment variable FOUNDRY_AGENT_NAME.
agent_version: The version of the agent (required for PromptAgents, optional for HostedAgents).
Can also be set via environment variable FOUNDRY_AGENT_VERSION.
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers for injecting dynamic context.
client_type: Custom client class to use (must be a subclass of ``RawFoundryAgentChatClient``).
Defaults to ``_FoundryAgentChatClient`` (full client middleware).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments passed to the Agent base class.
"""
# Create the client
actual_client_type = client_type or _FoundryAgentChatClient
if not issubclass(actual_client_type, RawFoundryAgentChatClient):
raise TypeError(
f"client_type must be a subclass of RawFoundryAgentChatClient, got {actual_client_type.__name__}"
)
client = actual_client_type(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
super().__init__(
client=client, # type: ignore[arg-type]
tools=tools, # type: ignore[arg-type]
context_providers=context_providers,
**kwargs,
)
async def configure_azure_monitor(
self,
enable_sensitive_data: bool = False,
**kwargs: Any,
) -> None:
"""Setup observability with Azure Monitor (Microsoft Foundry integration).
This method configures Azure Monitor for telemetry collection using the
connection string from the Foundry project client (accessed via the internal client).
Args:
enable_sensitive_data: Enable sensitive data logging (prompts, responses).
Should only be enabled in development/test environments. Default is False.
**kwargs: Additional arguments passed to configure_azure_monitor().
Raises:
ImportError: If azure-monitor-opentelemetry-exporter is not installed.
"""
from azure.core.exceptions import ResourceNotFoundError
from ._foundry_agent_client import RawFoundryAgentChatClient
client = self.client
if not isinstance(client, RawFoundryAgentChatClient):
raise TypeError("configure_azure_monitor requires a RawFoundryAgentChatClient-based client.")
try:
conn_string = await client.project_client.telemetry.get_application_insights_connection_string()
except ResourceNotFoundError:
logger.warning(
"No Application Insights connection string found for the Foundry project. "
"Please ensure Application Insights is configured in your project, "
"or call configure_otel_providers() manually with custom exporters."
)
return
try:
from azure.monitor.opentelemetry import configure_azure_monitor # type: ignore[import]
except ImportError as exc:
raise ImportError(
"azure-monitor-opentelemetry is required for Azure Monitor integration. "
"Install it with: pip install azure-monitor-opentelemetry"
) from exc
from agent_framework.observability import create_metric_views, create_resource, enable_instrumentation
if "resource" not in kwargs:
kwargs["resource"] = create_resource()
configure_azure_monitor(
connection_string=conn_string,
views=create_metric_views(),
**kwargs,
)
enable_instrumentation(enable_sensitive_data=enable_sensitive_data)
class FoundryAgent( # type: ignore[misc]
AgentMiddlewareLayer,
AgentTelemetryLayer,
RawFoundryAgent[FoundryAgentOptionsT],
):
"""Microsoft Foundry Agent with full middleware and telemetry support.
Connects to an existing PromptAgent or HostedAgent in Foundry.
This is the recommended class for production use.
Examples:
.. code-block:: python
from agent_framework.foundry import FoundryAgent
from azure.identity import AzureCliCredential
# Connect to a PromptAgent
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
tools=[my_function_tool],
)
result = await agent.run("Hello!")
# Connect to a HostedAgent (no version needed)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-hosted-agent",
credential=AzureCliCredential(),
)
# Custom client (e.g., raw client without client middleware)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-agent",
credential=AzureCliCredential(),
client_type=RawFoundryAgentChatClient,
)
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent with full middleware and telemetry.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
agent_name: The name of the Foundry agent to connect to.
agent_version: The version of the agent (for PromptAgents).
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers.
middleware: Optional agent-level middleware.
client_type: Custom client class (must subclass ``RawFoundryAgentChatClient``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments.
"""
super().__init__(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
tools=tools,
context_providers=context_providers,
middleware=middleware,
client_type=client_type,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
**kwargs,
)
@@ -13,25 +13,38 @@ import sys
from contextlib import AbstractAsyncContextManager
from typing import TYPE_CHECKING, Any, ClassVar
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import load_settings
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentSession,
BaseContextProvider,
Message,
SessionContext,
load_settings,
)
from azure.ai.projects.aio import AIProjectClient
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from openai.types.responses import ResponseInputItemParam
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import FoundryProjectSettings
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
from typing import Self, TypedDict # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
from typing_extensions import Self, TypedDict # pragma: no cover
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
from agent_framework import SupportsAgentRun
logger = logging.getLogger(__name__)
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundryProjectSettings(TypedDict, total=False):
"""Foundry project settings loaded from FOUNDRY_ environment variables."""
project_endpoint: str | None
class FoundryMemoryProvider(BaseContextProvider):
"""Foundry Memory context provider using the new BaseContextProvider hooks pattern.
@@ -1,49 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
import sys
from collections.abc import Sequence
from agent_framework import Content
if sys.version_info >= (3, 11):
from typing import TypedDict # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
logger = logging.getLogger("agent_framework.foundry")
class FoundryProjectSettings(TypedDict, total=False):
"""Foundry project settings loaded from FOUNDRY_ environment variables."""
project_endpoint: str | None
def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
"""Resolve file IDs from strings or hosted-file Content objects."""
if not file_ids:
return None
resolved: list[str] = []
for item in file_ids:
if isinstance(item, str):
if not item:
raise ValueError("file_ids must not contain empty strings.")
resolved.append(item)
elif isinstance(item, Content):
if item.type != "hosted_file":
raise ValueError(
f"Unsupported Content type {item.type!r} for code interpreter file_ids. "
"Only Content.from_hosted_file() is supported."
)
if item.file_id is None:
raise ValueError(
"Content.from_hosted_file() item is missing a file_id. "
"Ensure the Content object has a valid file_id before using it in file_ids."
)
resolved.append(item.file_id)
return resolved if resolved else None
Binary file not shown.

After

Width:  |  Height:  |  Size: 178 KiB

@@ -0,0 +1,413 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
import sys
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentResponse, ChatContext, ChatMiddleware, Message, tool
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import AzureCliCredential
from agent_framework_foundry._agent import (
FoundryAgent,
RawFoundryAgent,
RawFoundryAgentChatClient,
_FoundryAgentChatClient,
)
skip_if_foundry_agent_integration_tests_disabled = pytest.mark.skipif(
os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/")
or os.getenv("FOUNDRY_AGENT_NAME", "") == "",
reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_AGENT_NAME provided; skipping integration tests.",
)
_FOUNDRY_AGENT_ENV_VARS = (
"FOUNDRY_PROJECT_ENDPOINT",
"FOUNDRY_AGENT_NAME",
"FOUNDRY_AGENT_VERSION",
)
@pytest.fixture(autouse=True)
def clear_foundry_agent_settings_env(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> None:
"""Prevent unit tests from inheriting Foundry agent settings from the shell."""
if request.node.get_closest_marker("integration") is not None:
return
for env_var in _FOUNDRY_AGENT_ENV_VARS:
monkeypatch.delenv(env_var, raising=False)
def test_raw_foundry_agent_chat_client_init_requires_agent_name() -> None:
"""Test that agent_name is required."""
with pytest.raises(ValueError, match="Agent name is required"):
RawFoundryAgentChatClient(
project_client=MagicMock(),
)
def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None:
"""Test construction with agent_name and project_client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
assert client.agent_version == "1.0"
def test_raw_foundry_agent_chat_client_get_agent_reference_with_version() -> None:
"""Test agent reference includes version when provided."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="my-agent",
agent_version="2.0",
)
ref = client._get_agent_reference()
assert ref == {"name": "my-agent", "version": "2.0", "type": "agent_reference"}
def test_raw_foundry_agent_chat_client_get_agent_reference_without_version() -> None:
"""Test agent reference omits version for HostedAgents."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="hosted-agent",
)
ref = client._get_agent_reference()
assert ref == {"name": "hosted-agent", "type": "agent_reference"}
assert "version" not in ref
def test_raw_foundry_agent_chat_client_as_agent_preserves_client_type() -> None:
"""Test that as_agent() wraps the client in FoundryAgent using the same client class."""
class CustomClient(RawFoundryAgentChatClient):
pass
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = CustomClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
agent = client.as_agent(instructions="You are helpful.")
assert isinstance(agent, FoundryAgent)
assert agent.name == "test-agent"
assert isinstance(agent.client, CustomClient)
assert agent.client.project_client is mock_project
assert agent.client.agent_name == "test-agent"
assert agent.client.agent_version == "1.0"
named_agent = client.as_agent(name="display-name", instructions="You are helpful.")
assert named_agent.name == "display-name"
assert named_agent.client.agent_name == "test-agent"
async def test_raw_foundry_agent_chat_client_prepare_options_validates_tools() -> None:
"""Test that _prepare_options rejects non-FunctionTool objects."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
with pytest.raises(TypeError, match="Only FunctionTool objects are accepted"):
await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [{"type": "function", "function": {"name": "bad"}}]},
)
async def test_raw_foundry_agent_chat_client_prepare_options_accepts_function_tools() -> None:
"""Test that _prepare_options accepts FunctionTool objects."""
mock_project = MagicMock()
mock_openai = MagicMock()
mock_project.get_openai_client.return_value = mock_openai
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
with patch(
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
new_callable=AsyncMock,
return_value={},
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [my_func]},
)
assert "extra_body" in result
assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None:
"""Test that _check_model_presence does nothing (model is on service)."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
options: dict[str, Any] = {}
client._check_model_presence(options)
assert "model" not in options
def test_foundry_agent_chat_client_init() -> None:
"""Test construction of the full-middleware client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = _FoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
def test_raw_foundry_agent_init_creates_client() -> None:
"""Test that RawFoundryAgent creates a client internally."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_raw_foundry_agent_init_with_custom_client_type() -> None:
"""Test that client_type parameter is respected."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
client_type=RawFoundryAgentChatClient,
)
assert isinstance(agent.client, RawFoundryAgentChatClient)
def test_raw_foundry_agent_init_rejects_invalid_client_type() -> None:
"""Test that invalid client_type raises TypeError."""
with pytest.raises(TypeError, match="must be a subclass of RawFoundryAgentChatClient"):
RawFoundryAgent(
project_client=MagicMock(),
agent_name="test-agent",
client_type=object, # type: ignore[arg-type]
)
def test_raw_foundry_agent_init_with_function_tools() -> None:
"""Test that FunctionTool and callables are accepted."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
tools=[my_func],
)
assert agent.default_options.get("tools") is not None
def test_foundry_agent_init() -> None:
"""Test construction of the full-middleware agent."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_foundry_agent_init_with_middleware() -> None:
"""Test that agent-level middleware is accepted."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
class MyMiddleware(ChatMiddleware):
async def process(self, context: ChatContext) -> None:
pass
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
middleware=[MyMiddleware()],
)
assert agent.client is not None
async def test_foundry_agent_configure_azure_monitor() -> None:
"""Test configure_azure_monitor delegates through the underlying client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
mock_configure = MagicMock()
mock_views = MagicMock(return_value=[])
mock_resource = MagicMock()
mock_enable = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", mock_views),
patch("agent_framework.observability.create_resource", return_value=mock_resource),
patch("agent_framework.observability.enable_instrumentation", mock_enable),
):
await agent.configure_azure_monitor(enable_sensitive_data=True)
mock_project.telemetry.get_application_insights_connection_string.assert_called_once()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
assert call_kwargs["views"] == []
assert call_kwargs["resource"] is mock_resource
mock_enable.assert_called_once_with(enable_sensitive_data=True)
async def test_foundry_agent_configure_azure_monitor_resource_not_found() -> None:
"""Test configure_azure_monitor handles ResourceNotFoundError gracefully."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
side_effect=ResourceNotFoundError("No Application Insights found")
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
await agent.configure_azure_monitor()
mock_project.telemetry.get_application_insights_connection_string.assert_called_once()
async def test_foundry_agent_configure_azure_monitor_import_error() -> None:
"""Test configure_azure_monitor raises ImportError when Azure Monitor is unavailable."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
original_import = __import__
def _import_with_missing_azure_monitor(
name: str,
globals: dict[str, Any] | None = None,
locals: dict[str, Any] | None = None,
fromlist: tuple[str, ...] = (),
level: int = 0,
) -> Any:
if name == "azure.monitor.opentelemetry":
raise ImportError("No module named 'azure.monitor.opentelemetry'")
return original_import(name, globals, locals, fromlist, level)
with (
patch.dict(sys.modules, {"azure.monitor.opentelemetry": None}),
patch("builtins.__import__", side_effect=_import_with_missing_azure_monitor),
pytest.raises(ImportError, match="azure-monitor-opentelemetry is required"),
):
await agent.configure_azure_monitor()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_agent_integration_tests_disabled
async def test_foundry_agent_basic_run() -> None:
"""Smoke-test FoundryAgent against a real configured agent."""
async with FoundryAgent(credential=AzureCliCredential()) as agent:
response = await agent.run("Please respond with exactly: 'This is a response test.'")
assert isinstance(response, AgentResponse)
assert response.text is not None
assert "response test" in response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_agent_integration_tests_disabled
async def test_foundry_agent_custom_client_run() -> None:
"""Smoke-test FoundryAgent against a real configured agent."""
async with FoundryAgent(credential=AzureCliCredential(), client_type=RawFoundryAgentChatClient) as agent:
response = await agent.run("Please respond with exactly: 'This is a response test.'")
assert isinstance(response, AgentResponse)
assert response.text is not None
assert "response test" in response.text.lower()
@@ -0,0 +1,751 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import os
import sys
from functools import wraps
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatResponse, Content, Message, SupportsChatGetResponse, tool
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException
from agent_framework_openai import OpenAIContentFilterException
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import AzureCliCredential
from openai import BadRequestError
from pydantic import BaseModel
from pytest import param
from agent_framework_foundry import FoundryChatClient, RawFoundryChatClient
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str | None = None
@tool(approval_mode="never_require")
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
return f"The current weather in {location} is sunny."
skip_if_foundry_integration_tests_disabled = pytest.mark.skipif(
os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/")
or os.getenv("FOUNDRY_MODEL", "") == "",
reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_MODEL provided; skipping integration tests.",
)
_TEST_FOUNDRY_PROJECT_ENDPOINT = "https://test-project.services.ai.azure.com/"
_TEST_FOUNDRY_MODEL = "test-gpt-4o"
_FOUNDRY_CHAT_ENV_VARS = ("FOUNDRY_PROJECT_ENDPOINT", "FOUNDRY_MODEL")
@pytest.fixture(autouse=True)
def clear_foundry_chat_settings_env(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> None:
"""Prevent unit tests from inheriting Foundry chat settings from the shell."""
if request.node.get_closest_marker("integration") is not None:
return
for env_var in _FOUNDRY_CHAT_ENV_VARS:
monkeypatch.delenv(env_var, raising=False)
def _with_foundry_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:
debug_message = (
"Foundry debug: "
f"project_endpoint={os.getenv('FOUNDRY_PROJECT_ENDPOINT', '<unset>')}, "
f"model={os.getenv('FOUNDRY_MODEL', '<unset>')}"
)
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 _make_mock_openai_client() -> MagicMock:
client = MagicMock()
client.default_headers = {}
client.responses = MagicMock()
client.responses.create = AsyncMock()
client.responses.parse = AsyncMock()
client.files = MagicMock()
client.files.create = AsyncMock()
client.files.delete = AsyncMock()
client.vector_stores = MagicMock()
client.vector_stores.create = AsyncMock()
client.vector_stores.delete = AsyncMock()
client.vector_stores.files = MagicMock()
client.vector_stores.files.create_and_poll = AsyncMock()
return client
async def create_vector_store(client: FoundryChatClient) -> tuple[str, Content]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."),
purpose="user_data",
)
vector_store = await client.client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.client.vector_stores.files.create_and_poll(
vector_store_id=vector_store.id,
file_id=file.id,
poll_interval_ms=1000,
)
if result.last_error is not None:
raise RuntimeError(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, Content.from_hosted_vector_store(vector_store_id=vector_store.id)
async def delete_vector_store(client: FoundryChatClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after tests."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
def test_init() -> None:
mock_openai_client = _make_mock_openai_client()
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=mock_project_client, model=_TEST_FOUNDRY_MODEL)
assert client.model == _TEST_FOUNDRY_MODEL
assert isinstance(client, SupportsChatGetResponse)
assert client.project_client is mock_project_client
def test_init_with_default_header() -> None:
default_headers = {"X-Unit-Test": "test-guid"}
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(
project_client=project_client,
model=_TEST_FOUNDRY_MODEL,
default_headers=default_headers,
)
assert client.model == _TEST_FOUNDRY_MODEL
for key, value in default_headers.items():
assert client.default_headers is not None
assert key in client.default_headers
assert client.default_headers[key] == value
def test_init_with_project_endpoint_creates_project_client() -> None:
credential = MagicMock()
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
with patch("agent_framework_foundry._chat_client.AIProjectClient", return_value=project_client) as factory:
client = FoundryChatClient(
project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT,
model=_TEST_FOUNDRY_MODEL,
credential=credential,
allow_preview=True,
)
assert client.project_client is project_client
assert client.model == _TEST_FOUNDRY_MODEL
assert factory.call_args.kwargs["endpoint"] == _TEST_FOUNDRY_PROJECT_ENDPOINT
assert factory.call_args.kwargs["credential"] is credential
assert factory.call_args.kwargs["allow_preview"] is True
assert factory.call_args.kwargs["user_agent"] == AGENT_FRAMEWORK_USER_AGENT
def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("FOUNDRY_MODEL", raising=False)
mock_openai_client = _make_mock_openai_client()
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with pytest.raises(ValueError, match="Model is required"):
FoundryChatClient(project_client=mock_project_client)
def test_init_with_empty_project_source_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("FOUNDRY_PROJECT_ENDPOINT", raising=False)
with pytest.raises(ValueError, match="Either 'project_endpoint' or 'project_client' is required"):
FoundryChatClient(model=_TEST_FOUNDRY_MODEL)
def test_init_with_project_endpoint_requires_credential() -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryChatClient(
project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT,
model=_TEST_FOUNDRY_MODEL,
)
async def test_configure_azure_monitor() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
mock_configure = MagicMock()
mock_views = MagicMock(return_value=[])
mock_resource = MagicMock()
mock_enable = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", mock_views),
patch("agent_framework.observability.create_resource", return_value=mock_resource),
patch("agent_framework.observability.enable_instrumentation", mock_enable),
):
await client.configure_azure_monitor(enable_sensitive_data=True)
project_client.telemetry.get_application_insights_connection_string.assert_called_once()
mock_configure.assert_called_once()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
assert call_kwargs["views"] == []
assert call_kwargs["resource"] is mock_resource
mock_enable.assert_called_once_with(enable_sensitive_data=True)
async def test_configure_azure_monitor_resource_not_found() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
side_effect=ResourceNotFoundError("No Application Insights found")
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
await client.configure_azure_monitor()
project_client.telemetry.get_application_insights_connection_string.assert_called_once()
async def test_configure_azure_monitor_import_error() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
original_import = __import__
def _import_with_missing_azure_monitor(
name: str,
globals: dict[str, Any] | None = None,
locals: dict[str, Any] | None = None,
fromlist: tuple[str, ...] = (),
level: int = 0,
) -> Any:
if name == "azure.monitor.opentelemetry":
raise ImportError("No module named 'azure.monitor.opentelemetry'")
return original_import(name, globals, locals, fromlist, level)
with (
patch.dict(sys.modules, {"azure.monitor.opentelemetry": None}),
patch("builtins.__import__", side_effect=_import_with_missing_azure_monitor),
pytest.raises(ImportError, match="azure-monitor-opentelemetry is required"),
):
await client.configure_azure_monitor()
async def test_configure_azure_monitor_with_custom_resource() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
custom_resource = MagicMock()
mock_configure = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", return_value=[]),
patch("agent_framework.observability.create_resource") as mock_create_resource,
patch("agent_framework.observability.enable_instrumentation"),
):
await client.configure_azure_monitor(resource=custom_resource)
mock_create_resource.assert_not_called()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["resource"] is custom_resource
async def test_get_response_with_invalid_input() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
with pytest.raises(ChatClientInvalidRequestException, match="Messages are required"):
await client.get_response(messages=[])
async def test_web_search_tool_with_location() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
web_search_tool = FoundryChatClient.get_web_search_tool(
user_location={
"city": "Seattle",
"country": "US",
"region": "WA",
"timezone": "America/Los_Angeles",
}
)
assert web_search_tool.user_location.city == "Seattle"
assert web_search_tool.user_location.country == "US"
_, 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:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
code_tool = FoundryChatClient.get_code_interpreter_tool()
assert code_tool.container["type"] == "auto"
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
assert run_options["tools"] == [code_tool]
code_tool_with_files = FoundryChatClient.get_code_interpreter_tool(file_ids=["file1", "file2"])
assert code_tool_with_files.container.file_ids == ["file1", "file2"]
_, 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_hosted_file_search_tool_validation() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
with pytest.raises(ValueError, match="vector_store_ids"):
FoundryChatClient.get_file_search_tool(vector_store_ids=[])
file_search_tool = FoundryChatClient.get_file_search_tool(vector_store_ids=["vs_123"])
assert file_search_tool.vector_store_ids == ["vs_123"]
_, 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:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
function_call = Content.from_function_call(
call_id="test-call-id",
name="test_function",
arguments='{"param": "value"}',
additional_properties={"fc_id": "test-fc-id"},
)
function_result = Content.from_function_result(call_id="test-call-id", result="Function executed successfully")
messages = [
Message(role="user", text="Call a function"),
Message(role="assistant", contents=[function_call]),
Message(role="tool", contents=[function_result]),
]
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_content_filter_exception() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_error = BadRequestError(
message="Content filter error",
response=MagicMock(),
body={"error": {"code": "content_filter", "message": "Content filter error"}},
)
mock_error.code = "content_filter"
client.client.responses.create.side_effect = mock_error
with pytest.raises(OpenAIContentFilterException) as exc_info:
await client.get_response(messages=[Message(role="user", text="Test message")])
assert "content error" in str(exc_info.value)
async def test_response_format_parse_path() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_parsed_response = MagicMock()
mock_parsed_response.id = "parsed_response_123"
mock_parsed_response.text = "Parsed response"
mock_parsed_response.model = "test-model"
mock_parsed_response.created_at = 1000000000
mock_parsed_response.metadata = {}
mock_parsed_response.output_parsed = None
mock_parsed_response.usage = None
mock_parsed_response.finish_reason = None
mock_parsed_response.conversation = None
client.client.responses.parse = AsyncMock(return_value=mock_parsed_response)
response = await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
assert response.conversation_id == "parsed_response_123"
assert response.model == "test-model"
async def test_response_format_parse_path_with_conversation_id() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_parsed_response = MagicMock()
mock_parsed_response.id = "parsed_response_123"
mock_parsed_response.text = "Parsed response"
mock_parsed_response.model = "test-model"
mock_parsed_response.created_at = 1000000000
mock_parsed_response.metadata = {}
mock_parsed_response.output_parsed = None
mock_parsed_response.usage = None
mock_parsed_response.finish_reason = None
mock_parsed_response.conversation = MagicMock()
mock_parsed_response.conversation.id = "conversation_456"
client.client.responses.parse = AsyncMock(return_value=mock_parsed_response)
response = await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
assert response.conversation_id == "conversation_456"
assert response.model == "test-model"
async def test_bad_request_error_non_content_filter() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_error = BadRequestError(
message="Invalid request",
response=MagicMock(),
body={"error": {"code": "invalid_request", "message": "Invalid request"}},
)
mock_error.code = "invalid_request"
client.client.responses.parse = AsyncMock(side_effect=mock_error)
with pytest.raises(ChatClientException) as exc_info:
await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct},
)
assert "failed to complete the prompt" in str(exc_info.value)
def test_get_mcp_tool_with_project_connection_id() -> None:
tool_config = FoundryChatClient.get_mcp_tool(
name="Docs MCP",
project_connection_id="conn-123",
allowed_tools=["search_docs"],
)
assert tool_config["project_connection_id"] == "conn-123"
assert tool_config["allowed_tools"] == ["search_docs"]
assert tool_config["server_label"] == "Docs_MCP"
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("tool_choice", "none", True, id="tool_choice_none"),
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
},
"required": ["location", "conditions"],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
@_with_foundry_debug()
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
client = FoundryChatClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
options: dict[str, Any] = {option_name: option_value}
if option_name.startswith("tool_choice"):
options["tools"] = [get_weather]
response = await client.get_response(messages=messages, options=options, stream=True).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
if needs_validation:
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
text = response.text.lower()
assert "sunny" in text or "seattle" in text
elif option_name.startswith("response_format"):
if option_value == OutputStruct:
assert response.value is not None
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
assert response.value is None
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@_with_foundry_debug()
async def test_integration_web_search() -> None:
client = FoundryChatClient(credential=AzureCliCredential())
web_search_tool = FoundryChatClient.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]},
}
response = await client.get_response(stream=True, **content).get_final_response()
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@_with_foundry_debug()
async def test_integration_tool_rich_content_image() -> None:
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@tool(approval_mode="never_require")
def get_test_image() -> Content:
return Content.from_data(data=image_bytes, media_type="image/jpeg")
client = FoundryChatClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
messages = [Message(role="user", text="Call the get_test_image tool and describe what you see.")]
options: dict[str, Any] = {"tools": [get_test_image], "tool_choice": "auto"}
response = await client.get_response(messages=messages, options=options, stream=True).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
def test_get_code_interpreter_tool() -> None:
"""Test code interpreter tool creation."""
tool_obj = RawFoundryChatClient.get_code_interpreter_tool()
assert tool_obj is not None
def test_get_code_interpreter_tool_with_file_ids() -> None:
"""Test code interpreter tool with file IDs."""
tool_obj = RawFoundryChatClient.get_code_interpreter_tool(file_ids=["file-abc123"])
assert tool_obj is not None
def test_get_file_search_tool() -> None:
"""Test file search tool creation."""
tool_obj = RawFoundryChatClient.get_file_search_tool(vector_store_ids=["vs_abc123"])
assert tool_obj is not None
def test_get_file_search_tool_requires_vector_store_ids() -> None:
"""Test that empty vector_store_ids raises ValueError."""
with pytest.raises(ValueError, match="vector_store_ids"):
RawFoundryChatClient.get_file_search_tool(vector_store_ids=[])
def test_get_web_search_tool() -> None:
"""Test web search tool creation."""
tool_obj = RawFoundryChatClient.get_web_search_tool()
assert tool_obj is not None
def test_get_web_search_tool_with_location() -> None:
"""Test web search tool with user location."""
tool_obj = RawFoundryChatClient.get_web_search_tool(
user_location={"city": "Seattle", "country": "US"},
search_context_size="high",
)
assert tool_obj is not None
def test_get_image_generation_tool() -> None:
"""Test image generation tool creation."""
tool_obj = RawFoundryChatClient.get_image_generation_tool()
assert tool_obj is not None
def test_get_mcp_tool() -> None:
"""Test MCP tool creation."""
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="my_mcp",
url="https://mcp.example.com",
)
assert tool_obj is not None
def test_get_mcp_tool_with_connection_id() -> None:
"""Test MCP tool with project connection ID."""
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="github_mcp",
project_connection_id="conn_abc123",
description="GitHub MCP via Foundry",
)
assert tool_obj is not None
@@ -0,0 +1,501 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework_foundry._memory_provider import FoundryMemoryProvider
@pytest.fixture
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
mock_client.beta = AsyncMock()
mock_client.beta.memory_stores = AsyncMock()
mock_client.beta.memory_stores.search_memories = AsyncMock()
mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@pytest.fixture
def mock_credential() -> Mock:
"""Create a mock Azure credential."""
return Mock()
# -- Initialization tests ------------------------------------------------------
def test_init_with_all_params(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
source_id="custom_source",
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
context_prompt="Custom prompt",
update_delay=60,
)
assert provider.source_id == "custom_source"
assert provider.project_client is mock_project_client
assert provider.memory_store_name == "test_store"
assert provider.scope == "user_123"
assert provider.context_prompt == "Custom prompt"
assert provider.update_delay == 60
def test_init_default_source_id(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.source_id == FoundryMemoryProvider.DEFAULT_SOURCE_ID
def test_init_default_context_prompt(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.context_prompt == FoundryMemoryProvider.DEFAULT_CONTEXT_PROMPT
def test_init_default_update_delay(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.update_delay == 300
def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMock, mock_credential: Mock) -> None:
with patch("agent_framework_foundry._memory_provider.AIProjectClient") as mock_ai_project_client:
mock_ai_project_client.return_value = mock_project_client
provider = FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
credential=mock_credential, # type: ignore[arg-type]
allow_preview=True,
memory_store_name="test_store",
scope="user_123",
)
assert provider.project_client is mock_project_client
mock_ai_project_client.assert_called_once_with(
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
def test_init_requires_project_endpoint_without_project_client() -> None:
with (
patch("agent_framework_foundry._memory_provider.load_settings") as mock_load_settings,
patch.dict(os.environ, {}, clear=True),
pytest.raises(ValueError, match="project endpoint is required"),
):
mock_load_settings.return_value = {"project_endpoint": None}
FoundryMemoryProvider(
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_credential_without_project_client() -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_memory_store_name(mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="memory_store_name is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="",
scope="user_123",
)
def test_init_requires_scope(mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="scope is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="",
)
# -- before_run tests ----------------------------------------------------------
async def test_retrieves_static_memories_on_first_run(mock_project_client: AsyncMock) -> None:
mem1 = Mock()
mem1.memory_item.content = "User prefers Python"
mem2 = Mock()
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should call search_memories twice: once for static, once for contextual
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
async def test_contextual_memories_added_to_context(mock_project_client: AsyncMock) -> None:
# Mock static search (first call)
static_mem = Mock()
static_mem.memory_item.content = "User prefers Python"
static_result = Mock()
static_result.memories = [static_mem]
# Mock contextual search (second call)
contextual_mem = Mock()
contextual_mem.memory_item.content = "Last discussed async patterns"
contextual_result = Mock()
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Check that memories were added to context
assert provider.source_id in ctx.context_messages
added = ctx.context_messages[provider.source_id]
assert len(added) == 1
assert "User prefers Python" in added[0].text # type: ignore[operator]
assert "Last discussed async patterns" in added[0].text # type: ignore[operator]
assert provider.context_prompt in added[0].text # type: ignore[operator]
assert session.state[provider.source_id]["previous_search_id"] == "search-123"
async def test_empty_input_skips_contextual_search(mock_project_client: AsyncMock) -> None:
static_result = Mock()
static_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should only call search_memories once for static memories
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(mock_project_client: AsyncMock) -> None:
mock_search_result = Mock()
mock_search_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="test")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert provider.source_id not in ctx.context_messages
async def test_static_memories_only_retrieved_once(mock_project_client: AsyncMock) -> None:
static_mem = Mock()
static_mem.memory_item.content = "Static memory"
static_result = Mock()
static_result.memories = [static_mem]
contextual_result = Mock()
contextual_result.memories = []
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# First call
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(mock_project_client: AsyncMock) -> None:
mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# Should not raise exception
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# No memories added
assert provider.source_id not in ctx.context_messages
# -- after_run tests -----------------------------------------------------------
async def test_stores_input_and_response(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_poller.update_id = "update-456"
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="question")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="answer")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
assert call_kwargs["items"][0]["content"] == "question"
assert call_kwargs["items"][1]["content"] == "answer"
assert session.state[provider.source_id]["previous_update_id"] == "update-456"
async def test_only_stores_user_assistant_system(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text="hello"),
Message(role="tool", text="tool output"),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[Message(role="assistant", text="reply")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
assert items[1]["content"] == "reply"
async def test_skips_empty_messages(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text=""),
Message(role="user", text=" "),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
update_delay=60,
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(mock_project_client: AsyncMock) -> None:
mock_poller1 = Mock()
mock_poller1.update_id = "update-1"
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx1 = SessionContext(input_messages=[Message(role="user", text="first")], session_id="s1")
ctx1._response = AgentResponse(messages=[Message(role="assistant", text="response1")])
# First update
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx1, state=session.state.setdefault(provider.source_id, {})
)
assert session.state[provider.source_id]["previous_update_id"] == "update-1"
# Second update should use previous_update_id
ctx2 = SessionContext(input_messages=[Message(role="user", text="second")], session_id="s1")
ctx2._response = AgentResponse(messages=[Message(role="assistant", text="response2")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(mock_project_client: AsyncMock) -> None:
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
# Should not raise exception
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# -- Context manager tests -----------------------------------------------------
async def test_aenter_delegates_to_client(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
result = await provider.__aenter__()
assert result is provider
mock_project_client.__aenter__.assert_awaited_once()
async def test_aexit_delegates_to_client(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
await provider.__aexit__(None, None, None)
mock_project_client.__aexit__.assert_awaited_once()
async def test_async_with_syntax(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
async with provider as p:
assert p is provider
@@ -1,374 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for FoundryAgentClient and FoundryAgent classes."""
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework._tools import tool
class TestRawFoundryAgentChatClient:
"""Tests for RawFoundryAgentChatClient."""
def test_init_requires_agent_name(self) -> None:
"""Test that agent_name is required."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
with pytest.raises(ValueError, match="Agent name is required"):
RawFoundryAgentChatClient(
project_client=MagicMock(),
)
def test_init_with_agent_name(self) -> None:
"""Test construction with agent_name and project_client."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
assert client.agent_version == "1.0"
def test_get_agent_reference_with_version(self) -> None:
"""Test agent reference includes version when provided."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="my-agent",
agent_version="2.0",
)
ref = client._get_agent_reference()
assert ref == {"name": "my-agent", "version": "2.0", "type": "agent_reference"}
def test_get_agent_reference_without_version(self) -> None:
"""Test agent reference omits version for HostedAgents."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="hosted-agent",
)
ref = client._get_agent_reference()
assert ref == {"name": "hosted-agent", "type": "agent_reference"}
assert "version" not in ref
def test_as_agent_returns_foundry_agent_and_preserves_client_type(self) -> None:
"""Test that as_agent() wraps the client in FoundryAgent using the same client class."""
from agent_framework_foundry._foundry_agent import FoundryAgent
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
class CustomClient(RawFoundryAgentChatClient):
pass
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = CustomClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
agent = client.as_agent(instructions="You are helpful.")
assert isinstance(agent, FoundryAgent)
assert agent.name == "test-agent"
assert isinstance(agent.client, CustomClient)
assert agent.client.project_client is mock_project
assert agent.client.agent_name == "test-agent"
assert agent.client.agent_version == "1.0"
named_agent = client.as_agent(name="display-name", instructions="You are helpful.")
assert named_agent.name == "display-name"
assert named_agent.client.agent_name == "test-agent"
async def test_prepare_options_validates_tools(self) -> None:
"""Test that _prepare_options rejects non-FunctionTool objects."""
from agent_framework import Message
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
# A dict tool should be rejected
with pytest.raises(TypeError, match="Only FunctionTool objects are accepted"):
await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [{"type": "function", "function": {"name": "bad"}}]},
)
async def test_prepare_options_accepts_function_tools(self) -> None:
"""Test that _prepare_options accepts FunctionTool objects."""
from agent_framework import Message
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_openai = MagicMock()
mock_project.get_openai_client.return_value = mock_openai
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
# Should not raise — patch the parent's _prepare_options
with patch(
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
new_callable=AsyncMock,
return_value={},
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [my_func]},
)
assert "extra_body" in result
assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
def test_check_model_presence_is_noop(self) -> None:
"""Test that _check_model_presence does nothing (model is on service)."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
options: dict[str, Any] = {}
client._check_model_presence(options)
assert "model" not in options
class TestFoundryAgentChatClient:
"""Tests for _FoundryAgentChatClient (full middleware)."""
def test_init(self) -> None:
"""Test construction of the full-middleware client."""
from agent_framework_foundry._foundry_agent_client import _FoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = _FoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
class TestRawFoundryAgent:
"""Tests for RawFoundryAgent."""
def test_init_creates_client(self) -> None:
"""Test that RawFoundryAgent creates a client internally."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_init_with_custom_client_type(self) -> None:
"""Test that client_type parameter is respected."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
client_type=RawFoundryAgentChatClient,
)
assert isinstance(agent.client, RawFoundryAgentChatClient)
def test_init_rejects_invalid_client_type(self) -> None:
"""Test that invalid client_type raises TypeError."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
with pytest.raises(TypeError, match="must be a subclass of RawFoundryAgentChatClient"):
RawFoundryAgent(
project_client=MagicMock(),
agent_name="test-agent",
client_type=object, # type: ignore[arg-type]
)
def test_init_with_function_tools(self) -> None:
"""Test that FunctionTool and callables are accepted."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
tools=[my_func],
)
assert agent.default_options.get("tools") is not None
class TestFoundryAgent:
"""Tests for FoundryAgent (full middleware)."""
def test_init(self) -> None:
"""Test construction of the full-middleware agent."""
from agent_framework_foundry._foundry_agent import FoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_init_with_middleware(self) -> None:
"""Test that agent-level middleware is accepted."""
from agent_framework import ChatContext, ChatMiddleware
from agent_framework_foundry._foundry_agent import FoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
class MyMiddleware(ChatMiddleware):
async def process(self, context: ChatContext) -> None:
pass
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
middleware=[MyMiddleware()],
)
assert agent.client is not None
class TestFoundryChatClientToolMethods:
"""Tests for RawFoundryChatClient tool factory methods."""
def test_get_code_interpreter_tool(self) -> None:
"""Test code interpreter tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_code_interpreter_tool()
assert tool_obj is not None
def test_get_code_interpreter_tool_with_file_ids(self) -> None:
"""Test code interpreter tool with file IDs."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_code_interpreter_tool(file_ids=["file-abc123"])
assert tool_obj is not None
def test_get_file_search_tool(self) -> None:
"""Test file search tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_file_search_tool(vector_store_ids=["vs_abc123"])
assert tool_obj is not None
def test_get_file_search_tool_requires_vector_store_ids(self) -> None:
"""Test that empty vector_store_ids raises ValueError."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
with pytest.raises(ValueError, match="vector_store_ids"):
RawFoundryChatClient.get_file_search_tool(vector_store_ids=[])
def test_get_web_search_tool(self) -> None:
"""Test web search tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_web_search_tool()
assert tool_obj is not None
def test_get_web_search_tool_with_location(self) -> None:
"""Test web search tool with user location."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_web_search_tool(
user_location={"city": "Seattle", "country": "US"},
search_context_size="high",
)
assert tool_obj is not None
def test_get_image_generation_tool(self) -> None:
"""Test image generation tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_image_generation_tool()
assert tool_obj is not None
def test_get_mcp_tool(self) -> None:
"""Test MCP tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="my_mcp",
url="https://mcp.example.com",
)
assert tool_obj is not None
def test_get_mcp_tool_with_connection_id(self) -> None:
"""Test MCP tool with project connection ID."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="github_mcp",
project_connection_id="conn_abc123",
description="GitHub MCP via Foundry",
)
assert tool_obj is not None
@@ -1,507 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework_foundry._foundry_memory_provider import FoundryMemoryProvider
@pytest.fixture
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
mock_client.beta = AsyncMock()
mock_client.beta.memory_stores = AsyncMock()
mock_client.beta.memory_stores.search_memories = AsyncMock()
mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@pytest.fixture
def mock_credential() -> Mock:
"""Create a mock Azure credential."""
return Mock()
# -- Initialization tests ------------------------------------------------------
class TestInit:
"""Test FoundryMemoryProvider initialization."""
def test_init_with_all_params(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
source_id="custom_source",
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
context_prompt="Custom prompt",
update_delay=60,
)
assert provider.source_id == "custom_source"
assert provider.project_client is mock_project_client
assert provider.memory_store_name == "test_store"
assert provider.scope == "user_123"
assert provider.context_prompt == "Custom prompt"
assert provider.update_delay == 60
def test_init_default_source_id(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.source_id == FoundryMemoryProvider.DEFAULT_SOURCE_ID
def test_init_default_context_prompt(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.context_prompt == FoundryMemoryProvider.DEFAULT_CONTEXT_PROMPT
def test_init_default_update_delay(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.update_delay == 300
def test_init_with_project_endpoint_and_credential(
self, mock_project_client: AsyncMock, mock_credential: Mock
) -> None:
with patch("agent_framework_foundry._foundry_memory_provider.AIProjectClient") as mock_ai_project_client:
mock_ai_project_client.return_value = mock_project_client
provider = FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
credential=mock_credential, # type: ignore[arg-type]
allow_preview=True,
memory_store_name="test_store",
scope="user_123",
)
assert provider.project_client is mock_project_client
mock_ai_project_client.assert_called_once_with(
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
def test_init_requires_project_endpoint_without_project_client(self) -> None:
with (
patch("agent_framework_foundry._foundry_memory_provider.load_settings") as mock_load_settings,
patch.dict(os.environ, {}, clear=True),
pytest.raises(ValueError, match="project endpoint is required"),
):
mock_load_settings.return_value = {"project_endpoint": None}
FoundryMemoryProvider(
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_credential_without_project_client(self) -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_memory_store_name(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="memory_store_name is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="",
scope="user_123",
)
def test_init_requires_scope(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="scope is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="",
)
# -- before_run tests ----------------------------------------------------------
class TestBeforeRun:
"""Test before_run hook."""
async def test_retrieves_static_memories_on_first_run(self, mock_project_client: AsyncMock) -> None:
"""First call retrieves static (user profile) memories."""
mem1 = Mock()
mem1.memory_item.content = "User prefers Python"
mem2 = Mock()
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should call search_memories twice: once for static, once for contextual
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
async def test_contextual_memories_added_to_context(self, mock_project_client: AsyncMock) -> None:
"""Contextual search returns memories → messages added to context with prompt."""
# Mock static search (first call)
static_mem = Mock()
static_mem.memory_item.content = "User prefers Python"
static_result = Mock()
static_result.memories = [static_mem]
# Mock contextual search (second call)
contextual_mem = Mock()
contextual_mem.memory_item.content = "Last discussed async patterns"
contextual_result = Mock()
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Check that memories were added to context
assert provider.source_id in ctx.context_messages
added = ctx.context_messages[provider.source_id]
assert len(added) == 1
assert "User prefers Python" in added[0].text # type: ignore[operator]
assert "Last discussed async patterns" in added[0].text # type: ignore[operator]
assert provider.context_prompt in added[0].text # type: ignore[operator]
assert session.state[provider.source_id]["previous_search_id"] == "search-123"
async def test_empty_input_skips_contextual_search(self, mock_project_client: AsyncMock) -> None:
"""Empty input messages → only static search performed, no contextual search."""
static_result = Mock()
static_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should only call search_memories once for static memories
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
"""Empty search results → no messages added."""
mock_search_result = Mock()
mock_search_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="test")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert provider.source_id not in ctx.context_messages
async def test_static_memories_only_retrieved_once(self, mock_project_client: AsyncMock) -> None:
"""Static memories are only retrieved on the first call."""
static_mem = Mock()
static_mem.memory_item.content = "Static memory"
static_result = Mock()
static_result.memories = [static_mem]
contextual_result = Mock()
contextual_result.memories = []
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# First call
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Search exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# Should not raise exception
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# No memories added
assert provider.source_id not in ctx.context_messages
# -- after_run tests -----------------------------------------------------------
class TestAfterRun:
"""Test after_run hook."""
async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -> None:
"""Stores input+response messages via begin_update_memories."""
mock_poller = Mock()
mock_poller.update_id = "update-456"
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="question")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="answer")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
assert call_kwargs["items"][0]["content"] == "question"
assert call_kwargs["items"][1]["content"] == "answer"
assert session.state[provider.source_id]["previous_update_id"] == "update-456"
async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
"""Only stores user/assistant/system messages with text."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text="hello"),
Message(role="tool", text="tool output"),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[Message(role="assistant", text="reply")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
assert items[1]["content"] == "reply"
async def test_skips_empty_messages(self, mock_project_client: AsyncMock) -> None:
"""Skips messages with empty text."""
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text=""),
Message(role="user", text=" "),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
"""Uses the configured update_delay parameter."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
update_delay=60,
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
"""Uses previous_update_id for incremental updates."""
mock_poller1 = Mock()
mock_poller1.update_id = "update-1"
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx1 = SessionContext(input_messages=[Message(role="user", text="first")], session_id="s1")
ctx1._response = AgentResponse(messages=[Message(role="assistant", text="response1")])
# First update
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx1, state=session.state.setdefault(provider.source_id, {})
)
assert session.state[provider.source_id]["previous_update_id"] == "update-1"
# Second update should use previous_update_id
ctx2 = SessionContext(input_messages=[Message(role="user", text="second")], session_id="s1")
ctx2._response = AgentResponse(messages=[Message(role="assistant", text="response2")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Update exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
# Should not raise exception
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# -- Context manager tests -----------------------------------------------------
class TestContextManager:
"""Test __aenter__/__aexit__ delegation."""
async def test_aenter_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
result = await provider.__aenter__()
assert result is provider
mock_project_client.__aenter__.assert_awaited_once()
async def test_aexit_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
await provider.__aexit__(None, None, None)
mock_project_client.__aexit__.assert_awaited_once()
async def test_async_with_syntax(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
async with provider as p:
assert p is provider
@@ -7,7 +7,7 @@ configured for GAIA benchmark tasks using the OpenAI Responses API.
Required Environment Variables:
OPENAI_API_KEY: Your OpenAI API key
OPENAI_RESPONSES_MODEL_ID: Model to use with Responses API (e.g., gpt-4o, gpt-4o-mini)
OPENAI_RESPONSES_MODEL: Model to use with Responses API (e.g., gpt-4o, gpt-4o-mini)
Optional Environment Variables:
OPENAI_BASE_URL: Custom API base URL if using a proxy or compatible service
@@ -19,7 +19,7 @@ Authentication:
Example:
export OPENAI_API_KEY="sk-..."
export OPENAI_RESPONSES_MODEL_ID="gpt-4o"
export OPENAI_RESPONSES_MODEL="gpt-4o"
"""
from collections.abc import AsyncIterator
+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
+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")
```
@@ -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,
@@ -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."""
@@ -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."
)
@@ -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:
@@ -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",
@@ -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
@@ -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(
@@ -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]:
@@ -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
@@ -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,
@@ -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()
@@ -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
@@ -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]
@@ -57,8 +57,8 @@ Depending on the selected client, set the appropriate environment variables:
**For OpenAI clients:**
- `OPENAI_API_KEY`: Your OpenAI API key
- `OPENAI_CHAT_MODEL_ID`: The OpenAI model for `openai_chat` and `openai_assistants`
- `OPENAI_RESPONSES_MODEL_ID`: The OpenAI model for `openai_responses`
- `OPENAI_CHAT_MODEL`: The OpenAI model for `openai_chat` and `openai_assistants`
- `OPENAI_RESPONSES_MODEL`: The OpenAI model for `openai_responses`
**For Anthropic client (`anthropic`):**
- `ANTHROPIC_API_KEY`: Your Anthropic API key
@@ -4,8 +4,9 @@ import asyncio
import os
from agent_framework import Agent
from agent_framework.azure import AzureAISearchContextProvider, AzureOpenAIEmbeddingClient
from agent_framework.azure import AzureAISearchContextProvider
from agent_framework.foundry import FoundryChatClient
from agent_framework.openai import OpenAIEmbeddingClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
@@ -31,8 +32,8 @@ Prerequisites:
- AZURE_SEARCH_INDEX_NAME: Your search index name
- FOUNDRY_PROJECT_ENDPOINT: Your Azure AI Foundry project endpoint
- AZURE_AI_MODEL_DEPLOYMENT_NAME: Your model deployment name (e.g., "gpt-4o")
- AZURE_OPENAI_EMBEDDING_MODEL_ID: (Optional) Your embedding model for hybrid search (e.g., "text-embedding-3-small")
- AZURE_OPENAI_ENDPOINT: (Optional) Your Azure OpenAI resource URL, required if using an OpenAI embedding model for hybrid search
- AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME: (Optional) Your Azure OpenAI embedding deployment for hybrid search
- AZURE_OPENAI_ENDPOINT: (Optional) Your Azure OpenAI resource URL, required if using Azure OpenAI embeddings
"""
# Sample queries to demonstrate RAG
@@ -55,13 +56,13 @@ async def main() -> None:
project_endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"]
model_deployment = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o")
openai_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
embedding_model = os.environ.get("AZURE_OPENAI_EMBEDDING_MODEL_ID", "text-embedding-3-small")
embedding_deployment = os.environ.get("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")
embedding_client = None
if openai_endpoint and embedding_model:
embedding_client = AzureOpenAIEmbeddingClient(
endpoint=openai_endpoint,
model=embedding_model,
if openai_endpoint and embedding_deployment:
embedding_client = OpenAIEmbeddingClient(
azure_endpoint=openai_endpoint,
model=embedding_deployment,
credential=credential,
)
+1 -1
View File
@@ -85,7 +85,7 @@ Alternatively, set environment variables globally:
```bash
export OPENAI_API_KEY="your-key-here"
export OPENAI_CHAT_MODEL_ID="gpt-4o"
export OPENAI_CHAT_MODEL="gpt-4o"
```
## Using DevUI with Your Own Agents
@@ -2,55 +2,59 @@
# Run with: uv run samples/02-agents/embeddings/azure_openai_embeddings.py
import asyncio
import os
from agent_framework.azure import AzureOpenAIEmbeddingClient
from agent_framework.openai import OpenAIEmbeddingClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
load_dotenv()
"""Azure OpenAI Embedding Client Example
This sample demonstrates how to generate embeddings using the Azure OpenAI embedding client.
It supports both API key and Azure credential authentication.
"""This sample demonstrates Azure OpenAI embedding generation with ``OpenAIEmbeddingClient``.
Prerequisites:
Set the following environment variables or add them to a .env file:
- AZURE_OPENAI_ENDPOINT: Your Azure OpenAI endpoint URL
- AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME: The embedding model deployment name
- AZURE_OPENAI_API_KEY: Your API key (or use Azure credential instead)
Set the following environment variables or add them to a local ``.env`` file:
- ``AZURE_OPENAI_ENDPOINT``: Your Azure OpenAI endpoint URL
- ``AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME``: The embedding deployment name
- ``AZURE_OPENAI_API_VERSION``: Optional API version override
Sign in with ``az login`` before running the sample.
"""
load_dotenv()
async def main() -> None:
"""Generate embeddings with Azure OpenAI."""
# 1. Create a client using environment variables.
# Reads AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME,
# and AZURE_OPENAI_API_KEY from environment.
client = AzureOpenAIEmbeddingClient()
async with AzureCliCredential() as credential:
client = OpenAIEmbeddingClient(
model=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=credential,
)
# 2. Generate a single embedding.
result = await client.get_embeddings(["Hello, world!"])
print(f"Single embedding dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
print(f"Model: {result[0].model_id}")
print(f"Usage: {result.usage}")
print()
# 1. Generate a single embedding.
result = await client.get_embeddings(["Hello, world!"])
print(f"Single embedding dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
print(f"Model: {result[0].model}")
print(f"Usage: {result.usage}")
print()
# 3. Generate embeddings for multiple inputs.
texts = [
"The weather is sunny today.",
"It is raining outside.",
"Machine learning is fascinating.",
]
result = await client.get_embeddings(texts)
print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions")
print()
# 2. Generate embeddings for multiple inputs.
texts = [
"The weather is sunny today.",
"It is raining outside.",
"Machine learning is fascinating.",
]
result = await client.get_embeddings(texts)
print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions")
print(f"First embedding vector: {result[0].vector[:5]}")
print()
# 4. Generate embeddings with custom dimensions.
result = await client.get_embeddings(["Custom dimensions example"], options={"dimensions": 256})
print(f"Custom dimensions: {result[0].dimensions}")
# 3. Generate embeddings with custom dimensions.
result = await client.get_embeddings(["Custom dimensions example"], options={"dimensions": 256})
print(f"Custom dimensions: {result[0].dimensions}")
if __name__ == "__main__":
@@ -3,31 +3,32 @@
# Run with: uv run samples/02-agents/embeddings/openai_embeddings.py
import asyncio
import os
from agent_framework.openai import OpenAIEmbeddingClient
from dotenv import load_dotenv
load_dotenv()
"""OpenAI Embedding Client Example
This sample demonstrates how to generate embeddings using the OpenAI embedding client.
It shows single and batch embedding generation, as well as custom dimensions.
"""This sample demonstrates OpenAI embedding generation with explicit constructor settings.
Prerequisites:
Set the OPENAI_API_KEY environment variable or add it to a .env file.
Set ``OPENAI_API_KEY`` in your environment or in a local ``.env`` file.
"""
load_dotenv()
async def main() -> None:
"""Generate embeddings with OpenAI."""
client = OpenAIEmbeddingClient(model="text-embedding-3-small")
client = OpenAIEmbeddingClient(
model="text-embedding-3-small",
api_key=os.getenv("OPENAI_API_KEY"),
)
# 1. Generate a single embedding.
result = await client.get_embeddings(["Hello, world!"])
print(f"Single embedding dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
print(f"Model: {result[0].model_id}")
print(f"Model: {result[0].model}")
print(f"Usage: {result.usage}")
print()
@@ -39,7 +40,7 @@ async def main() -> None:
]
result = await client.get_embeddings(texts)
print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions")
print(f"First embedding vector: {result[0].vector[:5]}") # Print first 5 values of the first embedding
print(f"First embedding vector: {result[0].vector[:5]}")
print()
# 3. Generate embeddings with custom dimensions.
+1 -1
View File
@@ -17,7 +17,7 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to
## Prerequisites
- `OPENAI_API_KEY` environment variable
- `OPENAI_RESPONSES_MODEL_ID` environment variable
- `OPENAI_RESPONSES_MODEL` environment variable
For `mcp_github_pat.py`:
- `GITHUB_PAT` - Your GitHub Personal Access Token (create at https://github.com/settings/tokens)
@@ -25,7 +25,7 @@ The new usage tracking sample uses `OpenAIResponsesClient`, so set the usual Ope
```bash
export OPENAI_API_KEY="your-openai-api-key"
export OPENAI_RESPONSES_MODEL_ID="gpt-4.1-mini"
export OPENAI_RESPONSES_MODEL="gpt-4.1-mini"
```
Then run:
@@ -40,8 +40,8 @@ ENABLE_SENSITIVE_DATA=true
# OpenAI specific variables
# ==========================
OPENAI_API_KEY="..."
OPENAI_RESPONSES_MODEL_ID="gpt-4o-2024-08-06"
OPENAI_CHAT_MODEL_ID="gpt-4o-2024-08-06"
OPENAI_RESPONSES_MODEL="gpt-4o-2024-08-06"
OPENAI_CHAT_MODEL="gpt-4o-2024-08-06"
# Azure AI Foundry specific variables
# ====================================
@@ -1,12 +1,48 @@
# Azure Provider Samples
This folder contains Azure OpenAI chat completion samples for Agent Framework.
This folder contains Azure-backed samples for the generic OpenAI clients in
`agent_framework.openai`.
## Azure OpenAI ChatCompletionClient Samples
## Chat Completions API samples (`OpenAIChatCompletionClient`)
| File | Description |
|------|-------------|
| [`openai_chat_completion_client_azure_basic.py`](openai_chat_completion_client_azure_basic.py) | Azure OpenAI Chat Client Basic Example |
| [`openai_chat_completion_client_azure_with_explicit_settings.py`](openai_chat_completion_client_azure_with_explicit_settings.py) | Azure OpenAI Chat Client with Explicit Settings Example |
| [`openai_chat_completion_client_azure_with_function_tools.py`](openai_chat_completion_client_azure_with_function_tools.py) | Azure OpenAI Chat Client with Function Tools Example |
| [`openai_chat_completion_client_azure_with_session.py`](openai_chat_completion_client_azure_with_session.py) | Azure OpenAI Chat Client with Session Management Example |
| [`openai_chat_completion_client_basic.py`](openai_chat_completion_client_basic.py) | Basic Azure chat completions sample using explicit Azure settings and `credential=AzureCliCredential()`. |
| [`openai_chat_completion_client_with_explicit_settings.py`](openai_chat_completion_client_with_explicit_settings.py) | Azure chat completions sample with explicit settings. |
| [`openai_chat_completion_client_with_function_tools.py`](openai_chat_completion_client_with_function_tools.py) | Azure chat completions sample with function tools. |
| [`openai_chat_completion_client_with_session.py`](openai_chat_completion_client_with_session.py) | Azure chat completions sample with session management. |
## Responses API samples (`OpenAIChatClient`)
| File | Description |
|------|-------------|
| [`openai_client_basic.py`](openai_client_basic.py) | Basic Azure responses sample using explicit settings and `credential=AzureCliCredential()`. |
| [`openai_client_with_function_tools.py`](openai_client_with_function_tools.py) | Azure responses sample with function tools. |
| [`openai_client_with_session.py`](openai_client_with_session.py) | Azure responses sample with session management. |
| [`openai_client_with_structured_output.py`](openai_client_with_structured_output.py) | Azure responses sample with structured output. |
## Environment Variables
Set these before running the Azure provider samples:
- `AZURE_OPENAI_ENDPOINT`
- `AZURE_OPENAI_DEPLOYMENT_NAME`
Optionally, you can also set:
- `AZURE_OPENAI_API_KEY`
- `AZURE_OPENAI_API_VERSION`
- `AZURE_OPENAI_BASE_URL`
These Azure samples are written around explicit Azure inputs such as
`credential=AzureCliCredential()`, so they stay on Azure even if `OPENAI_API_KEY` is also present.
## Optional Dependencies
Credential-based samples require `azure-identity`:
```bash
pip install azure-identity
```
Run `az login` before executing the credential-based samples.
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
@@ -16,14 +17,12 @@ load_dotenv()
"""
Azure OpenAI Chat Client Basic Example
This sample demonstrates basic usage of OpenAIChatCompletionClient for direct chat-based
interactions, showing both streaming and non-streaming responses.
This sample demonstrates basic usage of OpenAIChatCompletionClient with explicit Azure
settings and a credential, showing both streaming and non-streaming responses.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
@@ -37,11 +36,14 @@ async def non_streaming_example() -> None:
"""Example of non-streaming response (get the complete result at once)."""
print("=== Non-streaming Response Example ===")
# Create agent with Azure Chat Client
# For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
# authentication option.
agent = Agent(
client=OpenAIChatCompletionClient(credential=AzureCliCredential()),
client=OpenAIChatCompletionClient(
model=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=AzureCliCredential(),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -56,11 +58,14 @@ async def streaming_example() -> None:
"""Example of streaming response (get results as they are generated)."""
print("=== Streaming Response Example ===")
# Create agent with Azure Chat Client
# For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
# authentication option.
agent = Agent(
client=OpenAIChatCompletionClient(credential=AzureCliCredential()),
client=OpenAIChatCompletionClient(
model=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=AzureCliCredential(),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -75,7 +80,7 @@ async def streaming_example() -> None:
async def main() -> None:
print("=== Basic Azure Chat Client Agent Example ===")
print("=== Basic Azure Chat Completion Client Agent Example ===")
await non_streaming_example()
await streaming_example()
@@ -15,16 +15,16 @@ from pydantic import Field
load_dotenv()
"""
Azure OpenAI Chat Client with Explicit Settings Example
OpenAI Chat Completion Client with Explicit Settings Example
This sample demonstrates creating Azure OpenAI Chat Client with explicit configuration
This samples connects to Azure OpenAI.
This sample demonstrates creating OpenAI Chat Completion Client with explicit configuration
settings rather than relying on environment variable defaults.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
@@ -39,13 +39,12 @@ async def main() -> None:
# For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
# authentication option.
_client = OpenAIChatCompletionClient(
model=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
credential=AzureCliCredential(),
)
agent = Agent(
client=_client,
client=OpenAIChatCompletionClient(
model=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
credential=AzureCliCredential(),
),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
@@ -0,0 +1,90 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
Azure OpenAI Chat Client Basic Example
This sample demonstrates basic usage of OpenAIChatClient with explicit Azure
settings and a credential, showing both streaming and non-streaming responses.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
async def non_streaming_example() -> None:
"""Example of non-streaming response (get the complete result at once)."""
print("=== Non-streaming Response Example ===")
agent = Agent(
client=OpenAIChatClient(
model=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=AzureCliCredential(),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Result: {result}\n")
async def streaming_example() -> None:
"""Example of streaming response (get results as they are generated)."""
print("=== Streaming Response Example ===")
agent = Agent(
client=OpenAIChatClient(
model=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
credential=AzureCliCredential(),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather in Portland?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
async for chunk in agent.run(query, stream=True):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
async def main() -> None:
print("=== Basic Azure OpenAI Chat Client Agent Example ===")
await non_streaming_example()
await streaming_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,137 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from datetime import datetime, timezone
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
Azure OpenAI Chat Client with Function Tools Example
This sample demonstrates function tool integration with Azure OpenAI Chat Client,
showing both agent-level and query-level tool configuration patterns.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
@tool(approval_mode="never_require")
def get_time() -> str:
"""Get the current UTC time."""
current_time = datetime.now(timezone.utc)
return f"The current UTC time is {current_time.strftime('%Y-%m-%d %H:%M:%S')}."
async def tools_on_agent_level() -> None:
"""Example showing tools defined when creating the agent."""
print("=== Tools Defined on Agent Level ===")
# Tools are provided when creating the agent
# The agent can use these tools for any query during its lifetime
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can provide weather and time information.",
tools=[get_weather, get_time], # Tools defined at agent creation
)
# First query - agent can use weather tool
query1 = "What's the weather like in New York?"
print(f"User: {query1}")
result1 = await agent.run(query1)
print(f"Agent: {result1}\n")
# Second query - agent can use time tool
query2 = "What's the current UTC time?"
print(f"User: {query2}")
result2 = await agent.run(query2)
print(f"Agent: {result2}\n")
# Third query - agent can use both tools if needed
query3 = "What's the weather in London and what's the current UTC time?"
print(f"User: {query3}")
result3 = await agent.run(query3)
print(f"Agent: {result3}\n")
async def tools_on_run_level() -> None:
"""Example showing tools passed to the run method."""
print("=== Tools Passed to Run Method ===")
# Agent created without tools
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
# No tools defined here
)
# First query with weather tool
query1 = "What's the weather like in Seattle?"
print(f"User: {query1}")
result1 = await agent.run(query1, tools=[get_weather]) # Tool passed to run method
print(f"Agent: {result1}\n")
# Second query with time tool
query2 = "What's the current UTC time?"
print(f"User: {query2}")
result2 = await agent.run(query2, tools=[get_time]) # Different tool for this query
print(f"Agent: {result2}\n")
# Third query with multiple tools
query3 = "What's the weather in Chicago and what's the current UTC time?"
print(f"User: {query3}")
result3 = await agent.run(query3, tools=[get_weather, get_time]) # Multiple tools
print(f"Agent: {result3}\n")
async def mixed_tools_example() -> None:
"""Example showing both agent-level tools and run-method tools."""
print("=== Mixed Tools Example (Agent + Run Method) ===")
# Agent created with some base tools
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a comprehensive assistant that can help with various information requests.",
tools=[get_weather], # Base tool available for all queries
)
# Query using both agent tool and additional run-method tools
query = "What's the weather in Denver and what's the current UTC time?"
print(f"User: {query}")
# Agent has access to get_weather (from creation) + additional tools from run method
result = await agent.run(
query,
tools=[get_time], # Additional tools for this specific query
)
print(f"Agent: {result}\n")
async def main() -> None:
print("=== Azure OpenAI Chat Client Agent with Function Tools Examples ===\n")
await tools_on_agent_level()
await tools_on_run_level()
await mixed_tools_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,152 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from random import randint
from typing import Annotated
from agent_framework import Agent, AgentSession, tool
from agent_framework.openai import OpenAIChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
Azure OpenAI Chat Client with Session Management Example
This sample demonstrates session management with Azure OpenAI Chat Client, showing
persistent conversation context and simplified response handling.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
async def example_with_automatic_session_creation() -> None:
"""Example showing automatic session creation."""
print("=== Automatic Session Creation Example ===")
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
# First conversation - no session provided, will be created automatically
query1 = "What's the weather like in Seattle?"
print(f"User: {query1}")
result1 = await agent.run(query1)
print(f"Agent: {result1.text}")
# Second conversation - still no session provided, will create another new session
query2 = "What was the last city I asked about?"
print(f"\nUser: {query2}")
result2 = await agent.run(query2)
print(f"Agent: {result2.text}")
print("Note: Each call creates a separate session, so the agent doesn't remember previous context.\n")
async def example_with_session_persistence_in_memory() -> None:
"""
Example showing session persistence across multiple conversations.
In this example, messages are stored in-memory.
"""
print("=== Session Persistence Example (In-Memory) ===")
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
# Create a new session that will be reused
session = agent.create_session()
# First conversation
query1 = "What's the weather like in Tokyo?"
print(f"User: {query1}")
result1 = await agent.run(query1, session=session, store=False)
print(f"Agent: {result1.text}")
# Second conversation using the same session - maintains context
query2 = "How about London?"
print(f"\nUser: {query2}")
result2 = await agent.run(query2, session=session, store=False)
print(f"Agent: {result2.text}")
# Third conversation - agent should remember both previous cities
query3 = "Which of the cities I asked about has better weather?"
print(f"\nUser: {query3}")
result3 = await agent.run(query3, session=session, store=False)
print(f"Agent: {result3.text}")
print("Note: The agent remembers context from previous messages in the same session.\n")
async def example_with_existing_session_id() -> None:
"""
Example showing how to work with an existing session ID from the service.
In this example, messages are stored on the server using OpenAI conversation state.
"""
print("=== Existing Session ID Example ===")
# First, create a conversation and capture the session ID
existing_session_id = None
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
# Start a conversation and get the session ID
session = agent.create_session()
query1 = "What's the weather in Paris?"
print(f"User: {query1}")
result1 = await agent.run(query1, session=session)
print(f"Agent: {result1.text}")
# The session ID is set after the first response
existing_session_id = session.service_session_id
print(f"Session ID: {existing_session_id}")
if existing_session_id:
print("\n--- Continuing with the same session ID in a new agent instance ---")
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_session_id)
query2 = "What was the last city I asked about?"
print(f"User: {query2}")
result2 = await agent.run(query2, session=session)
print(f"Agent: {result2.text}")
print("Note: The agent continues the conversation from the previous session by using session ID.\n")
async def main() -> None:
print("=== Azure OpenAI Chat Client Session Management Examples ===\n")
await example_with_automatic_session_creation()
await example_with_session_persistence_in_memory()
await example_with_existing_session_id()
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,93 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, AgentResponse
from agent_framework.openai import OpenAIChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from pydantic import BaseModel
# Load environment variables from .env file
load_dotenv()
"""
Azure OpenAI Chat Client with Structured Output Example
This sample demonstrates using structured output capabilities with Azure OpenAI Chat Client,
showing Pydantic model integration for type-safe response parsing and data extraction.
"""
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
city: str
description: str
async def non_streaming_example() -> None:
print("=== Non-streaming example ===")
# Create an Azure OpenAI Chat agent
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
name="CityAgent",
instructions="You are a helpful agent that describes cities in a structured format.",
)
# Ask the agent about a city
query = "Tell me about Paris, France"
print(f"User: {query}")
# Get structured response from the agent using response_format parameter
result = await agent.run(query, options={"response_format": OutputStruct})
# Access the structured output using the parsed value
if structured_data := result.value:
print("Structured Output Agent:")
print(f"City: {structured_data.city}")
print(f"Description: {structured_data.description}")
else:
print(f"Failed to parse response: {result.text}")
async def streaming_example() -> None:
print("=== Streaming example ===")
# Create an Azure OpenAI Chat agent
agent = Agent(
client=OpenAIChatClient(credential=AzureCliCredential()),
name="CityAgent",
instructions="You are a helpful agent that describes cities in a structured format.",
)
# Ask the agent about a city
query = "Tell me about Tokyo, Japan"
print(f"User: {query}")
# Get structured response from streaming agent using AgentResponse.from_update_generator
# This method collects all streaming updates and combines them into a single AgentResponse
result = await AgentResponse.from_update_generator(
agent.run(query, stream=True, options={"response_format": OutputStruct}),
output_format_type=OutputStruct,
)
# Access the structured output using the parsed value
if structured_data := result.value:
print("Structured Output (from streaming with AgentResponse.from_update_generator):")
print(f"City: {structured_data.city}")
print(f"Description: {structured_data.description}")
else:
print(f"Failed to parse response: {result.text}")
async def main() -> None:
print("=== Azure OpenAI Chat Client Agent with Structured Output ===")
await non_streaming_example()
await streaming_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -27,7 +27,7 @@ Both approaches allow you to extend the framework for your specific use cases wh
## Understanding Raw Client Classes
The framework provides `Raw...Client` classes (e.g., `RawOpenAIChatClient`, `RawOpenAIResponsesClient`, `RawAzureAIClient`) that are intermediate implementations without middleware, telemetry, or function invocation support.
The framework provides `Raw...Client` classes (e.g., `RawOpenAIChatClient`, `RawOpenAIChatCompletionClient`, `RawAzureAIClient`) that are intermediate implementations without middleware, telemetry, or function invocation support.
### Warning: Raw Clients Should Not Normally Be Used Directly
@@ -60,8 +60,8 @@ class MyCustomClient(
For most use cases, use the fully-featured public client classes which already have all layers correctly composed:
- `OpenAIChatClient` - OpenAI Chat completions with all layers
- `OpenAIResponsesClient` - OpenAI Responses API with all layers
- `OpenAIChatCompletionClient` - OpenAI Chat Completions API with all layers
- `OpenAIChatClient` - OpenAI Responses API with all layers
- `AzureOpenAIChatClient` - Azure OpenAI Chat with all layers
- `AzureOpenAIResponsesClient` - Azure OpenAI Responses with all layers
- `AzureAIClient` - Azure AI Project with all layers
@@ -1,67 +1,63 @@
# OpenAI Agent Framework Examples
# OpenAI Provider Samples
This folder contains examples demonstrating different ways to create and use agents with the OpenAI clients from the `agent_framework.openai` package.
This folder contains OpenAI provider samples for the generic clients in
`agent_framework.openai`.
## Examples
## Chat Completions API samples (`OpenAIChatCompletionClient`)
| File | Description |
|------|-------------|
| [`openai_assistants_basic.py`](openai_assistants_basic.py) | Basic usage of `OpenAIAssistantProvider` with streaming and non-streaming responses. |
| [`openai_assistants_provider_methods.py`](openai_assistants_provider_methods.py) | Demonstrates all `OpenAIAssistantProvider` methods: `create_agent()`, `get_agent()`, and `as_agent()`. |
| [`openai_assistants_with_code_interpreter.py`](openai_assistants_with_code_interpreter.py) | Using `OpenAIAssistantsClient.get_code_interpreter_tool()` with `OpenAIAssistantProvider` to execute Python code. |
| [`openai_assistants_with_existing_assistant.py`](openai_assistants_with_existing_assistant.py) | Working with pre-existing assistants using `get_agent()` and `as_agent()` methods. |
| [`openai_assistants_with_explicit_settings.py`](openai_assistants_with_explicit_settings.py) | Configuring `OpenAIAssistantProvider` with explicit settings including API key and model ID. |
| [`openai_assistants_with_file_search.py`](openai_assistants_with_file_search.py) | Using `OpenAIAssistantsClient.get_file_search_tool()` with `OpenAIAssistantProvider` for file search capabilities. |
| [`openai_assistants_with_function_tools.py`](openai_assistants_with_function_tools.py) | Function tools with `OpenAIAssistantProvider` at both agent-level and query-level. |
| [`openai_assistants_with_response_format.py`](openai_assistants_with_response_format.py) | Structured outputs with `OpenAIAssistantProvider` using Pydantic models. |
| [`openai_assistants_with_session.py`](openai_assistants_with_session.py) | Session management with `OpenAIAssistantProvider` for conversation context persistence. |
| [`openai_chat_client_basic.py`](openai_chat_client_basic.py) | The simplest way to create an agent using `Agent` with `OpenAIChatClient`. Shows both streaming and non-streaming responses for chat-based interactions with OpenAI models. |
| [`openai_chat_client_with_explicit_settings.py`](openai_chat_client_with_explicit_settings.py) | Shows how to initialize an agent with a specific chat client, configuring settings explicitly including API key and model ID. |
| [`openai_chat_client_with_function_tools.py`](openai_chat_client_with_function_tools.py) | Demonstrates how to use function tools with agents. Shows both agent-level tools (defined when creating the agent) and query-level tools (provided with specific queries). |
| [`openai_chat_client_with_local_mcp.py`](openai_chat_client_with_local_mcp.py) | Shows how to integrate OpenAI agents with local Model Context Protocol (MCP) servers for enhanced functionality and tool integration. |
| [`openai_chat_client_with_session.py`](openai_chat_client_with_session.py) | Demonstrates session management with OpenAI agents, including automatic session creation for stateless conversations and explicit session management for maintaining conversation context across multiple interactions. |
| [`openai_chat_client_with_web_search.py`](openai_chat_client_with_web_search.py) | Shows how to use `OpenAIChatClient.get_web_search_tool()` for web search capabilities with OpenAI agents. |
| [`openai_chat_client_with_runtime_json_schema.py`](openai_chat_client_with_runtime_json_schema.py) | Shows how to supply a runtime JSON Schema via `additional_chat_options` for structured output without defining a Pydantic model. |
| [`openai_responses_client_basic.py`](openai_responses_client_basic.py) | The simplest way to create an agent using `Agent` with `OpenAIResponsesClient`. Shows both streaming and non-streaming responses for structured response generation with OpenAI models. |
| [`openai_responses_client_image_analysis.py`](openai_responses_client_image_analysis.py) | Demonstrates how to use vision capabilities with agents to analyze images. |
| [`openai_responses_client_image_generation.py`](openai_responses_client_image_generation.py) | Demonstrates how to use `OpenAIResponsesClient.get_image_generation_tool()` to create images based on text descriptions. |
| [`openai_responses_client_reasoning.py`](openai_responses_client_reasoning.py) | Demonstrates how to use reasoning capabilities with OpenAI agents, showing how the agent can provide detailed reasoning for its responses. |
| [`openai_responses_client_streaming_image_generation.py`](openai_responses_client_streaming_image_generation.py) | Demonstrates streaming image generation with partial images for real-time image creation feedback and improved user experience. |
| [`openai_responses_client_with_agent_as_tool.py`](openai_responses_client_with_agent_as_tool.py) | Shows how to use the agent-as-tool pattern with OpenAI Responses Client, where one agent delegates work to specialized sub-agents wrapped as tools using `as_tool()`. Demonstrates hierarchical agent architectures. |
| [`openai_responses_client_with_code_interpreter.py`](openai_responses_client_with_code_interpreter.py) | Shows how to use `OpenAIResponsesClient.get_code_interpreter_tool()` to write and execute Python code. |
| [`openai_responses_client_with_code_interpreter_files.py`](openai_responses_client_with_code_interpreter_files.py) | Shows how to use code interpreter with uploaded files for data analysis. |
| [`openai_responses_client_with_explicit_settings.py`](openai_responses_client_with_explicit_settings.py) | Shows how to initialize an agent with a specific responses client, configuring settings explicitly including API key and model ID. |
| [`openai_responses_client_with_file_search.py`](openai_responses_client_with_file_search.py) | Demonstrates how to use `OpenAIResponsesClient.get_file_search_tool()` for searching through uploaded files. |
| [`openai_responses_client_with_function_tools.py`](openai_responses_client_with_function_tools.py) | Demonstrates how to use function tools with agents. Shows both agent-level tools (defined when creating the agent) and run-level tools (provided with specific queries). |
| [`openai_responses_client_with_hosted_mcp.py`](openai_responses_client_with_hosted_mcp.py) | Shows how to use `OpenAIResponsesClient.get_mcp_tool()` for hosted MCP servers, including approval workflows. |
| [`openai_responses_client_with_local_mcp.py`](openai_responses_client_with_local_mcp.py) | Shows how to integrate OpenAI agents with local Model Context Protocol (MCP) servers for enhanced functionality and tool integration. |
| [`openai_responses_client_with_runtime_json_schema.py`](openai_responses_client_with_runtime_json_schema.py) | Shows how to supply a runtime JSON Schema via `additional_chat_options` for structured output without defining a Pydantic model. |
| [`openai_responses_client_with_structured_output.py`](openai_responses_client_with_structured_output.py) | Demonstrates how to use structured outputs with OpenAI agents to get structured data responses in predefined formats. |
| [`openai_responses_client_with_session.py`](openai_responses_client_with_session.py) | Demonstrates session management with OpenAI agents, including automatic session creation for stateless conversations and explicit session management for maintaining conversation context across multiple interactions. |
| [`openai_responses_client_with_web_search.py`](openai_responses_client_with_web_search.py) | Shows how to use `OpenAIResponsesClient.get_web_search_tool()` for web search capabilities. |
| [`chat_completion_client_basic.py`](chat_completion_client_basic.py) | Basic non-streaming and streaming chat completion sample with an explicit `gpt-5.4-nano` model and API key. |
| [`chat_completion_client_with_explicit_settings.py`](chat_completion_client_with_explicit_settings.py) | Chat completion sample with explicit model and API key settings. |
| [`chat_completion_client_with_function_tools.py`](chat_completion_client_with_function_tools.py) | Function tools with agent-level and run-level patterns. |
| [`chat_completion_client_with_local_mcp.py`](chat_completion_client_with_local_mcp.py) | Local MCP integration with the chat completions client. |
| [`chat_completion_client_with_runtime_json_schema.py`](chat_completion_client_with_runtime_json_schema.py) | Runtime JSON schema output with the chat completions client. |
| [`chat_completion_client_with_session.py`](chat_completion_client_with_session.py) | Session management with the chat completions client. |
| [`chat_completion_client_with_web_search.py`](chat_completion_client_with_web_search.py) | Web search with the chat completions client. |
## Responses API samples (`OpenAIChatClient`)
| File | Description |
|------|-------------|
| [`client_basic.py`](client_basic.py) | Basic non-streaming and streaming responses sample with an explicit `gpt-5.4-nano` model and API key. |
| [`client_image_analysis.py`](client_image_analysis.py) | Analyze images with the responses client. |
| [`client_image_generation.py`](client_image_generation.py) | Generate images from text prompts. |
| [`client_reasoning.py`](client_reasoning.py) | Reasoning-focused sample for models such as `gpt-5`. |
| [`client_streaming_image_generation.py`](client_streaming_image_generation.py) | Streaming image generation sample. |
| [`client_with_agent_as_tool.py`](client_with_agent_as_tool.py) | Agent-as-tool orchestration pattern. |
| [`client_with_code_interpreter.py`](client_with_code_interpreter.py) | Code interpreter sample. |
| [`client_with_code_interpreter_files.py`](client_with_code_interpreter_files.py) | Code interpreter sample with uploaded files. |
| [`client_with_explicit_settings.py`](client_with_explicit_settings.py) | Responses client with explicit model and API key settings. |
| [`client_with_file_search.py`](client_with_file_search.py) | Hosted file search sample. |
| [`client_with_function_tools.py`](client_with_function_tools.py) | Function tools with agent-level and run-level patterns. |
| [`client_with_hosted_mcp.py`](client_with_hosted_mcp.py) | Hosted MCP tools and approval workflows. |
| [`client_with_local_mcp.py`](client_with_local_mcp.py) | Local MCP integration with the responses client. |
| [`client_with_local_shell.py`](client_with_local_shell.py) | Local shell tool sample. |
| [`client_with_runtime_json_schema.py`](client_with_runtime_json_schema.py) | Runtime JSON schema output with the responses client. |
| [`client_with_session.py`](client_with_session.py) | Session management with the responses client. |
| [`client_with_shell.py`](client_with_shell.py) | Hosted shell tool sample. |
| [`client_with_structured_output.py`](client_with_structured_output.py) | Structured output with Pydantic models. |
| [`client_with_web_search.py`](client_with_web_search.py) | Web search with the responses client. |
## Environment Variables
Make sure to set the following environment variables before running the examples:
Set these before running the OpenAI provider samples:
- `OPENAI_API_KEY`: Your OpenAI API key
- `OPENAI_CHAT_MODEL_ID`: The OpenAI model to use (e.g., `gpt-4o`, `gpt-4o-mini`, `gpt-3.5-turbo`)
- `OPENAI_RESPONSES_MODEL_ID`: The OpenAI model to use (e.g., `gpt-4o`, `gpt-4o-mini`, `gpt-3.5-turbo`)
- For image processing examples, use a vision-capable model like `gpt-4o` or `gpt-4o-mini`
- `OPENAI_API_KEY`
- `OPENAI_MODEL`
Optionally, you can set:
- `OPENAI_ORG_ID`: Your OpenAI organization ID (if applicable)
- `OPENAI_API_BASE_URL`: Your OpenAI base URL (if using a different base URL)
Optionally, you can also set:
- `OPENAI_ORG_ID`
- `OPENAI_BASE_URL`
If your shell also contains `AZURE_OPENAI_*` variables, these samples still stay on OpenAI as long as
`OPENAI_API_KEY` is present. To force Azure routing with the generic clients, pass an explicit Azure
input such as `credential`, `azure_endpoint`, or `api_version`, or use the Azure provider samples.
## Optional Dependencies
Some examples require additional dependencies:
Some samples need extra packages:
- **Image Generation Example**: The `openai_responses_client_image_generation.py` example requires PIL (Pillow) for image display. Install with:
```bash
# Using uv
uv add pillow
# Or using pip
pip install pillow
```
- `client_image_generation.py` and `client_streaming_image_generation.py` use Pillow for image display.
- MCP samples require the relevant MCP server/tooling you configure locally.
@@ -0,0 +1,85 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Chat Completion Client Basic Example
This sample demonstrates basic usage of OpenAIChatCompletionClient with explicit model and
API key settings, showing both streaming and non-streaming responses.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
async def non_streaming_example() -> None:
"""Example of non-streaming response (get the complete result at once)."""
print("=== Non-streaming Response Example ===")
agent = Agent(
client=OpenAIChatCompletionClient(
model="gpt-5.4-nano",
api_key=os.getenv("OPENAI_API_KEY"),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Result: {result}\n")
async def streaming_example() -> None:
"""Example of streaming response (get results as they are generated)."""
print("=== Streaming Response Example ===")
agent = Agent(
client=OpenAIChatCompletionClient(
model="gpt-5.4-nano",
api_key=os.getenv("OPENAI_API_KEY"),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Portland?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
async for chunk in agent.run(query, stream=True):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
async def main() -> None:
print("=== Basic OpenAI Chat Completion Client Agent Example ===")
await non_streaming_example()
await streaming_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -6,7 +6,7 @@ from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
from pydantic import Field
@@ -14,9 +14,9 @@ from pydantic import Field
load_dotenv()
"""
OpenAI Responses Client with Explicit Settings Example
OpenAI Chat Completion Client with Explicit Settings Example
This sample demonstrates creating OpenAI Responses Client with explicit configuration
This sample demonstrates creating OpenAI Chat Completion Client with explicit configuration
settings rather than relying on environment variable defaults.
"""
@@ -34,15 +34,13 @@ def get_weather(
async def main() -> None:
print("=== OpenAI Responses Client with Explicit Settings ===")
_client = OpenAIResponsesClient(
model=os.environ["OPENAI_MODEL"],
api_key=os.environ["OPENAI_API_KEY"],
)
print("=== OpenAI Chat Completion Client with Explicit Settings ===")
agent = Agent(
client=_client,
client=OpenAIChatCompletionClient(
model=os.environ["OPENAI_MODEL"],
api_key=os.environ["OPENAI_API_KEY"],
),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -6,7 +6,7 @@ from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
from pydantic import Field
@@ -14,9 +14,9 @@ from pydantic import Field
load_dotenv()
"""
OpenAI Responses Client with Function Tools Example
OpenAI Chat Completion Client with Function Tools Example
This sample demonstrates function tool integration with OpenAI Responses Client,
This sample demonstrates function tool integration with OpenAI Chat Completion Client,
showing both agent-level and query-level tool configuration patterns.
"""
@@ -47,7 +47,7 @@ async def tools_on_agent_level() -> None:
# Tools are provided when creating the agent
# The agent can use these tools for any query during its lifetime
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful assistant that can provide weather and time information.",
tools=[get_weather, get_time], # Tools defined at agent creation
)
@@ -77,7 +77,7 @@ async def tools_on_run_level() -> None:
# Agent created without tools
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful assistant.",
# No tools defined here
)
@@ -107,7 +107,7 @@ async def mixed_tools_example() -> None:
# Agent created with some base tools
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a comprehensive assistant that can help with various information requests.",
tools=[get_weather], # Base tool available for all queries
)
@@ -125,7 +125,7 @@ async def mixed_tools_example() -> None:
async def main() -> None:
print("=== OpenAI Responses Client Agent with Function Tools Examples ===\n")
print("=== OpenAI Chat Completion Client Agent with Function Tools Examples ===\n")
await tools_on_agent_level()
await tools_on_run_level()
@@ -3,17 +3,17 @@
import asyncio
from agent_framework import Agent, MCPStreamableHTTPTool
from agent_framework.openai import OpenAIChatClient
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Chat Client with Local MCP Example
OpenAI Chat Completion Client with Local MCP Example
This sample demonstrates integrating Model Context Protocol (MCP) tools with
OpenAI Chat Client for extended functionality and external service access.
OpenAI Chat Completion Client for extended functionality and external service access.
The Agent Framework now supports enhanced metadata extraction from MCP tool
results, including error states, token usage, costs, and other arbitrary
@@ -34,7 +34,7 @@ async def mcp_tools_on_run_level() -> None:
url="https://learn.microsoft.com/api/mcp",
) as mcp_server,
Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
name="DocsAgent",
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
) as agent,
@@ -60,7 +60,7 @@ async def mcp_tools_on_agent_level() -> None:
# The agent can use these tools for any query during its lifetime
# The agent will connect to the MCP server through its context manager.
async with Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
name="DocsAgent",
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=MCPStreamableHTTPTool( # Tools defined at agent creation
@@ -82,7 +82,7 @@ async def mcp_tools_on_agent_level() -> None:
async def main() -> None:
print("=== OpenAI Chat Client Agent with MCP Tools Examples ===\n")
print("=== OpenAI Chat Completion Client Agent with MCP Tools Examples ===\n")
await mcp_tools_on_agent_level()
await mcp_tools_on_run_level()
@@ -4,14 +4,14 @@ import asyncio
import json
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatCompletionClient, OpenAIChatOptions
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Chat Client Runtime JSON Schema Example
OpenAI Chat Completion Client Runtime JSON Schema Example
Demonstrates structured outputs when the schema is only known at runtime.
Uses additional_chat_options to pass a JSON Schema payload directly to OpenAI
@@ -38,7 +38,7 @@ async def non_streaming_example() -> None:
print("=== Non-streaming runtime JSON schema example ===")
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatCompletionClient[OpenAIChatOptions](),
name="RuntimeSchemaAgent",
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
)
@@ -72,7 +72,7 @@ async def streaming_example() -> None:
print("=== Streaming runtime JSON schema example ===")
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatCompletionClient(),
name="RuntimeSchemaAgent",
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
)
@@ -108,7 +108,7 @@ async def streaming_example() -> None:
async def main() -> None:
print("=== OpenAI Chat Client with runtime JSON Schema ===")
print("=== OpenAI Chat Completion Client with runtime JSON Schema ===")
await non_streaming_example()
await streaming_example()
@@ -5,7 +5,7 @@ from random import randint
from typing import Annotated
from agent_framework import Agent, AgentSession, InMemoryHistoryProvider, tool
from agent_framework.openai import OpenAIChatClient
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
from pydantic import Field
@@ -13,9 +13,9 @@ from pydantic import Field
load_dotenv()
"""
OpenAI Chat Client with Session Management Example
OpenAI Chat Completion Client with Session Management Example
This sample demonstrates session management with OpenAI Chat Client, showing
This sample demonstrates session management with OpenAI Chat Completion Client, showing
conversation sessions and message history preservation across interactions.
"""
@@ -37,7 +37,7 @@ async def example_with_automatic_session_creation() -> None:
print("=== Automatic Session Creation Example ===")
agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -62,7 +62,7 @@ async def example_with_session_persistence() -> None:
print("Using the same session across multiple conversations to maintain context.\n")
agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -95,7 +95,7 @@ async def example_with_existing_session_messages() -> None:
print("=== Existing Session Messages Example ===")
agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -118,7 +118,7 @@ async def example_with_existing_session_messages() -> None:
# Create a new agent instance but use the existing session with its message history
new_agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatCompletionClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -142,7 +142,7 @@ async def example_with_existing_session_messages() -> None:
async def main() -> None:
print("=== OpenAI Chat Client Agent Session Management Examples ===\n")
print("=== OpenAI Chat Completion Client Agent Session Management Examples ===\n")
await example_with_automatic_session_creation()
await example_with_session_persistence()
@@ -3,22 +3,22 @@
import asyncio
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
from agent_framework.openai import OpenAIChatCompletionClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Chat Client with Web Search Example
OpenAI Chat Completion Client with Web Search Example
This sample demonstrates using get_web_search_tool() with OpenAI Chat Client
This sample demonstrates using get_web_search_tool() with OpenAI Chat Completion Client
for real-time information retrieval and current data access.
"""
async def main() -> None:
client = OpenAIChatClient(model="gpt-4o-search-preview")
client = OpenAIChatCompletionClient(model="gpt-4o-search-preview")
# Create web search tool with location context
web_search_tool = client.get_web_search_tool(
@@ -1,12 +1,14 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
@@ -14,17 +16,15 @@ load_dotenv()
"""
OpenAI Chat Client Basic Example
This sample demonstrates basic usage of OpenAIChatClient for direct chat-based
interactions, showing both streaming and non-streaming responses.
This sample demonstrates basic usage of OpenAIChatClient with explicit model and
API key settings, showing both streaming and non-streaming responses.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, "The location to get the weather for."],
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
@@ -36,13 +36,16 @@ async def non_streaming_example() -> None:
print("=== Non-streaming Response Example ===")
agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatClient(
model="gpt-5.4-nano",
api_key=os.getenv("OPENAI_API_KEY"),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Seattle?"
query = "What's the weather in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Result: {result}\n")
@@ -53,13 +56,16 @@ async def streaming_example() -> None:
print("=== Streaming Response Example ===")
agent = Agent(
client=OpenAIChatClient(),
client=OpenAIChatClient(
model="gpt-5.4-nano",
api_key=os.getenv("OPENAI_API_KEY"),
),
name="WeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Portland?"
query = "What's the weather in Portland?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
async for chunk in agent.run(query, stream=True):
@@ -3,26 +3,26 @@
import asyncio
from agent_framework import Agent, Content
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client Image Analysis Example
OpenAI Chat Client Image Analysis Example
This sample demonstrates using OpenAI Responses Client for image analysis and vision tasks,
This sample demonstrates using OpenAI Chat Client for image analysis and vision tasks,
showing multi-modal content handling with text and images.
"""
async def main():
print("=== OpenAI Responses Agent with Image Analysis ===")
print("=== OpenAI Chat Client Agent with Image Analysis ===")
# 1. Create an OpenAI Responses agent with vision capabilities
# 1. Create an OpenAI Chat agent with vision capabilities
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
name="VisionAgent",
instructions="You are a image analysist, you get a image and need to respond with what you see in the picture.",
)
@@ -7,17 +7,17 @@ import urllib.request as urllib_request
from pathlib import Path
from agent_framework import Agent, Content
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client Image Generation Example
OpenAI Chat Client Image Generation Example
This sample demonstrates how to generate images using OpenAI's DALL-E models
through the Responses Client. Image generation capabilities enable AI to create visual content from text,
through the Chat Client. Image generation capabilities enable AI to create visual content from text,
making it ideal for creative applications, content creation, design prototyping,
and automated visual asset generation.
"""
@@ -57,10 +57,10 @@ def save_image(output: Content) -> None:
async def main() -> None:
print("=== OpenAI Responses Image Generation Agent Example ===")
print("=== OpenAI Chat Image Generation Agent Example ===")
# Create an agent with customized image generation options
client = OpenAIResponsesClient()
client = OpenAIChatClient()
agent = Agent(
client=client,
instructions="You are a helpful AI that can generate images.",
@@ -3,14 +3,14 @@
import asyncio
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient, OpenAIResponsesOptions
from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client Reasoning Example
OpenAI Chat Client Reasoning Example
This sample demonstrates advanced reasoning capabilities using OpenAI's gpt-5 models,
showing step-by-step reasoning process visualization and complex problem-solving.
@@ -25,7 +25,7 @@ In this case they are here: https://platform.openai.com/docs/api-reference/respo
agent = Agent(
client=OpenAIResponsesClient[OpenAIResponsesOptions](model_id="gpt-5"),
client=OpenAIChatClient[OpenAIChatOptions](model_id="gpt-5"),
name="MathHelper",
instructions="You are a personal math tutor. When asked a math question, "
"reason over how best to approach the problem and share your thought process.",
@@ -76,7 +76,7 @@ async def streaming_reasoning_example() -> None:
async def main() -> None:
print("\033[92m=== Basic OpenAI Responses Reasoning Agent Example ===\033[0m")
print("\033[92m=== Basic OpenAI Chat Reasoning Agent Example ===\033[0m")
await reasoning_example()
await streaming_reasoning_example()
@@ -7,12 +7,12 @@ from pathlib import Path
import anyio
from agent_framework import Agent, Content
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""OpenAI Responses Client Streaming Image Generation Example
"""OpenAI Chat Client Streaming Image Generation Example
Demonstrates streaming partial image generation using OpenAI's image generation tool.
Shows progressive image rendering with partial images for improved user experience.
Note: The number of partial images received depends on generation speed:
@@ -42,7 +42,7 @@ async def main():
"""Demonstrate streaming image generation with partial images."""
print("=== OpenAI Streaming Image Generation Example ===\n")
# Create agent with streaming image generation enabled
client = OpenAIResponsesClient()
client = OpenAIChatClient()
agent = Agent(
client=client,
instructions="You are a helpful agent that can generate images.",
@@ -4,14 +4,14 @@ import asyncio
from collections.abc import Awaitable, Callable
from agent_framework import Agent, FunctionInvocationContext
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client Agent-as-Tool Example
OpenAI Chat Client Agent-as-Tool Example
Demonstrates hierarchical agent architectures where one agent delegates
work to specialized sub-agents wrapped as tools using as_tool().
@@ -35,9 +35,9 @@ async def logging_middleware(
async def main() -> None:
print("=== OpenAI Responses Client Agent-as-Tool Pattern ===")
print("=== OpenAI Chat Client Agent-as-Tool Pattern ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create a specialized writer agent
writer = Agent(
@@ -6,25 +6,25 @@ from agent_framework import (
Agent,
Content,
)
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with Code Interpreter Example
OpenAI Chat Client with Code Interpreter Example
This sample demonstrates using get_code_interpreter_tool() with OpenAI Responses Client
This sample demonstrates using get_code_interpreter_tool() with OpenAI Chat Client
for Python code execution and mathematical problem solving.
"""
async def main() -> None:
"""Example showing how to use the code interpreter tool with OpenAI Responses."""
print("=== OpenAI Responses Agent with Code Interpreter Example ===")
"""Example showing how to use the code interpreter tool with OpenAI Chat."""
print("=== OpenAI Chat Client Agent with Code Interpreter Example ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
agent = Agent(
client=client,
instructions="You are a helpful assistant that can write and execute Python code to solve problems.",
@@ -5,7 +5,7 @@ import os
import tempfile
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
from openai import AsyncOpenAI
@@ -13,9 +13,9 @@ from openai import AsyncOpenAI
load_dotenv()
"""
OpenAI Responses Client with Code Interpreter and Files Example
OpenAI Chat Client with Code Interpreter and Files Example
This sample demonstrates using get_code_interpreter_tool() with OpenAI Responses Client
This sample demonstrates using get_code_interpreter_tool() with OpenAI Chat Client
for Python code execution and data analysis with uploaded files.
"""
@@ -69,8 +69,8 @@ async def main() -> None:
temp_file_path, file_id = await create_sample_file_and_upload(openai_client)
# Create agent using OpenAI Responses client
client = OpenAIResponsesClient()
# Create agent using OpenAI Chat client
client = OpenAIChatClient()
agent = Agent(
client=client,
instructions="You are a helpful assistant that can analyze data files using Python code.",
@@ -3,23 +3,23 @@
import asyncio
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with File Search Example
OpenAI Chat Client with File Search Example
This sample demonstrates using get_file_search_tool() with OpenAI Responses Client
This sample demonstrates using get_file_search_tool() with OpenAI Chat Client
for direct document-based question answering and information retrieval.
"""
# Helper functions
async def create_vector_store(client: OpenAIResponsesClient) -> tuple[str, str]:
async def create_vector_store(client: OpenAIChatClient) -> tuple[str, str]:
"""Create a vector store with sample documents."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="user_data"
@@ -35,14 +35,14 @@ async def create_vector_store(client: OpenAIResponsesClient) -> tuple[str, str]:
return file.id, vector_store.id
async def delete_vector_store(client: OpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
async def delete_vector_store(client: OpenAIChatClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after using it."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
async def main() -> None:
client = OpenAIResponsesClient()
client = OpenAIChatClient()
message = "What is the weather today? Do a file search to find the answer."
@@ -4,7 +4,7 @@ import asyncio
from typing import TYPE_CHECKING, Any
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
if TYPE_CHECKING:
@@ -14,10 +14,10 @@ if TYPE_CHECKING:
load_dotenv()
"""
OpenAI Responses Client with Hosted MCP Example
OpenAI Chat Client with Hosted MCP Example
This sample demonstrates integrating hosted Model Context Protocol (MCP) tools with
OpenAI Responses Client, including user approval workflows for function call security.
OpenAI Chat Client, including user approval workflows for function call security.
"""
@@ -102,7 +102,7 @@ async def run_hosted_mcp_without_session_and_specific_approval() -> None:
"""Example showing Mcp Tools with approvals without using a session."""
print("=== Mcp with approvals and without session ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create MCP tool with specific approval mode
mcp_tool = client.get_mcp_tool(
name="Microsoft Learn MCP",
@@ -135,7 +135,7 @@ async def run_hosted_mcp_without_approval() -> None:
"""Example showing Mcp Tools without approvals."""
print("=== Mcp without approvals ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create MCP tool that never requires approval
mcp_tool = client.get_mcp_tool(
name="Microsoft Learn MCP",
@@ -167,7 +167,7 @@ async def run_hosted_mcp_with_session() -> None:
"""Example showing Mcp Tools with approvals using a session."""
print("=== Mcp with approvals and with session ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create MCP tool that always requires approval
mcp_tool = client.get_mcp_tool(
name="Microsoft Learn MCP",
@@ -200,7 +200,7 @@ async def run_hosted_mcp_with_session_streaming() -> None:
"""Example showing Mcp Tools with approvals using a session."""
print("=== Mcp with approvals and with session ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create MCP tool that always requires approval
mcp_tool = client.get_mcp_tool(
name="Microsoft Learn MCP",
@@ -234,7 +234,7 @@ async def run_hosted_mcp_with_session_streaming() -> None:
async def main() -> None:
print("=== OpenAI Responses Client Agent with Hosted Mcp Tools Examples ===\n")
print("=== OpenAI Chat Client Agent with Hosted Mcp Tools Examples ===\n")
await run_hosted_mcp_without_approval()
await run_hosted_mcp_without_session_and_specific_approval()
@@ -3,17 +3,17 @@
import asyncio
from agent_framework import Agent, MCPStreamableHTTPTool
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with Local MCP Example
OpenAI Chat Client with Local MCP Example
This sample demonstrates integrating local Model Context Protocol (MCP) tools with
OpenAI Responses Client for direct response generation with external capabilities.
OpenAI Chat Client for direct response generation with external capabilities.
"""
@@ -27,7 +27,7 @@ async def streaming_with_mcp(show_raw_stream: bool = False) -> None:
# Tools are provided when creating the agent
# The agent can use these tools for any query during its lifetime
async with Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
name="DocsAgent",
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=MCPStreamableHTTPTool( # Tools defined at agent creation
@@ -65,7 +65,7 @@ async def run_with_mcp() -> None:
# Tools are provided when creating the agent
# The agent can use these tools for any query during its lifetime
async with Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
name="DocsAgent",
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=MCPStreamableHTTPTool( # Tools defined at agent creation
@@ -87,7 +87,7 @@ async def run_with_mcp() -> None:
async def main() -> None:
print("=== OpenAI Responses Client Agent with Function Tools Examples ===\n")
print("=== OpenAI Chat Client Agent with Function Tools Examples ===\n")
await run_with_mcp()
await streaming_with_mcp()
@@ -5,14 +5,14 @@ import subprocess
from typing import Any
from agent_framework import Agent, Message, tool
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with Local Shell Tool Example
OpenAI Chat Client with Local Shell Tool Example
This sample demonstrates implementing a local shell tool using get_shell_tool(func=...)
that wraps Python's subprocess module. Unlike the hosted shell tool (get_shell_tool()),
@@ -53,7 +53,7 @@ async def main() -> None:
print("=== OpenAI Agent with Local Shell Tool Example ===")
print("NOTE: Commands will execute on your local machine.\n")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
local_shell_tool = client.get_shell_tool(
func=run_bash,
)
@@ -4,7 +4,7 @@ import asyncio
import json
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
@@ -38,7 +38,7 @@ async def non_streaming_example() -> None:
print("=== Non-streaming runtime JSON schema example ===")
agent = Agent(
client=OpenAIChatClient[OpenAIChatOptions](),
client=OpenAIChatClient(),
name="RuntimeSchemaAgent",
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
)
@@ -5,7 +5,7 @@ from random import randint
from typing import Annotated
from agent_framework import Agent, AgentSession, tool
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
from pydantic import Field
@@ -13,9 +13,9 @@ from pydantic import Field
load_dotenv()
"""
OpenAI Responses Client with Session Management Example
OpenAI Chat Client with Session Management Example
This sample demonstrates session management with OpenAI Responses Client, showing
This sample demonstrates session management with OpenAI Chat Client, showing
persistent conversation context and simplified response handling.
"""
@@ -37,7 +37,7 @@ async def example_with_automatic_session_creation() -> None:
print("=== Automatic Session Creation Example ===")
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -64,7 +64,7 @@ async def example_with_session_persistence_in_memory() -> None:
print("=== Session Persistence Example (In-Memory) ===")
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -103,7 +103,7 @@ async def example_with_existing_session_id() -> None:
existing_session_id = None
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -124,7 +124,7 @@ async def example_with_existing_session_id() -> None:
print("\n--- Continuing with the same session ID in a new agent instance ---")
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
@@ -3,16 +3,16 @@
import asyncio
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with Shell Tool Example
OpenAI Chat Client with Shell Tool Example
This sample demonstrates using get_shell_tool() with OpenAI Responses Client
This sample demonstrates using get_shell_tool() with OpenAI Chat Client
for executing shell commands in a managed container environment hosted by OpenAI.
The shell tool allows the model to run commands like listing files, running scripts,
@@ -21,10 +21,10 @@ or performing system operations within a secure, sandboxed container.
async def main() -> None:
"""Example showing how to use the shell tool with OpenAI Responses."""
print("=== OpenAI Responses Agent with Shell Tool Example ===")
"""Example showing how to use the shell tool with OpenAI Chat."""
print("=== OpenAI Chat Client Agent with Shell Tool Example ===")
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create a hosted shell tool with the default auto container environment
shell_tool = client.get_shell_tool()
@@ -3,7 +3,7 @@
import asyncio
from agent_framework import Agent, AgentResponse
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
from pydantic import BaseModel
@@ -11,9 +11,9 @@ from pydantic import BaseModel
load_dotenv()
"""
OpenAI Responses Client with Structured Output Example
OpenAI Chat Client with Structured Output Example
This sample demonstrates using structured output capabilities with OpenAI Responses Client,
This sample demonstrates using structured output capabilities with OpenAI Chat Client,
showing Pydantic model integration for type-safe response parsing and data extraction.
"""
@@ -28,9 +28,9 @@ class OutputStruct(BaseModel):
async def non_streaming_example() -> None:
print("=== Non-streaming example ===")
# Create an OpenAI Responses agent
# Create an OpenAI Chat agent
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
name="CityAgent",
instructions="You are a helpful agent that describes cities in a structured format.",
)
@@ -54,9 +54,9 @@ async def non_streaming_example() -> None:
async def streaming_example() -> None:
print("=== Streaming example ===")
# Create an OpenAI Responses agent
# Create an OpenAI Chat agent
agent = Agent(
client=OpenAIResponsesClient(),
client=OpenAIChatClient(),
name="CityAgent",
instructions="You are a helpful agent that describes cities in a structured format.",
)
@@ -82,7 +82,7 @@ async def streaming_example() -> None:
async def main() -> None:
print("=== OpenAI Responses Agent with Structured Output ===")
print("=== OpenAI Chat Client Agent with Structured Output ===")
await non_streaming_example()
await streaming_example()
@@ -3,22 +3,22 @@
import asyncio
from agent_framework import Agent
from agent_framework.openai import OpenAIResponsesClient
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client with Web Search Example
OpenAI Chat Client with Web Search Example
This sample demonstrates using get_web_search_tool() with OpenAI Responses Client
This sample demonstrates using get_web_search_tool() with OpenAI Chat Client
for direct real-time information retrieval and current data access.
"""
async def main() -> None:
client = OpenAIResponsesClient()
client = OpenAIChatClient()
# Create web search tool with location context
web_search_tool = client.get_web_search_tool(
@@ -1,98 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants Basic Example
This sample demonstrates basic usage of OpenAIAssistantProvider with automatic
assistant lifecycle management, showing both streaming and non-streaming responses.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
async def non_streaming_example() -> None:
"""Example of non-streaming response (get the complete result at once)."""
print("=== Non-streaming Response Example ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Create a new assistant via the provider
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
try:
query = "What's the weather like in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Agent: {result}\n")
finally:
# Clean up the assistant from OpenAI
await client.beta.assistants.delete(agent.id)
async def streaming_example() -> None:
"""Example of streaming response (get results as they are generated)."""
print("=== Streaming Response Example ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Create a new assistant via the provider
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
try:
query = "What's the weather like in Portland?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
async for chunk in agent.run(query, stream=True):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
finally:
# Clean up the assistant from OpenAI
await client.beta.assistants.delete(agent.id)
async def main() -> None:
print("=== Basic OpenAI Assistants Provider Example ===")
await non_streaming_example()
await streaming_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,158 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistant Provider Methods Example
This sample demonstrates the methods available on the OpenAIAssistantProvider class:
- create_agent(): Create a new assistant on the service
- get_agent(): Retrieve an existing assistant by ID
- as_agent(): Wrap an SDK Assistant object without making HTTP calls
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
async def create_agent_example() -> None:
"""Create a new assistant using provider.create_agent()."""
print("\n--- create_agent() ---")
async with (
AsyncOpenAI() as client,
OpenAIAssistantProvider(client) as provider,
):
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather assistant.",
tools=[get_weather],
)
try:
print(f"Created: {agent.name} (ID: {agent.id})")
result = await agent.run("What's the weather in Seattle?")
print(f"Response: {result}")
finally:
await client.beta.assistants.delete(agent.id)
async def get_agent_example() -> None:
"""Retrieve an existing assistant by ID using provider.get_agent()."""
print("\n--- get_agent() ---")
async with (
AsyncOpenAI() as client,
OpenAIAssistantProvider(client) as provider,
):
# Create an assistant directly with SDK (simulating pre-existing assistant)
sdk_assistant = await client.beta.assistants.create(
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
name="ExistingAssistant",
instructions="You always respond with 'Hello!'",
)
try:
# Retrieve using provider
agent = await provider.get_agent(sdk_assistant.id)
print(f"Retrieved: {agent.name} (ID: {agent.id})")
result = await agent.run("Hi there!")
print(f"Response: {result}")
finally:
await client.beta.assistants.delete(sdk_assistant.id)
async def as_agent_example() -> None:
"""Wrap an SDK Assistant object using Agent(client=provider, ...)."""
print("\n--- as_agent() ---")
async with (
AsyncOpenAI() as client,
OpenAIAssistantProvider(client) as provider,
):
# Create assistant using SDK
sdk_assistant = await client.beta.assistants.create(
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
name="WrappedAssistant",
instructions="You respond with poetry.",
)
try:
# Wrap synchronously (no HTTP call)
agent = Agent(client=provider, agent=sdk_assistant)
print(f"Wrapped: {agent.name} (ID: {agent.id})")
result = await agent.run("Tell me about the sunset.")
print(f"Response: {result}")
finally:
await client.beta.assistants.delete(sdk_assistant.id)
async def multiple_agents_example() -> None:
"""Create and manage multiple assistants with a single provider."""
print("\n--- Multiple Agents ---")
async with (
AsyncOpenAI() as client,
OpenAIAssistantProvider(client) as provider,
):
weather_agent = await provider.create_agent(
name="WeatherSpecialist",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a weather specialist.",
tools=[get_weather],
)
greeter_agent = await provider.create_agent(
name="GreeterAgent",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a friendly greeter.",
)
try:
print(f"Created: {weather_agent.name}, {greeter_agent.name}")
greeting = await greeter_agent.run("Hello!")
print(f"Greeter: {greeting}")
weather = await weather_agent.run("What's the weather in Tokyo?")
print(f"Weather: {weather}")
finally:
await client.beta.assistants.delete(weather_agent.id)
await client.beta.assistants.delete(greeter_agent.id)
async def main() -> None:
print("OpenAI Assistant Provider Methods")
await create_agent_example()
await get_agent_example()
await as_agent_example()
await multiple_agents_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,81 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from agent_framework import AgentResponseUpdate, ChatResponseUpdate
from agent_framework.openai import OpenAIAssistantProvider, OpenAIAssistantsClient
from dotenv import load_dotenv
from openai import AsyncOpenAI
from openai.types.beta.threads.runs import (
CodeInterpreterToolCallDelta,
RunStepDelta,
RunStepDeltaEvent,
ToolCallDeltaObject,
)
from openai.types.beta.threads.runs.code_interpreter_tool_call_delta import CodeInterpreter
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with Code Interpreter Example
This sample demonstrates using get_code_interpreter_tool() with OpenAI Assistants
for Python code execution and mathematical problem solving.
"""
def get_code_interpreter_chunk(chunk: AgentResponseUpdate) -> str | None:
"""Helper method to access code interpreter data."""
if (
isinstance(chunk.raw_representation, ChatResponseUpdate)
and isinstance(chunk.raw_representation.raw_representation, RunStepDeltaEvent)
and isinstance(chunk.raw_representation.raw_representation.delta, RunStepDelta)
and isinstance(chunk.raw_representation.raw_representation.delta.step_details, ToolCallDeltaObject)
and chunk.raw_representation.raw_representation.delta.step_details.tool_calls
):
for tool_call in chunk.raw_representation.raw_representation.delta.step_details.tool_calls:
if (
isinstance(tool_call, CodeInterpreterToolCallDelta)
and isinstance(tool_call.code_interpreter, CodeInterpreter)
and tool_call.code_interpreter.input is not None
):
return tool_call.code_interpreter.input
return None
async def main() -> None:
"""Example showing how to use the code interpreter tool with OpenAI Assistants."""
print("=== OpenAI Assistants Provider with Code Interpreter Example ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
chat_client = OpenAIAssistantsClient(client=client)
agent = await provider.create_agent(
name="CodeHelper",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant that can write and execute Python code to solve problems.",
tools=[chat_client.get_code_interpreter_tool()],
)
try:
query = "Use code to get the factorial of 100?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
generated_code = ""
async for chunk in agent.run(query, stream=True):
if chunk.text:
print(chunk.text, end="", flush=True)
code_interpreter_chunk = get_code_interpreter_chunk(chunk)
if code_interpreter_chunk is not None:
generated_code += code_interpreter_chunk
print(f"\nGenerated code:\n{generated_code}")
finally:
await client.beta.assistants.delete(agent.id)
if __name__ == "__main__":
asyncio.run(main())
@@ -1,118 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import Agent, tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with Existing Assistant Example
This sample demonstrates working with pre-existing OpenAI Assistants
using the provider's get_agent() and as_agent() methods.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
async def example_get_agent_by_id() -> None:
"""Example: Using get_agent() to retrieve an existing assistant by ID."""
print("=== Get Existing Assistant by ID ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Create an assistant via SDK (simulating an existing assistant)
created_assistant = await client.beta.assistants.create(
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
name="WeatherAssistant",
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather for a given location.",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string", "description": "The location"}},
"required": ["location"],
},
},
}
],
)
print(f"Created assistant: {created_assistant.id}")
try:
# Use get_agent() to retrieve the existing assistant
agent = await provider.get_agent(
assistant_id=created_assistant.id,
tools=[get_weather], # Required: implementation for function tools
instructions="You are a helpful weather agent.",
)
result = await agent.run("What's the weather like in Tokyo?")
print(f"Agent: {result}\n")
finally:
await client.beta.assistants.delete(created_assistant.id)
print("Assistant deleted.\n")
async def example_as_agent_wrap_sdk_object() -> None:
"""Example: Using as_agent() to wrap an existing SDK Assistant object."""
print("=== Wrap Existing SDK Assistant Object ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Create and fetch an assistant via SDK
created_assistant = await client.beta.assistants.create(
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
name="SimpleAssistant",
instructions="You are a friendly assistant.",
)
print(f"Created assistant: {created_assistant.id}")
try:
# Use as_agent() to wrap the SDK object
agent = Agent(
client=provider,
agent=created_assistant,
instructions="You are an extremely helpful assistant. Be enthusiastic!",
)
result = await agent.run("Hello! What can you help me with?")
print(f"Agent: {result}\n")
finally:
await client.beta.assistants.delete(created_assistant.id)
print("Assistant deleted.\n")
async def main() -> None:
print("=== OpenAI Assistants Provider with Existing Assistant Examples ===\n")
await example_get_agent_by_id()
await example_as_agent_wrap_sdk_object()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,61 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with Explicit Settings Example
This sample demonstrates creating OpenAI Assistants with explicit configuration
settings rather than relying on environment variable defaults.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
async def main() -> None:
print("=== OpenAI Assistants Provider with Explicit Settings ===")
# Create client with explicit API key
client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
provider = OpenAIAssistantProvider(client)
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ["OPENAI_MODEL"],
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
try:
query = "What's the weather like in New York?"
print(f"Query: {query}")
result = await agent.run(query)
print(f"Result: {result}\n")
finally:
await client.beta.assistants.delete(agent.id)
if __name__ == "__main__":
asyncio.run(main())
@@ -1,78 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from agent_framework import Content
from agent_framework.openai import OpenAIAssistantProvider, OpenAIAssistantsClient
from dotenv import load_dotenv
from openai import AsyncOpenAI
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with File Search Example
This sample demonstrates using get_file_search_tool() with OpenAI Assistants
for document-based question answering and information retrieval.
"""
async def create_vector_store(client: AsyncOpenAI) -> tuple[str, Content]:
"""Create a vector store with sample documents."""
file = await client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="user_data"
)
vector_store = await client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
if result.last_error is not None:
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, Content.from_hosted_vector_store(vector_store_id=vector_store.id)
async def delete_vector_store(client: AsyncOpenAI, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after using it."""
await client.vector_stores.delete(vector_store_id=vector_store_id)
await client.files.delete(file_id=file_id)
async def main() -> None:
print("=== OpenAI Assistants Provider with File Search Example ===\n")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
chat_client = OpenAIAssistantsClient(client=client)
agent = await provider.create_agent(
name="SearchAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant that searches files in a knowledge base.",
tools=[chat_client.get_file_search_tool()],
)
try:
query = "What is the weather today? Do a file search to find the answer."
file_id, vector_store_content = await create_vector_store(client)
print(f"User: {query}")
print("Agent: ", end="", flush=True)
async for chunk in agent.run(
query,
stream=True,
options={"tool_resources": {"file_search": {"vector_store_ids": [vector_store_content.vector_store_id]}}},
):
if chunk.text:
print(chunk.text, end="", flush=True)
await delete_vector_store(client, file_id, vector_store_content.vector_store_id)
finally:
await client.beta.assistants.delete(agent.id)
if __name__ == "__main__":
asyncio.run(main())
@@ -1,159 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from datetime import datetime, timezone
from random import randint
from typing import Annotated
from agent_framework import tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with Function Tools Example
This sample demonstrates function tool integration with OpenAI Assistants,
showing both agent-level and query-level tool configuration patterns.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
@tool(approval_mode="never_require")
def get_time() -> str:
"""Get the current UTC time."""
current_time = datetime.now(timezone.utc)
return f"The current UTC time is {current_time.strftime('%Y-%m-%d %H:%M:%S')}."
async def tools_on_agent_level() -> None:
"""Example showing tools defined when creating the agent."""
print("=== Tools Defined on Agent Level ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Tools are provided when creating the agent
# The agent can use these tools for any query during its lifetime
agent = await provider.create_agent(
name="InfoAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant that can provide weather and time information.",
tools=[get_weather, get_time], # Tools defined at agent creation
)
try:
# First query - agent can use weather tool
query1 = "What's the weather like in New York?"
print(f"User: {query1}")
result1 = await agent.run(query1)
print(f"Agent: {result1}\n")
# Second query - agent can use time tool
query2 = "What's the current UTC time?"
print(f"User: {query2}")
result2 = await agent.run(query2)
print(f"Agent: {result2}\n")
# Third query - agent can use both tools if needed
query3 = "What's the weather in London and what's the current UTC time?"
print(f"User: {query3}")
result3 = await agent.run(query3)
print(f"Agent: {result3}\n")
finally:
await client.beta.assistants.delete(agent.id)
async def tools_on_run_level() -> None:
"""Example showing tools passed to the run method."""
print("=== Tools Passed to Run Method ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Agent created with base tools, additional tools can be passed at run time
agent = await provider.create_agent(
name="FlexibleAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful assistant.",
tools=[get_weather], # Base tool
)
try:
# First query using base weather tool
query1 = "What's the weather like in Seattle?"
print(f"User: {query1}")
result1 = await agent.run(query1)
print(f"Agent: {result1}\n")
# Second query with additional time tool
query2 = "What's the current UTC time?"
print(f"User: {query2}")
result2 = await agent.run(query2, tools=[get_time]) # Additional tool for this query
print(f"Agent: {result2}\n")
# Third query with both tools
query3 = "What's the weather in Chicago and what's the current UTC time?"
print(f"User: {query3}")
result3 = await agent.run(query3, tools=[get_time]) # Time tool adds to weather
print(f"Agent: {result3}\n")
finally:
await client.beta.assistants.delete(agent.id)
async def mixed_tools_example() -> None:
"""Example showing both agent-level tools and run-method tools."""
print("=== Mixed Tools Example (Agent + Run Method) ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# Agent created with some base tools
agent = await provider.create_agent(
name="ComprehensiveAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a comprehensive assistant that can help with various information requests.",
tools=[get_weather], # Base tool available for all queries
)
try:
# Query using both agent tool and additional run-method tools
query = "What's the weather in Denver and what's the current UTC time?"
print(f"User: {query}")
# Agent has access to get_weather (from creation) + additional tools from run method
result = await agent.run(
query,
tools=[get_time], # Additional tools for this specific query
)
print(f"Agent: {result}\n")
finally:
await client.beta.assistants.delete(agent.id)
async def main() -> None:
print("=== OpenAI Assistants Provider with Function Tools Examples ===\n")
await tools_on_agent_level()
await tools_on_run_level()
await mixed_tools_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,96 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import BaseModel, ConfigDict
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistant Provider Response Format Example
This sample demonstrates using OpenAIAssistantProvider with response_format
for structured outputs in two ways:
1. Setting default response_format at agent creation time (default_options)
2. Overriding response_format at runtime (options parameter in agent.run)
"""
class WeatherInfo(BaseModel):
"""Structured weather information."""
location: str
temperature: int
conditions: str
recommendation: str
model_config = ConfigDict(extra="forbid")
class CityInfo(BaseModel):
"""Structured city information."""
city_name: str
population: int
country: str
model_config = ConfigDict(extra="forbid")
async def main() -> None:
"""Example of using response_format at creation time and runtime."""
async with (
AsyncOpenAI() as client,
OpenAIAssistantProvider(client) as provider,
):
# Create agent with default response_format (WeatherInfo)
agent = await provider.create_agent(
name="StructuredReporter",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="Return structured JSON based on the requested format.",
default_options={"response_format": WeatherInfo},
)
try:
# Request 1: Uses default response_format from agent creation
print("--- Request 1: Using default response_format (WeatherInfo) ---")
query1 = "What's the weather like in Paris today?"
print(f"User: {query1}")
result1 = await agent.run(query1)
try:
weather = result1.value
print("Agent:")
print(f" Location: {weather.location}")
print(f" Temperature: {weather.temperature}")
print(f" Conditions: {weather.conditions}")
print(f" Recommendation: {weather.recommendation}")
except Exception:
print(f"Failed to parse response: {result1.text}")
# Request 2: Override response_format at runtime with CityInfo
print("\n--- Request 2: Runtime override with CityInfo ---")
query2 = "Tell me about Tokyo."
print(f"User: {query2}")
result2 = await agent.run(query2, options={"response_format": CityInfo})
try:
city = result2.value
print("Agent:")
print(f" City: {city.city_name}")
print(f" Population: {city.population}")
print(f" Country: {city.country}")
except Exception:
print(f"Failed to parse response: {result2.text}")
finally:
await client.beta.assistants.delete(agent.id)
if __name__ == "__main__":
asyncio.run(main())
@@ -1,172 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
from random import randint
from typing import Annotated
from agent_framework import AgentSession, tool
from agent_framework.openai import OpenAIAssistantProvider
from dotenv import load_dotenv
from openai import AsyncOpenAI
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Assistants with Session Management Example
This sample demonstrates session management with OpenAI Assistants, showing
persistent conversation sessions and context preservation across interactions.
"""
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C."
async def example_with_automatic_session_creation() -> None:
"""Example showing automatic session creation (service-managed session)."""
print("=== Automatic Session Creation Example ===")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
try:
# First conversation - no session provided, will be created automatically
query1 = "What's the weather like in Seattle?"
print(f"User: {query1}")
result1 = await agent.run(query1)
print(f"Agent: {result1.text}")
# Second conversation - still no session provided, will create another new session
query2 = "What was the last city I asked about?"
print(f"\nUser: {query2}")
result2 = await agent.run(query2)
print(f"Agent: {result2.text}")
print("Note: Each call creates a separate session, so the agent doesn't remember previous context.\n")
finally:
await client.beta.assistants.delete(agent.id)
async def example_with_session_persistence() -> None:
"""Example showing session persistence across multiple conversations."""
print("=== Session Persistence Example ===")
print("Using the same session across multiple conversations to maintain context.\n")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
try:
# Create a new session that will be reused
session = agent.create_session()
# First conversation
query1 = "What's the weather like in Tokyo?"
print(f"User: {query1}")
result1 = await agent.run(query1, session=session)
print(f"Agent: {result1.text}")
# Second conversation using the same session - maintains context
query2 = "How about London?"
print(f"\nUser: {query2}")
result2 = await agent.run(query2, session=session)
print(f"Agent: {result2.text}")
# Third conversation - agent should remember both previous cities
query3 = "Which of the cities I asked about has better weather?"
print(f"\nUser: {query3}")
result3 = await agent.run(query3, session=session)
print(f"Agent: {result3.text}")
print("Note: The agent remembers context from previous messages in the same session.\n")
finally:
await client.beta.assistants.delete(agent.id)
async def example_with_existing_session_id() -> None:
"""Example showing how to work with an existing session ID from the service."""
print("=== Existing Session ID Example ===")
print("Using a specific session ID to continue an existing conversation.\n")
client = AsyncOpenAI()
provider = OpenAIAssistantProvider(client)
# First, create a conversation and capture the session ID
existing_session_id = None
assistant_id = None
agent = await provider.create_agent(
name="WeatherAssistant",
model=os.environ.get("OPENAI_MODEL", "gpt-4"),
instructions="You are a helpful weather agent.",
tools=[get_weather],
)
assistant_id = agent.id
try:
# Start a conversation and get the session ID
session = agent.create_session()
query1 = "What's the weather in Paris?"
print(f"User: {query1}")
result1 = await agent.run(query1, session=session)
print(f"Agent: {result1.text}")
# The session ID is set after the first response
existing_session_id = session.service_session_id
print(f"Session ID: {existing_session_id}")
if existing_session_id:
print("\n--- Continuing with the same session ID using get_agent ---")
# Get the existing assistant by ID
agent2 = await provider.get_agent(
assistant_id=assistant_id,
tools=[get_weather], # Must provide function implementations
)
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_session_id)
query2 = "What was the last city I asked about?"
print(f"User: {query2}")
result2 = await agent2.run(query2, session=session)
print(f"Agent: {result2.text}")
print("Note: The agent continues the conversation from the previous session.\n")
finally:
if assistant_id:
await client.beta.assistants.delete(assistant_id)
async def main() -> None:
print("=== OpenAI Assistants Provider Session Management Examples ===\n")
await example_with_automatic_session_creation()
await example_with_session_persistence()
await example_with_existing_session_id()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,132 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from collections.abc import Awaitable, Callable
from random import randint
from typing import Annotated
from agent_framework import (
Agent,
ChatContext,
ChatResponse,
Message,
MiddlewareTermination,
Role,
chat_middleware,
tool,
)
from agent_framework.openai import OpenAIResponsesClient
from dotenv import load_dotenv
from pydantic import Field
# Load environment variables from .env file
load_dotenv()
"""
OpenAI Responses Client Basic Example
This sample demonstrates basic usage of OpenAIResponsesClient for structured
response generation, showing both streaming and non-streaming responses.
"""
@chat_middleware
async def security_and_override_middleware(
context: ChatContext,
call_next: Callable[[], Awaitable[None]],
) -> None:
"""Function-based middleware that implements security filtering and response override."""
print("[SecurityMiddleware] Processing input...")
# Security check - block sensitive information
blocked_terms = ["password", "secret", "api_key", "token"]
for message in context.messages:
if message.text:
message_lower = message.text.lower()
for term in blocked_terms:
if term in message_lower:
print(f"[SecurityMiddleware] BLOCKED: Found '{term}' in message")
# Override the response instead of calling AI
context.result = ChatResponse(
messages=[
Message(
role=Role.ASSISTANT,
text="I cannot process requests containing sensitive information. "
"Please rephrase your question without including passwords, secrets, or other "
"sensitive data.",
)
]
)
# Terminate middleware execution with the blocked response
raise MiddlewareTermination(result=context.result)
# Continue to next middleware or AI execution
await call_next()
print("[SecurityMiddleware] Response generated.")
print(type(context.result))
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production;
# see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
async def non_streaming_example() -> None:
"""Example of non-streaming response (get the complete result at once)."""
print("=== Non-streaming Response Example ===")
agent = Agent(
client=OpenAIResponsesClient(),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Result: {result}\n")
async def streaming_example() -> None:
"""Example of streaming response (get results as they are generated)."""
print("=== Streaming Response Example ===")
agent = Agent(
client=OpenAIResponsesClient(
middleware=[security_and_override_middleware],
),
instructions="You are a helpful weather agent.",
tools=get_weather,
)
query = "What's the weather like in Portland?"
print(f"User: {query}")
print("Agent: ", end="", flush=True)
response = agent.run(query, stream=True)
async for chunk in response:
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
print(f"Final Result: {await response.get_final_response()}")
async def main() -> None:
print("=== Basic OpenAI Responses Client Agent Example ===")
await streaming_example()
await non_streaming_example()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,6 +1,6 @@
# OpenAI Configuration
OPENAI_API_KEY=
OPENAI_CHAT_MODEL_ID=
OPENAI_CHAT_MODEL=
# Agent 365 Agentic Authentication Configuration
USE_ANONYMOUS_MODE=
@@ -21,7 +21,7 @@ export USE_ANONYMOUS_MODE=True # set to false if using auth
# OpenAI
export OPENAI_API_KEY="..."
export OPENAI_CHAT_MODEL_ID="..."
export OPENAI_CHAT_MODEL="..."
```
## Installing Dependencies

Some files were not shown because too many files have changed in this diff Show More