From c4293c34d17d0d9bb162750919d76b82d0563d13 Mon Sep 17 00:00:00 2001 From: shaohuzhang1 <80892890+shaohuzhang1@users.noreply.github.com> Date: Wed, 9 Jul 2025 16:14:44 +0800 Subject: [PATCH] fix: ai chat node mcp (#3531) --- .../ai_chat_step_node/impl/base_chat_node.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py index 66556b40f..492e2984f 100644 --- a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py +++ b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py @@ -104,16 +104,17 @@ def write_context_stream(node_variable: Dict, workflow_variable: Dict, node: INo async def _yield_mcp_response(chat_model, message_list, mcp_servers): - async with MultiServerMCPClient(json.loads(mcp_servers)) as client: - agent = create_react_agent(chat_model, client.get_tools()) - response = agent.astream({"messages": message_list}, stream_mode='messages') - async for chunk in response: - if isinstance(chunk[0], ToolMessage): - content = tool_message_template % (chunk[0].name, chunk[0].content) - chunk[0].content = content - yield chunk[0] - if isinstance(chunk[0], AIMessageChunk): - yield chunk[0] + client = MultiServerMCPClient(json.loads(mcp_servers)) + tools = await client.get_tools() + agent = create_react_agent(chat_model, tools) + response = agent.astream({"messages": message_list}, stream_mode='messages') + async for chunk in response: + if isinstance(chunk[0], ToolMessage): + content = tool_message_template % (chunk[0].name, chunk[0].content) + chunk[0].content = content + yield chunk[0] + if isinstance(chunk[0], AIMessageChunk): + yield chunk[0] def mcp_response_generator(chat_model, message_list, mcp_servers):