[shaohuzhang1]

fix: retry empty OpenAI streams with non-streaming invoke

[shaohuzhang1]

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.