Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
134 changes: 78 additions & 56 deletions libs/core/langchain_core/language_models/_compat_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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] = {}
Expand Down Expand Up @@ -633,39 +701,15 @@ 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)

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,
Expand All @@ -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] = {}
Expand Down Expand Up @@ -713,39 +756,17 @@ 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)

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,
Expand Down Expand Up @@ -816,6 +837,7 @@ async def amessage_to_events(


__all__ = [
"BlockStreamTracker",
"CompatBlock",
"achunks_to_events",
"amessage_to_events",
Expand Down
18 changes: 16 additions & 2 deletions libs/core/langchain_core/language_models/chat_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -667,8 +667,16 @@ def _iter_v2_events(
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(
Expand Down Expand Up @@ -698,8 +706,14 @@ async def _aiter_v2_events(
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(
Expand Down
28 changes: 28 additions & 0 deletions libs/core/langchain_core/language_models/stream_events.py
Original file line number Diff line number Diff line change
@@ -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",
]
Original file line number Diff line number Diff line change
@@ -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"
Loading