Python: Reduce Azure chat client import overhead (#4744)

* Reduce Azure chat client import overhead

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

* Fix Azure chat client type annotations and add _parse_text_from_openai tests

- Move Choice and ChunkChoice imports under TYPE_CHECKING to avoid
  runtime import cost (from __future__ annotations is already present)
- Restore proper typed signature (Choice | ChunkChoice) instead of Any
- Add direct unit tests for _parse_text_from_openai covering:
  - Choice with message content
  - ChunkChoice with delta content
  - Refusal branch for both Choice and ChunkChoice
  - No content/no refusal returning None
  - None delta (async content filtering) returning None

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

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-18 09:05:42 +01:00
committed by GitHub
Unverified
parent c74b1b08eb
commit 192a283c9a
2 changed files with 90 additions and 6 deletions
@@ -8,9 +8,6 @@ import sys
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Generic, cast
from openai.lib.azure import AsyncAzureOpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from pydantic import BaseModel
from agent_framework import (
@@ -23,8 +20,7 @@ from agent_framework import (
FunctionInvocationLayer,
)
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai import OpenAIChatOptions
from agent_framework.openai._chat_client import RawOpenAIChatClient
from agent_framework.openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from .._settings import load_settings
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
@@ -48,6 +44,10 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from openai.lib.azure import AsyncAzureOpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
from agent_framework._middleware import MiddlewareTypes
logger: logging.Logger = logging.getLogger(__name__)
@@ -297,7 +297,9 @@ class AzureOpenAIChatClient( # type: ignore[misc]
For docs see:
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
"""
message = choice.message if isinstance(choice, Choice) else choice.delta
message = getattr(choice, "message", None)
if message is None:
message = getattr(choice, "delta", None)
# When you enable asynchronous content filtering in Azure OpenAI, you may receive empty deltas
if message is None: # type: ignore
return None
@@ -652,6 +652,88 @@ async def test_streaming_with_none_delta(
assert any(msg.contents for msg in results)
# region _parse_text_from_openai direct unit tests
def test_parse_text_from_openai_with_choice_message(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai correctly reads message from a Choice."""
client = AzureOpenAIChatClient()
choice = Choice(
index=0,
message=ChatCompletionMessage(content="hello", role="assistant"),
finish_reason="stop",
)
result = client._parse_text_from_openai(choice)
assert result is not None
assert result.type == "text"
assert result.text == "hello"
def test_parse_text_from_openai_with_chunk_choice_delta(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai correctly reads delta from a ChunkChoice."""
client = AzureOpenAIChatClient()
choice = ChunkChoice(
index=0,
delta=ChunkChoiceDelta(content="streamed", role="assistant"),
finish_reason=None,
)
result = client._parse_text_from_openai(choice)
assert result is not None
assert result.type == "text"
assert result.text == "streamed"
def test_parse_text_from_openai_refusal_choice(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai returns refusal text from a Choice."""
client = AzureOpenAIChatClient()
choice = Choice(
index=0,
message=ChatCompletionMessage(content=None, role="assistant", refusal="I cannot help with that"),
finish_reason="stop",
)
result = client._parse_text_from_openai(choice)
assert result is not None
assert result.type == "text"
assert result.text == "I cannot help with that"
def test_parse_text_from_openai_refusal_chunk_choice(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai returns refusal text from a ChunkChoice."""
client = AzureOpenAIChatClient()
choice = ChunkChoice(
index=0,
delta=ChunkChoiceDelta(content=None, role="assistant", refusal="I cannot help with that"),
finish_reason=None,
)
result = client._parse_text_from_openai(choice)
assert result is not None
assert result.type == "text"
assert result.text == "I cannot help with that"
def test_parse_text_from_openai_no_content_no_refusal(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai returns None when no content or refusal."""
client = AzureOpenAIChatClient()
choice = Choice(
index=0,
message=ChatCompletionMessage(content=None, role="assistant"),
finish_reason="stop",
)
result = client._parse_text_from_openai(choice)
assert result is None
def test_parse_text_from_openai_none_delta(azure_openai_unit_test_env: dict[str, str]) -> None:
"""Test _parse_text_from_openai returns None when delta is None (async content filtering)."""
client = AzureOpenAIChatClient()
choice = ChunkChoice.model_construct(index=0, delta=None, finish_reason=None)
result = client._parse_text_from_openai(choice)
assert result is None
# endregion
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
async def test_cmc_with_conversation_id(
mock_create: AsyncMock,