diff --git a/apps/application/tests/test_langchain_compatibility.py b/apps/application/tests/test_langchain_compatibility.py index ffd29b409a1..b0f107e0110 100644 --- a/apps/application/tests/test_langchain_compatibility.py +++ b/apps/application/tests/test_langchain_compatibility.py @@ -16,6 +16,8 @@ from botocore.exceptions import ClientError from django.test import SimpleTestCase from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage, ToolMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_openai import ChatOpenAI from openai import RateLimitError from application.workflow.backend import sandbox_shell @@ -61,6 +63,52 @@ def llm_node(): class LangChainMessageCompatibilityTests(SimpleTestCase): + def test_empty_stream_falls_back_to_invoke(self): + model = openai_model(streaming=True) + expected = AIMessage(content="fallback") + + def empty_stream(*args, **kwargs): + return + yield + + with ( + patch.object(ChatOpenAI, "_stream", new=empty_stream), + patch.object( + ChatOpenAI, + "_generate", + return_value=ChatResult(generations=[ChatGeneration(message=expected)]), + ) as generate, + ): + chunks = list(model.stream([HumanMessage(content="hello")])) + + self.assertEqual("".join(chunk.content for chunk in chunks), "fallback") + generate.assert_called_once() + + def test_empty_async_stream_falls_back_to_ainvoke(self): + model = openai_model(streaming=True) + expected = AIMessage(content="fallback") + + async def empty_stream(*args, **kwargs): + return + yield + + with ( + patch.object(ChatOpenAI, "_astream", new=empty_stream), + patch.object( + ChatOpenAI, + "_agenerate", + new=AsyncMock(return_value=ChatResult(generations=[ChatGeneration(message=expected)])), + ) as agenerate, + ): + chunks = asyncio.run(self._collect_async_stream(model)) + + self.assertEqual("".join(chunk.content for chunk in chunks), "fallback") + agenerate.assert_awaited_once() + + @staticmethod + async def _collect_async_stream(model): + return [chunk async for chunk in model.astream([HumanMessage(content="hello")])] + async def test_complete_messages_preserve_text_tools_metadata_and_input(self): contents = [ "你好", diff --git a/apps/models_provider/impl/base_chat_open_ai.py b/apps/models_provider/impl/base_chat_open_ai.py index f35641257f8..3ecc911a509 100644 --- a/apps/models_provider/impl/base_chat_open_ai.py +++ b/apps/models_provider/impl/base_chat_open_ai.py @@ -5,6 +5,7 @@ from langchain_core.language_models import LanguageModelInput from langchain_core.messages import ( + AIMessage, BaseMessage, get_buffer_string, BaseMessageChunk, @@ -130,16 +131,24 @@ def get_num_tokens(self, text: str) -> int: ) def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGenerationChunk]: + chunk = None for chunk in super()._stream(*args, **kwargs): if chunk.message.usage_metadata is not None: self.usage_metadata = chunk.message.usage_metadata yield chunk + if chunk is None: + message = cast(AIMessage, self.invoke(*args, **{**kwargs, "stream": False})) + yield ChatGenerationChunk(message=AIMessageChunk(**message.model_dump(exclude={"type"}))) async def _astream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ChatGenerationChunk]: + chunk = None async for chunk in super()._astream(*args, **kwargs): if chunk.message.usage_metadata is not None: self.usage_metadata = chunk.message.usage_metadata yield chunk + if chunk is None: + message = cast(AIMessage, await self.ainvoke(*args, **{**kwargs, "stream": False})) + yield ChatGenerationChunk(message=AIMessageChunk(**message.model_dump(exclude={"type"}))) def _convert_chunk_to_generation_chunk( self,