Skip to content

fix: AI dialogue node calls MCP, AI's reply content is included in the code block of MCP call result, markdown parsing error #3846 #3894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shaohuzhang1 merged 1 commit into from
Aug 20, 2025
Merged

fix: AI dialogue node calls MCP, AI's reply content is included in the code block of MCP call result, markdown parsing error #3846 #3894

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import re
import time
from functools import reduce
from types import AsyncGeneratorType
from typing import List, Dict

from django.db.models import QuerySet
Expand All @@ -33,13 +32,25 @@
<strong>Called MCP Tool: <em>%s</em></strong>
</summary>

```json
%s
```

</details>

"""

tool_message_json_template = """
```json
%s
```
"""


def generate_tool_message_template(name, context):
if '```' in context:
return tool_message_template % (name, context)
else:
return tool_message_template % (name, tool_message_json_template % (context))


def _write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow, answer: str,
reasoning_content: str):
Expand Down Expand Up @@ -109,7 +120,7 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers):
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)
content = generate_tool_message_template(chunk[0].name, chunk[0].content)
chunk[0].content = content
yield chunk[0]
if isinstance(chunk[0], AIMessageChunk):
Copy link
Copy Markdown
Contributor Author

@shaohuzhang1 shaohuzhang1 Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code has a few issues and areas for improvement:

  1. Redundancy: The generate_tool_message_template function is called twice within the _yield_mcp_response method. This redundancy can be removed.

  2. Security Concerns: If user input (reasoning_content) contains malicious JSON, it could lead to injection attacks. Consider sanitizing this input before converting it to JSON.

  3. Code Duplication: There are similar template string constructions in both places where you're creating tool messages. Can these be combined into a single function?

Here's an optimized version of the code:

@@ -11,7 +11,6 @@
 import re
 import time
 from functools import reduce
-from types import AsyncGeneratorType
 from typing import List, Dict
 
 from django.db.models import QuerySet
@@ -33,13 +32,25 @@
         <strong>Called MCP Tool: <em>%s</em></strong>
     </summary>
 
-```json
 %s
-```
+
 </details>
 
 """
 
+def format_tool_message(content):
+    """Formats the content as a Markdown block with triple backticks."""
+    return f"```json\n{json.dumps(content, indent=4)}\n```"
 
 def _write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow, answer: str,
                    reasoning_content: str):
@@ -109,7 +120,7 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers):
         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)
+                content = format_tool_message(chunk[0].content)
                 chunk[0].content = content
                 yield chunk[0]
             if isinstance(chunk[0], AIMessageChunk):

Key Changes:

  • Single Function for Formatting: A new format_tool_message function has been introduced to handle the formatting of Tool Message data.
  • Simplified Code: The code duplication between the two places has been eliminated, improving readability and maintainability.

This should help resolve the redundancies and improve security while keeping the code concise.

Expand Down
1 change: 0 additions & 1 deletion ui/src/workflow/common/NodeContainer.vue
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,7 +297,6 @@ function clickNodes(item: any) {
x: anchorData.value?.x + width / 2 + 200,
y: anchorData.value?.y - item.height
})
console.log(nodeModel)
props.nodeModel.graphModel.addEdge({
type: 'app-edge',
sourceNodeId: props.nodeModel.id,
Copy link
Copy Markdown
Contributor Author

@shaohuzhang1 shaohuzhang1 Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is no irregularity or issue in the provided code snippet. The function clickNodes appears to be correctly setting up data and parameters needed for adding an edge to a graph using properties passed from props. A possible optimization suggestion would be to ensure that anchorData.value?.x, width, and other variables used in the calculation have defined values in order to prevent errors at runtime.

Expand Down