fix(backends/anyllm): stream tool_use blocks and map finish_reason on the streaming path
Preserve AnyLLM streaming tool calls and finish reasons.
This commit is contained in:
Abhay Singh
+111
-22
@@ -362,42 +362,131 @@ class AnyLLMBackend(Backend):
|
||||
},
|
||||
)
|
||||
|
||||
yield StreamEvent(
|
||||
event_type="content_block_start",
|
||||
data={
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
)
|
||||
|
||||
stream_response = await self.llm.acompletion(**kwargs)
|
||||
output_tokens = 0
|
||||
# Stream text immediately in a single text block, but BUFFER tool
|
||||
# calls and emit them as complete blocks at the end. OpenAI streams
|
||||
# parallel tool calls interleaved by index (index 0 and 1 introduced
|
||||
# together, then a fragment for 0, then for 1), while Anthropic
|
||||
# requires each content block to be fully emitted — start, deltas,
|
||||
# stop — before the next opens. Reassembling per index and flushing
|
||||
# complete blocks keeps every delta inside its own block's start/stop
|
||||
# for any interleaving. (The previous version pre-opened one text
|
||||
# block and dropped tool calls entirely; a naive open-on-new-index
|
||||
# instead mis-sequenced parallel calls, emitting a fragment for an
|
||||
# already-stopped block.)
|
||||
current_block_index = -1
|
||||
text_block_open = False
|
||||
# provider tool index -> {"id", "name", "arguments"}, first-seen order
|
||||
tool_calls: dict[int, dict[str, Any]] = {}
|
||||
tool_order: list[int] = []
|
||||
stop_reason = "end_turn"
|
||||
|
||||
async for chunk in cast(AsyncIterator[Any], stream_response):
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
delta = chunk.choices[0].delta
|
||||
if hasattr(delta, "content") and delta.content:
|
||||
if not (hasattr(chunk, "choices") and chunk.choices):
|
||||
continue
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# Map OpenAI finish_reason to the Anthropic stop_reason so a tool
|
||||
# call or a length truncation is not reported as end_turn.
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
if finish_reason == "tool_calls":
|
||||
stop_reason = "tool_use"
|
||||
elif finish_reason == "length":
|
||||
stop_reason = "max_tokens"
|
||||
elif finish_reason == "stop":
|
||||
stop_reason = "end_turn"
|
||||
|
||||
if getattr(delta, "tool_calls", None):
|
||||
for tc in delta.tool_calls:
|
||||
idx = tc.index if getattr(tc, "index", None) is not None else 0
|
||||
buf = tool_calls.get(idx)
|
||||
if buf is None:
|
||||
buf = {"id": None, "name": "", "arguments": ""}
|
||||
tool_calls[idx] = buf
|
||||
tool_order.append(idx)
|
||||
if getattr(tc, "id", None):
|
||||
buf["id"] = tc.id
|
||||
func = getattr(tc, "function", None)
|
||||
if func is not None:
|
||||
if getattr(func, "name", None):
|
||||
buf["name"] = func.name
|
||||
if getattr(func, "arguments", None):
|
||||
buf["arguments"] += func.arguments
|
||||
|
||||
elif getattr(delta, "content", None):
|
||||
if not text_block_open:
|
||||
current_block_index += 1
|
||||
text_block_open = True
|
||||
yield StreamEvent(
|
||||
event_type="content_block_delta",
|
||||
event_type="content_block_start",
|
||||
data={
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "text_delta", "text": delta.content},
|
||||
"type": "content_block_start",
|
||||
"index": current_block_index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
)
|
||||
output_tokens += 1
|
||||
yield StreamEvent(
|
||||
event_type="content_block_delta",
|
||||
data={
|
||||
"type": "content_block_delta",
|
||||
"index": current_block_index,
|
||||
"delta": {"type": "text_delta", "text": delta.content},
|
||||
},
|
||||
)
|
||||
output_tokens += 1
|
||||
|
||||
yield StreamEvent(
|
||||
event_type="content_block_stop",
|
||||
data={"type": "content_block_stop", "index": 0},
|
||||
)
|
||||
# Close the text block before any tool blocks (Anthropic orders
|
||||
# content blocks sequentially, text then tool_use).
|
||||
if text_block_open:
|
||||
yield StreamEvent(
|
||||
event_type="content_block_stop",
|
||||
data={"type": "content_block_stop", "index": current_block_index},
|
||||
)
|
||||
|
||||
# Flush each buffered tool call as a complete, self-contained block:
|
||||
# start, one input_json_delta with the reassembled arguments, stop.
|
||||
for idx in tool_order:
|
||||
buf = tool_calls[idx]
|
||||
current_block_index += 1
|
||||
tool_id = buf["id"] or f"toolu_{uuid.uuid4().hex[:24]}"
|
||||
yield StreamEvent(
|
||||
event_type="content_block_start",
|
||||
data={
|
||||
"type": "content_block_start",
|
||||
"index": current_block_index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tool_id,
|
||||
"name": buf["name"],
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
if buf["arguments"]:
|
||||
yield StreamEvent(
|
||||
event_type="content_block_delta",
|
||||
data={
|
||||
"type": "content_block_delta",
|
||||
"index": current_block_index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": buf["arguments"],
|
||||
},
|
||||
},
|
||||
)
|
||||
output_tokens += 1
|
||||
yield StreamEvent(
|
||||
event_type="content_block_stop",
|
||||
data={"type": "content_block_stop", "index": current_block_index},
|
||||
)
|
||||
|
||||
yield StreamEvent(
|
||||
event_type="message_delta",
|
||||
data={
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": {"output_tokens": output_tokens},
|
||||
},
|
||||
)
|
||||
|
||||
@@ -419,6 +419,173 @@ async def test_stream_message_yields_events_and_error(monkeypatch: pytest.Monkey
|
||||
assert error_events[-1].data["error"]["message"] == "stream broke"
|
||||
|
||||
|
||||
def _tool_call_delta(*, index, tc_id=None, name=None, arguments=None): # noqa: ANN001, ANN202
|
||||
"""Build an OpenAI-style streaming tool_call delta chunk."""
|
||||
func = SimpleNamespace(name=name, arguments=arguments)
|
||||
tc = SimpleNamespace(index=index, id=tc_id, function=func)
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(tool_calls=[tc]), finish_reason=None)]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_message_emits_tool_use_blocks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A tool call streamed over any-llm must surface as an Anthropic tool_use block.
|
||||
|
||||
Regression: the streamer only handled text deltas, so ``tools`` were
|
||||
forwarded upstream but any tool call the model streamed back was dropped and
|
||||
the client saw an empty turn with stop_reason=end_turn. The block must open,
|
||||
stream its arguments as input_json_delta, and the turn must end tool_use.
|
||||
"""
|
||||
backend, instance = make_backend(monkeypatch)
|
||||
instance.response = FakeAsyncStream(
|
||||
[
|
||||
_tool_call_delta(index=0, tc_id="call_abc", name="get_weather"),
|
||||
_tool_call_delta(index=0, arguments='{"city":'),
|
||||
_tool_call_delta(index=0, arguments='"paris"}'),
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(), finish_reason="tool_calls")]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
events = [
|
||||
event async for event in backend.stream_message({"model": "claude", "messages": []}, {})
|
||||
]
|
||||
types = [e.event_type for e in events]
|
||||
|
||||
# The tool call is buffered and flushed as one complete block: start, a
|
||||
# single input_json_delta with the reassembled arguments, then stop.
|
||||
assert types == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
|
||||
start = next(e for e in events if e.event_type == "content_block_start")
|
||||
assert start.data["content_block"]["type"] == "tool_use"
|
||||
assert start.data["content_block"]["id"] == "call_abc"
|
||||
assert start.data["content_block"]["name"] == "get_weather"
|
||||
|
||||
arg_deltas = [e for e in events if e.event_type == "content_block_delta"]
|
||||
assert [d.data["delta"]["type"] for d in arg_deltas] == ["input_json_delta"]
|
||||
joined = "".join(d.data["delta"]["partial_json"] for d in arg_deltas)
|
||||
assert joined == '{"city":"paris"}'
|
||||
|
||||
message_delta = next(e for e in events if e.event_type == "message_delta")
|
||||
assert message_delta.data["delta"]["stop_reason"] == "tool_use"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_message_handles_parallel_tool_calls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Interleaved parallel tool calls must produce valid, disjoint Anthropic blocks.
|
||||
|
||||
OpenAI can introduce two tool indices in one chunk and then stream argument
|
||||
fragments for each across later chunks. Each Anthropic tool_use block must be
|
||||
fully framed (exactly one start and stop, arguments reassembled) with no
|
||||
delta emitted after that block's stop.
|
||||
"""
|
||||
backend, instance = make_backend(monkeypatch)
|
||||
instance.response = FakeAsyncStream(
|
||||
[
|
||||
# One chunk introduces BOTH tool indices at once.
|
||||
SimpleNamespace(
|
||||
choices=[
|
||||
SimpleNamespace(
|
||||
delta=SimpleNamespace(
|
||||
tool_calls=[
|
||||
SimpleNamespace(
|
||||
index=0,
|
||||
id="call_0",
|
||||
function=SimpleNamespace(name="alpha", arguments='{"a":'),
|
||||
),
|
||||
SimpleNamespace(
|
||||
index=1,
|
||||
id="call_1",
|
||||
function=SimpleNamespace(name="beta", arguments='{"b":'),
|
||||
),
|
||||
]
|
||||
),
|
||||
finish_reason=None,
|
||||
)
|
||||
]
|
||||
),
|
||||
# Interleaved argument fragments: index 0, then index 1.
|
||||
_tool_call_delta(index=0, arguments="1}"),
|
||||
_tool_call_delta(index=1, arguments="2}"),
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(), finish_reason="tool_calls")]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
events = [
|
||||
event async for event in backend.stream_message({"model": "claude", "messages": []}, {})
|
||||
]
|
||||
|
||||
# Each tool block index must have exactly one start and one stop, and no
|
||||
# delta may appear after that index's stop.
|
||||
stopped: set[int] = set()
|
||||
starts: dict[int, int] = {}
|
||||
stops: dict[int, int] = {}
|
||||
args: dict[int, str] = {}
|
||||
for e in events:
|
||||
if e.event_type == "content_block_start":
|
||||
idx = e.data["index"]
|
||||
starts[idx] = starts.get(idx, 0) + 1
|
||||
assert e.data["content_block"]["type"] == "tool_use"
|
||||
elif e.event_type == "content_block_delta":
|
||||
idx = e.data["index"]
|
||||
assert idx not in stopped, f"delta for block {idx} after its stop"
|
||||
args[idx] = args.get(idx, "") + e.data["delta"]["partial_json"]
|
||||
elif e.event_type == "content_block_stop":
|
||||
idx = e.data["index"]
|
||||
stops[idx] = stops.get(idx, 0) + 1
|
||||
stopped.add(idx)
|
||||
|
||||
assert starts == {0: 1, 1: 1}
|
||||
assert stops == {0: 1, 1: 1}
|
||||
assert args == {0: '{"a":1}', 1: '{"b":2}'}
|
||||
|
||||
block0 = next(
|
||||
e for e in events if e.event_type == "content_block_start" and e.data["index"] == 0
|
||||
)
|
||||
block1 = next(
|
||||
e for e in events if e.event_type == "content_block_start" and e.data["index"] == 1
|
||||
)
|
||||
assert block0.data["content_block"]["name"] == "alpha"
|
||||
assert block0.data["content_block"]["id"] == "call_0"
|
||||
assert block1.data["content_block"]["name"] == "beta"
|
||||
assert block1.data["content_block"]["id"] == "call_1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_message_maps_length_finish_reason(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A truncated (length) text stream must report stop_reason=max_tokens."""
|
||||
backend, instance = make_backend(monkeypatch)
|
||||
instance.response = FakeAsyncStream(
|
||||
[
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="hi"), finish_reason=None)]
|
||||
),
|
||||
SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(), finish_reason="length")]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
events = [
|
||||
event async for event in backend.stream_message({"model": "claude", "messages": []}, {})
|
||||
]
|
||||
message_delta = next(e for e in events if e.event_type == "message_delta")
|
||||
assert message_delta.data["delta"]["stop_reason"] == "max_tokens"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_openai_message_maps_choices_and_tool_calls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
Reference in New Issue
Block a user