mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 16:08:41 +00:00
fix(client): emit streamed tool calls once with complete args (#5408)
* fix(client): emit streamed tool calls once with complete args
When a model streams a tool call as chunks (name and id first, then
argument fragments), DeerFlowClient.stream() emitted a tool_calls event
per chunk, each parsed from that fragment alone. Consumers got
`args={}` and then a call with no name or id, and never the complete
call, because the values snapshot skipped the already-streamed id.
Hold tool calls from AIMessageChunk.tool_call_chunks and emit them once
from the values snapshot, where the message is complete.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* test(client): pin additional_kwargs on the deferred tool_calls event
Also note why a tool-call chunk without a message id keeps the per-chunk
event.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
parent
1dd48d14d2
commit
529b0e05ec
@ -30,7 +30,7 @@ from typing import Any, Literal
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage, ToolMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from deerflow.agents.lead_agent.agent import _authorize_model_name, build_middlewares
|
||||
@ -949,6 +949,10 @@ class DeerFlowClient:
|
||||
# Cross-mode handoff: ids already streamed via LangGraph ``messages``
|
||||
# mode so the ``values`` path skips re-synthesis of the same message.
|
||||
streamed_ids: set[str] = set()
|
||||
# AI messages whose tool calls arrived as streamed fragments. The
|
||||
# arguments only parse once the message is complete, so their
|
||||
# tool_calls event is emitted from the values snapshot instead.
|
||||
pending_tool_call_ids: set[str] = set()
|
||||
# The same message id carries identical cumulative ``usage_metadata``
|
||||
# in both the final ``messages`` chunk and the values snapshot —
|
||||
# count it only on whichever arrives first.
|
||||
@ -1041,7 +1045,12 @@ class DeerFlowClient:
|
||||
)
|
||||
sent_additional_kwargs = bool(additional_kwargs_delta)
|
||||
|
||||
if msg_chunk.tool_calls:
|
||||
# A chunk without an id can't be matched to its values
|
||||
# snapshot, so it keeps the per-chunk event below.
|
||||
if isinstance(msg_chunk, AIMessageChunk) and msg_chunk.tool_call_chunks and msg_id:
|
||||
streamed_ids.add(msg_id)
|
||||
pending_tool_call_ids.add(msg_id)
|
||||
elif msg_chunk.tool_calls:
|
||||
if msg_id:
|
||||
streamed_ids.add(msg_id)
|
||||
additional_kwargs_delta = None if sent_additional_kwargs else _unsent_additional_kwargs(msg_id, additional_kwargs)
|
||||
@ -1074,7 +1083,10 @@ class DeerFlowClient:
|
||||
_account_usage(msg_id, getattr(msg, "usage_metadata", None))
|
||||
additional_kwargs = self._serialize_additional_kwargs(msg)
|
||||
additional_kwargs_delta = _unsent_additional_kwargs(msg_id, additional_kwargs)
|
||||
if additional_kwargs_delta:
|
||||
if msg_id in pending_tool_call_ids and msg.tool_calls:
|
||||
pending_tool_call_ids.discard(msg_id)
|
||||
yield self._ai_tool_calls_event(msg_id, msg.tool_calls, additional_kwargs_delta)
|
||||
elif additional_kwargs_delta:
|
||||
# Metadata-only follow-up: ``messages-tuple`` has no
|
||||
# dedicated attribution event, so clients should
|
||||
# merge this empty-content AI event by message id
|
||||
|
||||
@ -621,6 +621,56 @@ class TestStream:
|
||||
call_kwargs = agent.stream.call_args.kwargs
|
||||
assert "messages" in call_kwargs["stream_mode"]
|
||||
|
||||
def test_stream_emits_streamed_tool_calls_once_with_complete_args(self, client):
|
||||
"""Tool-call arguments streamed in fragments are emitted once, complete.
|
||||
|
||||
Each chunk only parses to a partial call (``args={}``, or no name/id), so
|
||||
the tool_calls event comes from the values snapshot, not from the chunks.
|
||||
"""
|
||||
call = {"name": "bash", "args": {"command": "ls -la"}, "id": "call-1"}
|
||||
attribution = {"version": 1, "kind": "tool_batch", "shared_attribution": False, "actions": []}
|
||||
assembled = AIMessage(content="", id="ai-1", tool_calls=[call], additional_kwargs={"token_usage_attribution": attribution})
|
||||
agent = MagicMock()
|
||||
agent.stream.return_value = iter(
|
||||
[
|
||||
(
|
||||
"messages",
|
||||
(
|
||||
AIMessageChunk(
|
||||
content="",
|
||||
id="ai-1",
|
||||
tool_call_chunks=[{"name": "bash", "args": "", "id": "call-1", "index": 0}],
|
||||
),
|
||||
{},
|
||||
),
|
||||
),
|
||||
(
|
||||
"messages",
|
||||
(
|
||||
AIMessageChunk(
|
||||
content="",
|
||||
id="ai-1",
|
||||
tool_call_chunks=[{"name": None, "args": '{"command": "ls -la"}', "id": None, "index": 0}],
|
||||
),
|
||||
{},
|
||||
),
|
||||
),
|
||||
("values", {"messages": [HumanMessage(content="hi", id="h-1"), assembled]}),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(client, "_ensure_agent"),
|
||||
patch.object(client, "_agent", agent),
|
||||
):
|
||||
events = list(client.stream("hi", thread_id="t-stream-tools"))
|
||||
|
||||
tool_call_events = _tool_call_events(events)
|
||||
assert len(tool_call_events) == 1
|
||||
assert tool_call_events[0].data["id"] == "ai-1"
|
||||
assert [(tc["name"], tc["args"], tc["id"]) for tc in tool_call_events[0].data["tool_calls"]] == [("bash", {"command": "ls -la"}, "call-1")]
|
||||
assert tool_call_events[0].data["additional_kwargs"] == {"token_usage_attribution": attribution}
|
||||
|
||||
def test_stream_emits_additional_kwargs_updates_for_streamed_ai_messages(self, client):
|
||||
"""stream() emits a follow-up AI event when attribution metadata arrives via values."""
|
||||
assembled = AIMessage(
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user