Bugfix/stream custom tool return direct (FlowiseAI#3003)

stream custom tool return direct
exdatic · Aug 12, 2024 · b9f0ec3 · b9f0ec3
1 parent f57dc24
commit b9f0ec3
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 1 deletion.
diff --git a/packages/components/nodes/agents/ConversationalAgent/ConversationalAgent.ts b/packages/components/nodes/agents/ConversationalAgent/ConversationalAgent.ts
@@ -136,6 +136,17 @@ class ConversationalAgent_Agents implements INode {
                 options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
                 usedTools = res.usedTools
             }
+            // If the tool is set to returnDirect, stream the output to the client
+            if (res.usedTools && res.usedTools.length) {
+                let inputTools = nodeData.inputs?.tools
+                inputTools = flatten(inputTools)
+                for (const tool of res.usedTools) {
+                    const inputTool = inputTools.find((inputTool: Tool) => inputTool.name === tool.tool)
+                    if (inputTool && inputTool.returnDirect) {
+                        options.socketIO.to(options.socketIOClientId).emit('token', tool.toolOutput)
+                    }
+                }
+            }
         } else {
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
             if (res.sourceDocuments) {

diff --git a/packages/components/nodes/agents/ToolAgent/ToolAgent.ts b/packages/components/nodes/agents/ToolAgent/ToolAgent.ts
@@ -1,4 +1,5 @@
 import { flatten } from 'lodash'
+import { Tool } from '@langchain/core/tools'
 import { BaseMessage } from '@langchain/core/messages'
 import { ChainValues } from '@langchain/core/utils/types'
 import { RunnableSequence } from '@langchain/core/runnables'
@@ -125,6 +126,17 @@ class ToolAgent_Agents implements INode {
                 options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
                 usedTools = res.usedTools
             }
+            // If the tool is set to returnDirect, stream the output to the client
+            if (res.usedTools && res.usedTools.length) {
+                let inputTools = nodeData.inputs?.tools
+                inputTools = flatten(inputTools)
+                for (const tool of res.usedTools) {
+                    const inputTool = inputTools.find((inputTool: Tool) => inputTool.name === tool.tool)
+                    if (inputTool && inputTool.returnDirect) {
+                        options.socketIO.to(options.socketIOClientId).emit('token', tool.toolOutput)
+                    }
+                }
+            }
         } else {
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
             if (res.sourceDocuments) {

diff --git a/packages/components/nodes/agents/XMLAgent/XMLAgent.ts b/packages/components/nodes/agents/XMLAgent/XMLAgent.ts
@@ -142,6 +142,17 @@ class XMLAgent_Agents implements INode {
                 options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
                 usedTools = res.usedTools
             }
+            // If the tool is set to returnDirect, stream the output to the client
+            if (res.usedTools && res.usedTools.length) {
+                let inputTools = nodeData.inputs?.tools
+                inputTools = flatten(inputTools)
+                for (const tool of res.usedTools) {
+                    const inputTool = inputTools.find((inputTool: Tool) => inputTool.name === tool.tool)
+                    if (inputTool && inputTool.returnDirect) {
+                        options.socketIO.to(options.socketIOClientId).emit('token', tool.toolOutput)
+                    }
+                }
+            }
         } else {
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
             if (res.sourceDocuments) {

diff --git a/packages/components/nodes/tools/CustomTool/CustomTool.ts b/packages/components/nodes/tools/CustomTool/CustomTool.ts
@@ -18,7 +18,7 @@ class CustomTool_Tools implements INode {
     constructor() {
         this.label = 'Custom Tool'
         this.name = 'customTool'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'CustomTool'
         this.icon = 'customtool.svg'
         this.category = 'Tools'
@@ -29,6 +29,13 @@ class CustomTool_Tools implements INode {
                 name: 'selectedTool',
                 type: 'asyncOptions',
                 loadMethod: 'listTools'
+            },
+            {
+                label: 'Return Direct',
+                name: 'returnDirect',
+                description: 'Return the output of the tool directly to the user',
+                type: 'boolean',
+                optional: true
             }
         ]
         this.baseClasses = [this.type, 'Tool', ...getBaseClasses(DynamicStructuredTool)]
@@ -66,6 +73,7 @@ class CustomTool_Tools implements INode {
         const customToolName = nodeData.inputs?.customToolName as string
         const customToolDesc = nodeData.inputs?.customToolDesc as string
         const customToolSchema = nodeData.inputs?.customToolSchema as string
+        const customToolReturnDirect = nodeData.inputs?.returnDirect as boolean
 
         const appDataSource = options.appDataSource as DataSource
         const databaseEntities = options.databaseEntities as IDatabaseEntity
@@ -97,6 +105,7 @@ class CustomTool_Tools implements INode {
             let dynamicStructuredTool = new DynamicStructuredTool(obj)
             dynamicStructuredTool.setVariables(variables)
             dynamicStructuredTool.setFlowObject(flow)
+            dynamicStructuredTool.returnDirect = customToolReturnDirect
 
             return dynamicStructuredTool
         } catch (e) {