mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-12 11:21:37 +00:00
fix(core): emit per-block lifecycle for string-keyed content blocks
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.
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user