From 4062a8db8950ac876335527560809dae9024cec4 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Tue, 21 Apr 2026 12:02:33 -0400 Subject: [PATCH] fix(core): emit per-block lifecycle for string-keyed content blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streams where content_blocks carries string index identifiers (e.g. OpenAI responses/v1 mode with 'lc_rs_305f30', 'lc_txt_1') collapsed all blocks to wire index 0 because _iter_protocol_blocks fell back to positional i for non-int indices. Every block appeared as a delta of block 0, and only one content-block-finish event fired at the end of the stream. Keep the raw block key (int or string) internally, allocate sequential uint wire indices per distinct block, and finish the previously-open block when a new block key appears — matching the protocol's no-interleave rule. --- .../language_models/_compat_bridge.py | 140 +++++++++++------- .../language_models/test_compat_bridge.py | 92 ++++++++++++ .../chat_models/test_responses_stream.py | 55 +++++++ 3 files changed, 233 insertions(+), 54 deletions(-) diff --git a/libs/core/langchain_core/language_models/_compat_bridge.py b/libs/core/langchain_core/language_models/_compat_bridge.py index e1967563430..6ec1a988f95 100644 --- a/libs/core/langchain_core/language_models/_compat_bridge.py +++ b/libs/core/langchain_core/language_models/_compat_bridge.py @@ -108,11 +108,15 @@ def _to_finalized_block(block: CompatBlock) -> FinalizedContentBlock: # --------------------------------------------------------------------------- -def _iter_protocol_blocks(msg: BaseMessage) -> list[tuple[int, CompatBlock]]: +def _iter_protocol_blocks(msg: BaseMessage) -> list[tuple[Any, CompatBlock]]: """Read per-chunk protocol blocks from `msg.content_blocks`. - Returns `(index, block)` pairs. Block indices come from each - block's `index` field when present, falling back to positional. + Returns `(key, block)` pairs. The key is the block's stable identifier + across the stream: the block's `index` field when present (can be an + int or a string — some providers use string identifiers like + `"lc_rs_305f30"`), or the positional index within the message as a + fallback. Callers are responsible for allocating wire-level `uint` + indices; this helper only surfaces the source-side identity. For finalized :class:`AIMessage`, also surfaces `invalid_tool_calls` — which `AIMessage.content_blocks` currently omits from its return @@ -123,22 +127,21 @@ def _iter_protocol_blocks(msg: BaseMessage) -> list[tuple[int, CompatBlock]]: except Exception: return [] - result: list[tuple[int, CompatBlock]] = [] + result: list[tuple[Any, CompatBlock]] = [] for i, block in enumerate(raw): if not isinstance(block, dict): continue - raw_idx = block.get("index", i) - idx = raw_idx if isinstance(raw_idx, int) else i - result.append((idx, dict(block))) + key = block.get("index", i) + result.append((key, dict(block))) if not isinstance(msg, AIMessageChunk): # Finalized AIMessage: pull invalid_tool_calls from the dedicated # field — AIMessage.content_blocks does not currently include them. for itc in getattr(msg, "invalid_tool_calls", None) or []: itc_block: CompatBlock = {"type": "invalid_tool_call"} - for key in ("id", "name", "args", "error"): - if itc.get(key) is not None: - itc_block[key] = itc[key] + for key_name in ("id", "name", "args", "error"): + if itc.get(key_name) is not None: + itc_block[key_name] = itc[key_name] result.append((len(result), itc_block)) return result @@ -393,28 +396,23 @@ def _build_message_finish( return finish_data -def _finish_all_blocks( - state: dict[int, CompatBlock], -) -> tuple[list[MessagesData], bool]: - """Emit `content-block-finish` events for every open block. +def _finalize_and_build_finish( + wire_idx: int, + block: CompatBlock, +) -> tuple[MessagesData, bool]: + """Finalize a block and wrap it in a `content-block-finish` event. - Returns the event list plus a flag indicating whether any finalized - block was a valid `tool_call` (used for finish-reason inference). + Returns the event plus a flag indicating whether the finalized block + was a valid `tool_call` (used for finish-reason inference). """ - events: list[MessagesData] = [] - has_valid_tool_call = False - for idx in sorted(state): - finalized = _finalize_block(state[idx]) - if finalized.get("type") == "tool_call": - has_valid_tool_call = True - events.append( - ContentBlockFinishData( - event="content-block-finish", - index=idx, - content_block=finalized, - ) - ) - return events, has_valid_tool_call + finalized = _finalize_block(block) + has_valid_tool_call = finalized.get("type") == "tool_call" + event = ContentBlockFinishData( + event="content-block-finish", + index=wire_idx, + content_block=finalized, + ) + return event, has_valid_tool_call # --------------------------------------------------------------------------- @@ -429,6 +427,13 @@ def chunks_to_events( ) -> Iterator[MessagesData]: """Convert a stream of `ChatGenerationChunk` to protocol events. + Blocks stream one at a time: when a chunk carries a different block + identifier than the currently-open one, the open block is finished + before the new block starts, matching the protocol's no-interleave + rule. Source-side identifiers (from the block's `index` field, which + may be int or string) are translated to sequential `uint` wire + indices. + Args: chunks: Iterator of `ChatGenerationChunk` from `_stream()`. message_id: Optional stable message ID. @@ -437,8 +442,11 @@ def chunks_to_events( `MessagesData` lifecycle events. """ started = False - state: dict[int, CompatBlock] = {} - first_seen: set[int] = set() + open_key: Any = None + open_block: CompatBlock | None = None + open_wire_idx: int = 0 + next_wire_idx = 0 + has_valid_tool_call = False usage: dict[str, Any] | None = None response_metadata: dict[str, Any] = {} finish_reason: FinishReason = "stop" @@ -455,21 +463,29 @@ def chunks_to_events( started = True yield _build_message_start(msg, message_id) - for idx, block in _iter_protocol_blocks(msg): - if idx not in first_seen: - first_seen.add(idx) + for key, block in _iter_protocol_blocks(msg): + if key != open_key: + if open_block is not None: + event, tc = _finalize_and_build_finish(open_wire_idx, open_block) + has_valid_tool_call = has_valid_tool_call or tc + yield event + open_key = key + open_wire_idx = next_wire_idx + next_wire_idx += 1 + open_block = dict(block) yield ContentBlockStartData( event="content-block-start", - index=idx, + index=open_wire_idx, content_block=_start_skeleton(block), ) + else: + open_block = _accumulate(open_block, block) if _should_emit_delta(block): yield ContentBlockDeltaData( event="content-block-delta", - index=idx, + index=open_wire_idx, content_block=_to_protocol_block(block), ) - state[idx] = _accumulate(state.get(idx), block) if msg.usage_metadata: usage = _accumulate_usage(usage, msg.usage_metadata) @@ -482,8 +498,11 @@ def chunks_to_events( if not started: return - finish_events, has_valid_tool_call = _finish_all_blocks(state) - yield from finish_events + if open_block is not None: + event, tc = _finalize_and_build_finish(open_wire_idx, open_block) + has_valid_tool_call = has_valid_tool_call or tc + yield event + yield _build_message_finish( finish_reason=finish_reason, has_valid_tool_call=has_valid_tool_call, @@ -499,8 +518,11 @@ async def achunks_to_events( ) -> AsyncIterator[MessagesData]: """Async variant of :func:`chunks_to_events`.""" started = False - state: dict[int, CompatBlock] = {} - first_seen: set[int] = set() + open_key: Any = None + open_block: CompatBlock | None = None + open_wire_idx: int = 0 + next_wire_idx = 0 + has_valid_tool_call = False usage: dict[str, Any] | None = None response_metadata: dict[str, Any] = {} finish_reason: FinishReason = "stop" @@ -517,21 +539,29 @@ async def achunks_to_events( started = True yield _build_message_start(msg, message_id) - for idx, block in _iter_protocol_blocks(msg): - if idx not in first_seen: - first_seen.add(idx) + for key, block in _iter_protocol_blocks(msg): + if key != open_key: + if open_block is not None: + event, tc = _finalize_and_build_finish(open_wire_idx, open_block) + has_valid_tool_call = has_valid_tool_call or tc + yield event + open_key = key + open_wire_idx = next_wire_idx + next_wire_idx += 1 + open_block = dict(block) yield ContentBlockStartData( event="content-block-start", - index=idx, + index=open_wire_idx, content_block=_start_skeleton(block), ) + else: + open_block = _accumulate(open_block, block) if _should_emit_delta(block): yield ContentBlockDeltaData( event="content-block-delta", - index=idx, + index=open_wire_idx, content_block=_to_protocol_block(block), ) - state[idx] = _accumulate(state.get(idx), block) if msg.usage_metadata: usage = _accumulate_usage(usage, msg.usage_metadata) @@ -544,9 +574,11 @@ async def achunks_to_events( if not started: return - finish_events, has_valid_tool_call = _finish_all_blocks(state) - for event in finish_events: + if open_block is not None: + event, tc = _finalize_and_build_finish(open_wire_idx, open_block) + has_valid_tool_call = has_valid_tool_call or tc yield event + yield _build_message_finish( finish_reason=finish_reason, has_valid_tool_call=has_valid_tool_call, @@ -583,16 +615,16 @@ def message_to_events( yield _build_message_start(msg, message_id) has_valid_tool_call = False - for idx, block in _iter_protocol_blocks(msg): + for wire_idx, (_key, block) in enumerate(_iter_protocol_blocks(msg)): yield ContentBlockStartData( event="content-block-start", - index=idx, + index=wire_idx, content_block=_start_skeleton(block), ) if _should_emit_delta(block): yield ContentBlockDeltaData( event="content-block-delta", - index=idx, + index=wire_idx, content_block=_to_protocol_block(block), ) finalized = _finalize_block(block) @@ -600,7 +632,7 @@ def message_to_events( has_valid_tool_call = True yield ContentBlockFinishData( event="content-block-finish", - index=idx, + index=wire_idx, content_block=finalized, ) diff --git a/libs/core/tests/unit_tests/language_models/test_compat_bridge.py b/libs/core/tests/unit_tests/language_models/test_compat_bridge.py index 341f24f98d1..40a2288fa32 100644 --- a/libs/core/tests/unit_tests/language_models/test_compat_bridge.py +++ b/libs/core/tests/unit_tests/language_models/test_compat_bridge.py @@ -150,6 +150,98 @@ def test_chunks_to_events_empty_iterator() -> None: assert list(chunks_to_events(iter([]))) == [] +def test_chunks_to_events_block_transitions_close_previous_block() -> None: + """String-keyed blocks that transition mid-stream each get their own lifecycle. + + Regression test for OpenAI `responses/v1` style streams where + `content_blocks` uses string identifiers (e.g. `"lc_rs_305f30"`) to + distinguish blocks. Each distinct block must get its own + `content-block-start` / `content-block-finish` pair, with sequential + `uint` wire indices, and blocks must not interleave. + """ + chunks = [ + ChatGenerationChunk( + message=AIMessageChunk( + content=[ + {"type": "reasoning", "reasoning": "hmm", "index": "rs_a"}, + ], + id="msg-1", + ) + ), + ChatGenerationChunk( + message=AIMessageChunk( + content=[ + {"type": "reasoning", "reasoning": " then", "index": "rs_a"}, + ], + id="msg-1", + ) + ), + ChatGenerationChunk( + message=AIMessageChunk( + content=[ + {"type": "reasoning", "reasoning": "different", "index": "rs_b"}, + ], + id="msg-1", + ) + ), + ChatGenerationChunk( + message=AIMessageChunk( + content=[ + {"type": "text", "text": "answer: ", "index": "txt_1"}, + ], + id="msg-1", + ) + ), + ChatGenerationChunk( + message=AIMessageChunk( + content=[ + {"type": "text", "text": "42", "index": "txt_1"}, + ], + id="msg-1", + ) + ), + ] + + events = list(chunks_to_events(iter(chunks), message_id="msg-1")) + + starts = [e for e in events if e["event"] == "content-block-start"] + finishes = [e for e in events if e["event"] == "content-block-finish"] + assert [s["content_block"]["type"] for s in starts] == [ + "reasoning", + "reasoning", + "text", + ] + assert [f["content_block"]["type"] for f in finishes] == [ + "reasoning", + "reasoning", + "text", + ] + # Wire indices are sequential uints regardless of source-side keys. + assert [s["index"] for s in starts] == [0, 1, 2] + assert [f["index"] for f in finishes] == [0, 1, 2] + + # Finish events must be interleaved with starts (no-interleave rule): + # block 0 finishes before block 1 starts, etc. + lifecycle = [ + (e["event"], e["index"]) + for e in events + if e["event"] in ("content-block-start", "content-block-finish") + ] + assert lifecycle == [ + ("content-block-start", 0), + ("content-block-finish", 0), + ("content-block-start", 1), + ("content-block-finish", 1), + ("content-block-start", 2), + ("content-block-finish", 2), + ] + + # Each finish carries the accumulated content for its block. + assert finishes[0]["content_block"]["reasoning"] == "hmm then" + assert finishes[1]["content_block"]["reasoning"] == "different" + assert finishes[2]["content_block"]["text"] == "answer: 42" + + def test_chunks_to_events_tool_call_multichunk() -> None: """Partial tool-call args across chunks finalize to a single tool_call.""" chunks = [ diff --git a/libs/partners/openai/tests/unit_tests/chat_models/test_responses_stream.py b/libs/partners/openai/tests/unit_tests/chat_models/test_responses_stream.py index 6e9a3f52575..fe58fb8a1a5 100644 --- a/libs/partners/openai/tests/unit_tests/chat_models/test_responses_stream.py +++ b/libs/partners/openai/tests/unit_tests/chat_models/test_responses_stream.py @@ -762,6 +762,61 @@ def test_responses_stream(output_version: str, expected_content: list[dict]) -> assert dumped == payload["input"][idx] +def test_responses_stream_v2_emits_reasoning_lifecycle() -> None: + """`stream_v2` must emit `content-block-finish` events for reasoning blocks. + + Regression test: the protocol bridge should surface the full lifecycle + (`content-block-start` / `content-block-delta` / `content-block-finish`) + for every reasoning block observed on the wire, not just text blocks. + """ + llm = ChatOpenAI( + model="o4-mini", use_responses_api=True, output_version="v1" + ) + mock_client = MagicMock() + + def mock_create(*args: Any, **kwargs: Any) -> MockSyncContextManager: + return MockSyncContextManager(responses_stream) + + mock_client.responses.create = mock_create + + with patch.object(llm, "root_client", mock_client): + events = list(llm.stream_v2("test")) + + reasoning_starts = [ + e + for e in events + if e["event"] == "content-block-start" + and e["content_block"]["type"] == "reasoning" + ] + reasoning_finishes = [ + e + for e in events + if e["event"] == "content-block-finish" + and e["content_block"]["type"] == "reasoning" + ] + + # The mock stream carries four reasoning summary parts (two per reasoning + # item, across two reasoning items), which surface as four reasoning + # content blocks in `output_version="v1"`. + assert len(reasoning_starts) == 4, ( + f"expected 4 reasoning start events, got {len(reasoning_starts)}" + ) + assert len(reasoning_finishes) == 4, ( + f"expected 4 reasoning finish events, got {len(reasoning_finishes)}: " + f"all finish events = " + f"{[e['content_block']['type'] for e in events if e['event'] == 'content-block-finish']}" + ) + + # Finish events must carry the accumulated reasoning text. + reasoning_texts = [f["content_block"]["reasoning"] for f in reasoning_finishes] + assert reasoning_texts == [ + "reasoning block one", + "another reasoning block", + "more reasoning", + "still more reasoning", + ] + + def test_responses_stream_with_image_generation_multiple_calls() -> None: """Test that streaming with image_generation tool works across multiple calls.