Skip to content

fix: retry empty OpenAI streams with non-streaming invoke #7736

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wxg0103 merged 1 commit into from
Nov 9, 2026
Merged

fix: retry empty OpenAI streams with non-streaming invoke #7736

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions apps/application/tests/test_langchain_compatibility.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 = [
"你好",
Expand Down
9 changes: 9 additions & 0 deletions apps/models_provider/impl/base_chat_open_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from langchain_core.language_models import LanguageModelInput
from langchain_core.messages import (
AIMessage,
BaseMessage,
get_buffer_string,
BaseMessageChunk,
Expand Down Expand Up @@ -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,

@shaohuzhang1 shaohuzhang1 Nov 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's my review of the diff:

Potential Issues

1. Recursion risk (most important)

In _stream, when no chunks arrive, it calls self.invoke(..., stream=False). This is fine as long as invoke dispatches to _generate and not back to _stream. Verify that:

  • Self.invoke does not route through ChatModel.stream/_stream when stream=False.
  • Subclasses that override _stream don't end up calling this method again.

If any subclass or middleware calls .stream() while invoke() internally re-enters _stream, you'll get unbounded recursion.

2. AIMessageChunk(**message.model_dump(exclude={"type"}))

  • Excluding only "type" may not be sufficient: AIMessage can carry fields (id, usage_metadata, response_metadata, tool_calls, additional_kwargs, etc.) that don't map cleanly to AIMessageChunk, or whose shape differs (e.g., tool_calls should be tool_call_chunks). Safer conversion is suggested in LangChain practice: 
    AIMessageChunk(content=message.content,
               id=message.id,
               tool_call_chunks=message.tool_calls_to_chunks() if hasattr(...) else None,
               usage_metadata=message.usage_metadata,
               response_metadata=message.response_metadata,
               additional_kwargs=message.additional_kwargs)
    or use ChatGenerationChunk(message=AIMessageChunk(...)) with explicit mapping.
  • If AIMessage contains a type field set via Pydantic, model_dump() will include it; the exclude works but coupling the code to that one field name is fragile.

3. Duplicate semantics: yielding one "whole" chunk

Yielding a single non-streaming message as a ChatGenerationChunk violates the streaming contract for callers that expect incremental content (they will see one big chunk). This may break token-by-token UI consumers. Document this behavior, or consider that an empty stream before completion is more honest.

4. chunk = None typing / lint

  • chunk inferred as ChatGenerationChunk after assignment inside the loop; initializing to None makes the type Optional[ChatGenerationChunk]. Fine at runtime, but annotate explicitly: chunk: ChatGenerationChunk | None = None.
  • The loop variable shadows the sentinel, which is intentional here but easy to misread — consider renaming (last_chunk).

5. usage_metadata side-effect timing

self.usage_metadata is only set when chunk.message.usage_metadata exists. In the fallback path, the code doesn't update self.usage_metadata from the fallback message. Add:

if message.usage_metadata is not None:
    self.usage_metadata = message.usage_metadata

6. Sync vs async parity

Both branches mirror each other well. Ensure ainvoke(*args, **{**kwargs, "stream": False}) doesn't accidentally pass stream=False twice if caller already included stream=True/other values — the {**kwargs, "stream": False} handles this correctly, good.

7. Exception handling / partial streams

If the parent stream raises mid-way, the fallback block won't execute (good), but any partially accumulated state stays. Consider whether you want to still fire the fallback on unusual early termination. Probably not — document intent in a comment.

8. Import ordering / alphabetical convention

Minor: AIMessage inserted between BaseMessage-adjacent imports; many linters enforce isort ordering — check ruff/isort passes. Also several existing imports in the block are out of order (get_buffer_string before BaseMessageChunk); pre-existing, but worth noting.

9. Locking / generator finalization

These are generators; if the consumer abandons iteration after the first yielded chunk, chunk is None fallback never matters. No leak concern, just note the fallback only triggers on an empty completed stream.

10. Cost/latency

On truly empty streams you make a second API call (non-streaming). That doubles cost in this edge case. Log a warning so this doesn't go unnoticed, e.g.:

logger.warning("%s produced an empty stream; falling back to non-streaming invoke", type(self).__name__)

Optimization suggestions

  • Factor the common logic into a helper _finalize_empty_stream(message) to avoid sync/async duplication.
  • Prefer the model-class-provided converter (_convert_chunk_to_generation_chunk or a message_to_chunk() util) over model_dump.
  • Add unit tests for: empty stream, stream with usage_metadata only on last chunk, async variant, and both with tool calls (verify tool-call chunk conversion).

Summary

The gate if chunk is None is a reasonable guard against empty streaming responses, but the fallback has three real risks: (a) recursion if invoke re-enters _stream, (b) message→chunk conversion via model_dump(exclude={"type"}) may drop/mis-shape tool_calls, and (c) it silently doubles API calls without updating usage_metadata or logging. Fix those and it's straightforward.

Expand Down
Loading