This commit is contained in:
朱潮 2025-12-14 09:43:24 +08:00
parent 367e292854
commit 671569026e

View File

@ -144,7 +144,7 @@ async def enhanced_generate_stream_response(
chunk_id = 0 chunk_id = 0
message_tag = "" message_tag = ""
function_call = {} function_call = {}
function_name = "" tool_id = ""
async for msg, metadata in agent.astream({"messages": messages}, stream_mode="messages"): async for msg, metadata in agent.astream({"messages": messages}, stream_mode="messages"):
new_content = "" new_content = ""
@ -153,13 +153,15 @@ async def enhanced_generate_stream_response(
if msg.tool_call_chunks: # 检查工具调用块 if msg.tool_call_chunks: # 检查工具调用块
if message_tag != "TOOL_CALL": if message_tag != "TOOL_CALL":
message_tag = "TOOL_CALL" message_tag = "TOOL_CALL"
function_name = ""
if msg.tool_call_chunks[0]["name"]: if msg.tool_call_chunks[0]["name"]:
function_name = msg.tool_call_chunks[0]["name"] tool_id = msg.tool_call_chunks[0]["id"]
if function_name and msg.tool_call_chunks[0]["args"]: if tool_id not in function_call:
if function_name not in function_call: function_call[tool_id] = {
function_call[function_name] = "" "function_name": msg.tool_call_chunks[0]["name"],
function_call[function_name] += msg.tool_call_chunks[0]["args"] "arguments": ""
}
if tool_id and tool_id in function_call and msg.tool_call_chunks[0]["args"]:
function_call[tool_id]["arguments"] += msg.tool_call_chunks[0]["args"]
elif len(msg.content) > 0: elif len(msg.content) > 0:
preamble_completed.set() preamble_completed.set()
await output_queue.put(("preamble_done", None)) await output_queue.put(("preamble_done", None))
@ -173,8 +175,8 @@ async def enhanced_generate_stream_response(
("finish_reason" in msg.response_metadata and msg.response_metadata["finish_reason"] == "tool_calls") or \ ("finish_reason" in msg.response_metadata and msg.response_metadata["finish_reason"] == "tool_calls") or \
("stop_reason" in msg.response_metadata and msg.response_metadata["stop_reason"] == "tool_use") ("stop_reason" in msg.response_metadata and msg.response_metadata["stop_reason"] == "tool_use")
): ):
for function_name, args in function_call.items(): for args in function_call.values():
new_content = f"[{message_tag}] {function_name}\n{args}" new_content = f"[{message_tag}] {args['function_name']}\n{args['arguments']}"
message_tag = "TOOL_CALL_FINISH" message_tag = "TOOL_CALL_FINISH"
elif isinstance(msg, ToolMessage) and len(msg.content) > 0: elif isinstance(msg, ToolMessage) and len(msg.content) > 0:
message_tag = "TOOL_RESPONSE" message_tag = "TOOL_RESPONSE"