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:
Nick Hollon
2026-04-21 12:02:33 -04:00
parent 2f14205acb
commit 4062a8db89
3 changed files with 233 additions and 54 deletions

View File

@@ -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,
)

View File

@@ -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 = [

View File

@@ -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.