From cc78501296b24840e06e3e25835516729778b4d7 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Tue, 9 Jun 2026 17:13:50 -0400 Subject: [PATCH] feat(core): extract reusable BlockStreamTracker from compat bridge --- .../language_models/_compat_bridge.py | 134 ++++++++++-------- .../language_models/chat_models.py | 18 ++- .../language_models/stream_events.py | 28 ++++ .../test_block_stream_tracker.py | 65 +++++++++ 4 files changed, 187 insertions(+), 58 deletions(-) create mode 100644 libs/core/langchain_core/language_models/stream_events.py create mode 100644 libs/core/tests/unit_tests/language_models/test_block_stream_tracker.py diff --git a/libs/core/langchain_core/language_models/_compat_bridge.py b/libs/core/langchain_core/language_models/_compat_bridge.py index 43490f0717c..67626c86cbb 100644 --- a/libs/core/langchain_core/language_models/_compat_bridge.py +++ b/libs/core/langchain_core/language_models/_compat_bridge.py @@ -563,6 +563,75 @@ def _finalize_and_build_finish( ) +# --------------------------------------------------------------------------- +# Shared block-lifecycle tracker +# --------------------------------------------------------------------------- + + +class BlockStreamTracker: + """Allocate wire indices and emit content-block lifecycle events. + + Single source of truth for the start/delta/accumulate/finalize + mechanics shared by the compat bridge and provider-native converters. + Source-side keys (a block's `index` field, int or str, or a + positional fallback) are mapped to sequential `uint` wire indices. + """ + + def __init__(self) -> None: + self._blocks: dict[Any, tuple[int, CompatBlock]] = {} + self._next_wire_idx = 0 + + def feed(self, key: Any, block: CompatBlock) -> Iterator[MessagesData]: + """Yield lifecycle events for a per-chunk block. + + Yields `content-block-start` the first time `key` is seen, then + `content-block-delta` when the block carries fresh content, + accumulating per-index state for later finalization. + """ + if key not in self._blocks: + wire_idx = self._next_wire_idx + self._next_wire_idx += 1 + self._blocks[key] = (wire_idx, dict(block)) + yield ContentBlockStartData( + event="content-block-start", + index=wire_idx, + content=_start_skeleton(block), + ) + else: + wire_idx, existing = self._blocks[key] + self._blocks[key] = (wire_idx, _accumulate(existing, block)) + if _should_emit_delta(block): + wire_idx, current = self._blocks[key] + is_block_delta = block.get("type") in ( + "tool_call_chunk", + "server_tool_call_chunk", + ) + delta_source = current if is_block_delta else block + yield ContentBlockDeltaData( + event="content-block-delta", + index=wire_idx, + delta=_to_content_delta(delta_source or block), + ) + + def finish_block(self, key: Any) -> Iterator[MessagesData]: + """Finalize and close one block at its true stop boundary. + + Native converters call this on the provider's block-stop signal. + A no-op for unknown / already-closed keys. + """ + entry = self._blocks.pop(key, None) + if entry is None: + return + wire_idx, block = entry + yield _finalize_and_build_finish(wire_idx, block) + + def finish_all(self) -> Iterator[MessagesData]: + """Finalize every still-open block, in wire-index order.""" + for wire_idx, block in self._blocks.values(): + yield _finalize_and_build_finish(wire_idx, block) + self._blocks.clear() + + # --------------------------------------------------------------------------- # Main generators # --------------------------------------------------------------------------- @@ -590,8 +659,7 @@ def chunks_to_events( `MessagesData` lifecycle events. """ started = False - blocks: dict[Any, tuple[int, CompatBlock]] = {} - next_wire_idx = 0 + tracker = BlockStreamTracker() usage: dict[str, Any] | None = None response_metadata: dict[str, Any] = {} additional_kwargs: dict[str, Any] = {} @@ -633,30 +701,7 @@ def chunks_to_events( yield _build_message_start(msg, message_id) for key, block in _iter_protocol_blocks(msg): - if key not in blocks: - wire_idx = next_wire_idx - next_wire_idx += 1 - blocks[key] = (wire_idx, dict(block)) - yield ContentBlockStartData( - event="content-block-start", - index=wire_idx, - content=_start_skeleton(block), - ) - else: - wire_idx, existing = blocks[key] - blocks[key] = (wire_idx, _accumulate(existing, block)) - if _should_emit_delta(block): - wire_idx, current = blocks[key] - is_block_delta = block.get("type") in ( - "tool_call_chunk", - "server_tool_call_chunk", - ) - delta_source = current if is_block_delta else block - yield ContentBlockDeltaData( - event="content-block-delta", - index=wire_idx, - delta=_to_content_delta(delta_source or block), - ) + yield from tracker.feed(key, block) if msg.usage_metadata: usage = _accumulate_usage(usage, msg.usage_metadata) @@ -664,8 +709,7 @@ def chunks_to_events( if not started: return - for wire_idx, block in blocks.values(): - yield _finalize_and_build_finish(wire_idx, block) + yield from tracker.finish_all() yield _build_message_finish( usage=usage, @@ -681,8 +725,7 @@ async def achunks_to_events( ) -> AsyncIterator[MessagesData]: """Async variant of `chunks_to_events`.""" started = False - blocks: dict[Any, tuple[int, CompatBlock]] = {} - next_wire_idx = 0 + tracker = BlockStreamTracker() usage: dict[str, Any] | None = None response_metadata: dict[str, Any] = {} additional_kwargs: dict[str, Any] = {} @@ -713,30 +756,8 @@ async def achunks_to_events( yield _build_message_start(msg, message_id) for key, block in _iter_protocol_blocks(msg): - if key not in blocks: - wire_idx = next_wire_idx - next_wire_idx += 1 - blocks[key] = (wire_idx, dict(block)) - yield ContentBlockStartData( - event="content-block-start", - index=wire_idx, - content=_start_skeleton(block), - ) - else: - wire_idx, existing = blocks[key] - blocks[key] = (wire_idx, _accumulate(existing, block)) - if _should_emit_delta(block): - wire_idx, current = blocks[key] - is_block_delta = block.get("type") in ( - "tool_call_chunk", - "server_tool_call_chunk", - ) - delta_source = current if is_block_delta else block - yield ContentBlockDeltaData( - event="content-block-delta", - index=wire_idx, - delta=_to_content_delta(delta_source or block), - ) + for event in tracker.feed(key, block): + yield event if msg.usage_metadata: usage = _accumulate_usage(usage, msg.usage_metadata) @@ -744,8 +765,8 @@ async def achunks_to_events( if not started: return - for wire_idx, block in blocks.values(): - yield _finalize_and_build_finish(wire_idx, block) + for event in tracker.finish_all(): + yield event yield _build_message_finish( usage=usage, @@ -816,6 +837,7 @@ async def amessage_to_events( __all__ = [ + "BlockStreamTracker", "CompatBlock", "achunks_to_events", "amessage_to_events", diff --git a/libs/core/langchain_core/language_models/chat_models.py b/libs/core/langchain_core/language_models/chat_models.py index 7c21d12110f..1884417c89a 100644 --- a/libs/core/langchain_core/language_models/chat_models.py +++ b/libs/core/langchain_core/language_models/chat_models.py @@ -667,8 +667,16 @@ class BaseChatModel(BaseLanguageModel[AIMessage], ABC): getattr(self, "_stream_chat_model_events", None), ) if native is not None: + # Thread the stream's message id (the LangChain run id) into the + # native producer so its `message-start` carries the same id the + # bridge path emits — keeping the two paths consistent for + # consumers that correlate v3 events by `message-start.id`. event_iter: Iterator[MessagesData] = native( - messages, stop=stop, run_manager=run_manager, **kwargs + messages, + stop=stop, + run_manager=run_manager, + message_id=stream.message_id, + **kwargs, ) else: event_iter = chunks_to_events( @@ -698,8 +706,14 @@ class BaseChatModel(BaseLanguageModel[AIMessage], ABC): getattr(self, "_astream_chat_model_events", None), ) if native is not None: + # See `_iter_v2_events`: thread the stream's message id into the + # native producer for `message-start` consistency with the bridge. event_iter: AsyncIterator[MessagesData] = native( - messages, stop=stop, run_manager=run_manager, **kwargs + messages, + stop=stop, + run_manager=run_manager, + message_id=stream.message_id, + **kwargs, ) else: event_iter = achunks_to_events( diff --git a/libs/core/langchain_core/language_models/stream_events.py b/libs/core/langchain_core/language_models/stream_events.py new file mode 100644 index 00000000000..b40351ee8f1 --- /dev/null +++ b/libs/core/langchain_core/language_models/stream_events.py @@ -0,0 +1,28 @@ +"""Shared primitives for provider-native streaming-event converters. + +Provider packages implementing `_stream_chat_model_events` import these +to drive raw API events into the same content-block lifecycle the compat +bridge produces, so native and bridged paths emit identical event shapes. +""" + +from langchain_core.language_models._compat_bridge import ( + BlockStreamTracker, + finalize_tool_call_chunk, +) +from langchain_core.language_models._compat_bridge import ( + _accumulate_usage as accumulate_usage, +) +from langchain_core.language_models._compat_bridge import ( + _build_message_finish as build_message_finish, +) +from langchain_core.language_models._compat_bridge import ( + _iter_protocol_blocks as iter_protocol_blocks, +) + +__all__ = [ + "BlockStreamTracker", + "accumulate_usage", + "build_message_finish", + "finalize_tool_call_chunk", + "iter_protocol_blocks", +] diff --git a/libs/core/tests/unit_tests/language_models/test_block_stream_tracker.py b/libs/core/tests/unit_tests/language_models/test_block_stream_tracker.py new file mode 100644 index 00000000000..4dc47907338 --- /dev/null +++ b/libs/core/tests/unit_tests/language_models/test_block_stream_tracker.py @@ -0,0 +1,65 @@ +"""Unit tests for the shared BlockStreamTracker.""" + +from typing import Any + +from langchain_core.language_models.stream_events import BlockStreamTracker + + +def test_feed_text_emits_start_then_delta() -> None: + tracker = BlockStreamTracker() + events: list[Any] = list( + tracker.feed(0, {"type": "text", "text": "Hi", "index": 0}) + ) + assert events[0]["event"] == "content-block-start" + assert events[0]["index"] == 0 + assert events[0]["content"] == {"type": "text", "text": ""} + assert events[1]["event"] == "content-block-delta" + assert events[1]["index"] == 0 + assert events[1]["delta"] == {"type": "text-delta", "text": "Hi"} + + +def test_feed_allocates_sequential_wire_indices() -> None: + tracker = BlockStreamTracker() + list(tracker.feed("a", {"type": "text", "text": "x", "index": "a"})) + second: list[Any] = list( + tracker.feed("b", {"type": "text", "text": "y", "index": "b"}) + ) + assert second[0]["index"] == 1 # second distinct key -> wire index 1 + + +def test_finish_block_finalizes_tool_call_at_boundary() -> None: + tracker = BlockStreamTracker() + list( + tracker.feed( + 0, + { + "type": "tool_call_chunk", + "id": "t1", + "name": "f", + "args": '{"a":', + "index": 0, + }, + ) + ) + list(tracker.feed(0, {"type": "tool_call_chunk", "args": "1}", "index": 0})) + finished: list[Any] = list(tracker.finish_block(0)) + assert len(finished) == 1 + assert finished[0]["event"] == "content-block-finish" + assert finished[0]["content"]["type"] == "tool_call" + assert finished[0]["content"]["args"] == {"a": 1} + # finished block is closed: finish_all must not re-emit it + assert list(tracker.finish_all()) == [] + + +def test_finish_block_unknown_key_is_noop() -> None: + assert list(BlockStreamTracker().finish_block("nope")) == [] + + +def test_finish_all_finalizes_open_blocks_in_wire_order() -> None: + tracker = BlockStreamTracker() + list(tracker.feed(0, {"type": "text", "text": "hello", "index": 0})) + list(tracker.feed(1, {"type": "reasoning", "reasoning": "why", "index": 1})) + finished: list[Any] = list(tracker.finish_all()) + assert [f["index"] for f in finished] == [0, 1] + assert finished[0]["content"]["text"] == "hello" + assert finished[1]["content"]["reasoning"] == "why"