Skip to content

Commit 6f94d61

Browse files
HenryHengZJHenryHengZJ
authored
Updated ChatMessage component to display calledTools (#5365)
1 parent 62d3406 commit 6f94d61

File tree

3 files changed

+116
-5
lines changed

3 files changed

+116
-5
lines changed

‎packages/components/nodes/agentflow/Agent/Agent.ts‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1676,7 +1676,12 @@ class Agent_Agentflow implements INode {
16761676
const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
16771677

16781678
if (response.tool_calls) {
1679-
sseStreamer.streamCalledToolsEvent(chatId, response.tool_calls)
1679+
const formattedToolCalls = response.tool_calls.map((toolCall: any) => ({
1680+
tool: toolCall.name || 'tool',
1681+
toolInput: toolCall.args,
1682+
toolOutput: ''
1683+
}))
1684+
sseStreamer.streamCalledToolsEvent(chatId, flatten(formattedToolCalls))
16801685
}
16811686

16821687
if (response.usage_metadata) {
@@ -1736,7 +1741,12 @@ class Agent_Agentflow implements INode {
17361741

17371742
// Stream tool calls if available
17381743
if (sseStreamer) {
1739-
sseStreamer.streamCalledToolsEvent(chatId, JSON.stringify(response.tool_calls))
1744+
const formattedToolCalls = response.tool_calls.map((toolCall: any) => ({
1745+
tool: toolCall.name || 'tool',
1746+
toolInput: toolCall.args,
1747+
toolOutput: ''
1748+
}))
1749+
sseStreamer.streamCalledToolsEvent(chatId, flatten(formattedToolCalls))
17401750
}
17411751

17421752
// Remove tool calls with no id
@@ -2045,7 +2055,12 @@ class Agent_Agentflow implements INode {
20452055

20462056
// Stream tool calls if available
20472057
if (sseStreamer) {
2048-
sseStreamer.streamCalledToolsEvent(chatId, JSON.stringify(response.tool_calls))
2058+
const formattedToolCalls = response.tool_calls.map((toolCall: any) => ({
2059+
tool: toolCall.name || 'tool',
2060+
toolInput: toolCall.args,
2061+
toolOutput: ''
2062+
}))
2063+
sseStreamer.streamCalledToolsEvent(chatId, flatten(formattedToolCalls))
20492064
}
20502065

20512066
// Remove tool calls with no id

‎packages/components/nodes/agentflow/LLM/LLM.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
updateFlowState
1414
} from '../utils'
1515
import { processTemplateVariables } from '../../../src/utils'
16+
import { flatten } from 'lodash'
1617

1718
class LLM_Agentflow implements INode {
1819
label: string
@@ -892,7 +893,12 @@ class LLM_Agentflow implements INode {
892893
const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
893894

894895
if (response.tool_calls) {
895-
sseStreamer.streamCalledToolsEvent(chatId, response.tool_calls)
896+
const formattedToolCalls = response.tool_calls.map((toolCall: any) => ({
897+
tool: toolCall.name || 'tool',
898+
toolInput: toolCall.args,
899+
toolOutput: ''
900+
}))
901+
sseStreamer.streamCalledToolsEvent(chatId, flatten(formattedToolCalls))
896902
}
897903

898904
if (response.usage_metadata) {

‎packages/ui/src/views/chatmessage/ChatMessage.jsx‎

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -687,11 +687,57 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
687687
setMessages((prevMessages) => {
688688
let allMessages = [...cloneDeep(prevMessages)]
689689
if (allMessages[allMessages.length - 1].type === 'userMessage') return allMessages
690+
691+
// When usedTools are received, check if there are matching calledTools to replace
692+
const lastMessage = allMessages[allMessages.length - 1]
693+
if (lastMessage.calledTools && lastMessage.calledTools.length > 0) {
694+
// Replace calledTools with usedTools for matching tool names
695+
const updatedCalledTools = lastMessage.calledTools.map((calledTool) => {
696+
const matchingUsedTool = usedTools.find((usedTool) => usedTool.tool === calledTool.tool)
697+
return matchingUsedTool || calledTool
698+
})
699+
700+
// Remove calledTools that have been replaced by usedTools
701+
const remainingCalledTools = updatedCalledTools.filter(
702+
(calledTool) => !usedTools.some((usedTool) => usedTool.tool === calledTool.tool)
703+
)
704+
705+
allMessages[allMessages.length - 1].calledTools = remainingCalledTools.length > 0 ? remainingCalledTools : undefined
706+
}
707+
690708
allMessages[allMessages.length - 1].usedTools = usedTools
691709
return allMessages
692710
})
693711
}
694712

713+
const updateLastMessageCalledTools = (calledTools) => {
714+
setMessages((prevMessages) => {
715+
let allMessages = [...cloneDeep(prevMessages)]
716+
if (allMessages[allMessages.length - 1].type === 'userMessage') return allMessages
717+
allMessages[allMessages.length - 1].calledTools = calledTools
718+
return allMessages
719+
})
720+
}
721+
722+
const cleanupCalledTools = () => {
723+
setMessages((prevMessages) => {
724+
let allMessages = [...cloneDeep(prevMessages)]
725+
if (allMessages[allMessages.length - 1].type === 'userMessage') return allMessages
726+
727+
// Remove any remaining calledTools when the stream ends
728+
const lastMessage = allMessages[allMessages.length - 1]
729+
if (lastMessage && lastMessage.calledTools && lastMessage.calledTools.length > 0) {
730+
// Only remove if there are still calledTools and no matching usedTools
731+
const hasUsedTools = lastMessage.usedTools && lastMessage.usedTools.length > 0
732+
if (!hasUsedTools) {
733+
allMessages[allMessages.length - 1].calledTools = undefined
734+
}
735+
}
736+
737+
return allMessages
738+
})
739+
}
740+
695741
const updateLastMessageFileAnnotations = (fileAnnotations) => {
696742
setMessages((prevMessages) => {
697743
let allMessages = [...cloneDeep(prevMessages)]
@@ -710,6 +756,7 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
710756
if (lastAgentReasoning && lastAgentReasoning.length > 0) {
711757
allMessages[allMessages.length - 1].agentReasoning = lastAgentReasoning.filter((reasoning) => !reasoning.nextAgent)
712758
}
759+
allMessages[allMessages.length - 1].calledTools = undefined
713760
return allMessages
714761
})
715762
setTimeout(() => {
@@ -1038,6 +1085,9 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
10381085
case 'usedTools':
10391086
updateLastMessageUsedTools(payload.data)
10401087
break
1088+
case 'calledTools':
1089+
updateLastMessageCalledTools(payload.data)
1090+
break
10411091
case 'fileAnnotations':
10421092
updateLastMessageFileAnnotations(payload.data)
10431093
break
@@ -1085,12 +1135,14 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
10851135
handleTTSAbort(payload.data)
10861136
break
10871137
case 'end':
1138+
cleanupCalledTools()
10881139
setLocalStorageChatflow(chatflowid, chatId)
10891140
closeResponse()
10901141
break
10911142
}
10921143
},
10931144
async onclose() {
1145+
cleanupCalledTools()
10941146
closeResponse()
10951147
},
10961148
async onerror(err) {
@@ -1102,6 +1154,7 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
11021154
}
11031155

11041156
const closeResponse = () => {
1157+
cleanupCalledTools()
11051158
setLoading(false)
11061159
setUserInput('')
11071160
setUploadedFiles([])
@@ -1202,6 +1255,7 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
12021255
}
12031256
if (message.sourceDocuments) obj.sourceDocuments = message.sourceDocuments
12041257
if (message.usedTools) obj.usedTools = message.usedTools
1258+
if (message.calledTools) obj.calledTools = message.calledTools
12051259
if (message.fileAnnotations) obj.fileAnnotations = message.fileAnnotations
12061260
if (message.agentReasoning) obj.agentReasoning = message.agentReasoning
12071261
if (message.action) obj.action = message.action
@@ -2423,6 +2477,42 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
24232477
sessionId={chatId}
24242478
/>
24252479
)}
2480+
{message.calledTools && (
2481+
<div
2482+
style={{
2483+
display: 'block',
2484+
flexDirection: 'row',
2485+
width: '100%'
2486+
}}
2487+
>
2488+
{message.calledTools.map((tool, index) => {
2489+
return tool ? (
2490+
<Chip
2491+
size='small'
2492+
key={`called-${index}`}
2493+
label={tool.tool}
2494+
component='a'
2495+
sx={{
2496+
mr: 1,
2497+
mt: 1,
2498+
borderColor: 'primary.main',
2499+
color: 'primary.main',
2500+
backgroundColor: 'rgba(25, 118, 210, 0.1)',
2501+
opacity: 0.9,
2502+
'&:hover': {
2503+
backgroundColor: 'rgba(25, 118, 210, 0.2)',
2504+
opacity: 1
2505+
}
2506+
}}
2507+
variant='outlined'
2508+
clickable
2509+
icon={<CircularProgress size={15} color='primary' />}
2510+
onClick={() => onSourceDialogClick(tool, 'Called Tools')}
2511+
/>
2512+
) : null
2513+
})}
2514+
</div>
2515+
)}
24262516
{message.usedTools && (
24272517
<div
24282518
style={{
@@ -2435,7 +2525,7 @@ const ChatMessage = ({ open, chatflowid, isAgentCanvas, isDialog, previews, setP
24352525
return tool ? (
24362526
<Chip
24372527
size='small'
2438-
key={index}
2528+
key={`used-${index}`}
24392529
label={tool.tool}
24402530
component='a'
24412531
sx={{

0 commit comments

Comments
 (0)