diff --git a/python/packages/core/agent_framework/azure/_chat_client.py b/python/packages/core/agent_framework/azure/_chat_client.py index 21c38f6b57..2ae21d124c 100644 --- a/python/packages/core/agent_framework/azure/_chat_client.py +++ b/python/packages/core/agent_framework/azure/_chat_client.py @@ -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 diff --git a/python/packages/core/tests/azure/test_azure_chat_client.py b/python/packages/core/tests/azure/test_azure_chat_client.py index dd50c48db4..80902de0af 100644 --- a/python/packages/core/tests/azure/test_azure_chat_client.py +++ b/python/packages/core/tests/azure/test_azure_chat_client.py @@ -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,