openai[patch]: support Responses streaming in AzureChatOpenAI (#31641)

Resolves #31303,
#31624

Author: ccurme

libs/partners/openai/langchain_openai/chat_models/azure.py

@@ -4,14 +4,14 @@
 
 import logging
 import os
-from collections.abc import Awaitable
+from collections.abc import AsyncIterator, Awaitable, Iterator
 from typing import Any, Callable, Optional, TypedDict, TypeVar, Union
 
 import openai
 from langchain_core.language_models import LanguageModelInput
 from langchain_core.language_models.chat_models import LangSmithParams
 from langchain_core.messages import BaseMessage
-from langchain_core.outputs import ChatResult
+from langchain_core.outputs import ChatGenerationChunk, ChatResult
 from langchain_core.runnables import Runnable
 from langchain_core.utils import from_env, secret_from_env
 from langchain_core.utils.pydantic import is_basemodel_subclass
@@ -736,6 +736,24 @@ def _create_chat_result(
 
         return chat_result
 
+    def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGenerationChunk]:
+        """Route to Chat Completions or Responses API."""
+        if self._use_responses_api({**kwargs, **self.model_kwargs}):
+            return super()._stream_responses(*args, **kwargs)
+        else:
+            return super()._stream(*args, **kwargs)
+
+    async def _astream(
+        self, *args: Any, **kwargs: Any
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        """Route to Chat Completions or Responses API."""
+        if self._use_responses_api({**kwargs, **self.model_kwargs}):
+            async for chunk in super()._astream_responses(*args, **kwargs):
+                yield chunk
+        else:
+            async for chunk in super()._astream(*args, **kwargs):
+                yield chunk
+
     def with_structured_output(
         self,
         schema: Optional[_DictOrPydanticClass] = None,

libs/partners/openai/tests/integration_tests/chat_models/test_azure_standard.py

@@ -2,6 +2,7 @@
 
 import os
 
+import pytest
 from langchain_core.language_models import BaseChatModel
 from langchain_tests.integration_tests import ChatModelIntegrationTests
 
@@ -39,6 +40,38 @@ def supports_json_mode(self) -> bool:
         return True
 
 
+class TestAzureOpenAIResponses(ChatModelIntegrationTests):
+    @property
+    def chat_model_class(self) -> type[BaseChatModel]:
+        return AzureChatOpenAI
+
+    @property
+    def chat_model_params(self) -> dict:
+        return {
+            "deployment_name": os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+            "model": "gpt-4o-mini",
+            "openai_api_version": OPENAI_API_VERSION,
+            "azure_endpoint": OPENAI_API_BASE,
+            "use_responses_api": True,
+        }
+
+    @property
+    def supports_image_inputs(self) -> bool:
+        return True
+
+    @property
+    def supports_image_urls(self) -> bool:
+        return True
+
+    @property
+    def supports_json_mode(self) -> bool:
+        return True
+
+    @pytest.mark.xfail(reason="Unsupported.")
+    def test_stop_sequence(self, model: BaseChatModel) -> None:
+        super().test_stop_sequence(model)
+
+
 class TestAzureOpenAIStandardLegacy(ChatModelIntegrationTests):
     """Test a legacy model."""