Bugfix/Prevent streaming of chatflow tool and chain tool (#3257)

prevent streaming of chatflow tool and chain tool
FlowiseAI · Sep 30, 2024 · a80ba4b · a80ba4b
1 parent d2577e5
commit a80ba4b
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 40 deletions.
diff --git a/packages/components/nodes/tools/ChainTool/core.ts b/packages/components/nodes/tools/ChainTool/core.ts
@@ -14,41 +14,17 @@ export class ChainTool extends DynamicTool {
         super({
             ...rest,
             func: async (input, runManager) => {
-                // prevent sending SSE events of the sub-chain
-                const sseStreamer = runManager?.handlers.find((handler) => handler instanceof CustomChainHandler)?.sseStreamer
-                if (runManager) {
-                    const callbacks = runManager.handlers
-                    for (let i = 0; i < callbacks.length; i += 1) {
-                        if (callbacks[i] instanceof CustomChainHandler) {
-                            ;(callbacks[i] as any).sseStreamer = undefined
-                        }
-                    }
-                }
+                const childManagers = runManager?.getChild()
+                const handlers = childManagers?.handlers?.filter((handler) => !(handler instanceof CustomChainHandler)) || []
+                if (childManagers) childManagers.handlers = handlers
 
                 if ((chain as any).prompt && (chain as any).prompt.promptValues) {
                     const promptValues = handleEscapeCharacters((chain as any).prompt.promptValues, true)
-
-                    const values = await chain.call(promptValues, runManager?.getChild())
-                    if (runManager && sseStreamer) {
-                        const callbacks = runManager.handlers
-                        for (let i = 0; i < callbacks.length; i += 1) {
-                            if (callbacks[i] instanceof CustomChainHandler) {
-                                ;(callbacks[i] as any).sseStreamer = sseStreamer
-                            }
-                        }
-                    }
+                    const values = await chain.call(promptValues, childManagers)
                     return values?.text
                 }
 
-                const values = chain.run(input, runManager?.getChild())
-                if (runManager && sseStreamer) {
-                    const callbacks = runManager.handlers
-                    for (let i = 0; i < callbacks.length; i += 1) {
-                        if (callbacks[i] instanceof CustomChainHandler) {
-                            ;(callbacks[i] as any).sseStreamer = sseStreamer
-                        }
-                    }
-                }
+                const values = chain.run(input, childManagers)
                 return values
             }
         })

diff --git a/packages/components/nodes/tools/ChatflowTool/ChatflowTool.ts b/packages/components/nodes/tools/ChatflowTool/ChatflowTool.ts
@@ -7,7 +7,6 @@ import { StructuredTool } from '@langchain/core/tools'
 import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
 import { availableDependencies, defaultAllowBuiltInDep, getCredentialData, getCredentialParam } from '../../../src/utils'
 import { v4 as uuidv4 } from 'uuid'
-import { CustomChainHandler } from '../../../src'
 
 class ChatflowTool_Tools implements INode {
     label: string
@@ -24,7 +23,7 @@ class ChatflowTool_Tools implements INode {
     constructor() {
         this.label = 'Chatflow Tool'
         this.name = 'ChatflowTool'
-        this.version = 4.0
+        this.version = 5.0
         this.type = 'ChatflowTool'
         this.icon = 'chatflowTool.svg'
         this.category = 'Tools'
@@ -58,6 +57,12 @@ class ChatflowTool_Tools implements INode {
                 placeholder:
                     'State of the Union QA - useful for when you need to ask questions about the most recent state of the union address.'
             },
+            {
+                label: 'Return Direct',
+                name: 'returnDirect',
+                type: 'boolean',
+                optional: true
+            },
             {
                 label: 'Override Config',
                 name: 'overrideConfig',
@@ -135,6 +140,7 @@ class ChatflowTool_Tools implements INode {
         const _name = nodeData.inputs?.name as string
         const description = nodeData.inputs?.description as string
         const useQuestionFromChat = nodeData.inputs?.useQuestionFromChat as boolean
+        const returnDirect = nodeData.inputs?.returnDirect as boolean
         const customInput = nodeData.inputs?.customInput as string
         const overrideConfig =
             typeof nodeData.inputs?.overrideConfig === 'string' &&
@@ -168,6 +174,7 @@ class ChatflowTool_Tools implements INode {
             name,
             baseURL,
             description,
+            returnDirect,
             chatflowid: selectedChatflowId,
             startNewSession,
             headers,
@@ -206,6 +213,7 @@ class ChatflowTool extends StructuredTool {
     constructor({
         name,
         description,
+        returnDirect,
         input,
         chatflowid,
         startNewSession,
@@ -215,6 +223,7 @@ class ChatflowTool extends StructuredTool {
     }: {
         name: string
         description: string
+        returnDirect: boolean
         input: string
         chatflowid: string
         startNewSession: boolean
@@ -231,6 +240,7 @@ class ChatflowTool extends StructuredTool {
         this.headers = headers
         this.chatflowid = chatflowid
         this.overrideConfig = overrideConfig
+        this.returnDirect = returnDirect
     }
 
     async call(
@@ -249,15 +259,6 @@ class ChatflowTool extends StructuredTool {
         } catch (e) {
             throw new Error(`Received tool input did not match expected schema: ${JSON.stringify(arg)}`)
         }
-        // iterate over the callbacks and the sse streamer
-        if (config.callbacks instanceof CallbackManager) {
-            const callbacks = config.callbacks.handlers
-            for (let i = 0; i < callbacks.length; i += 1) {
-                if (callbacks[i] instanceof CustomChainHandler) {
-                    ;(callbacks[i] as any).sseStreamer = undefined
-                }
-            }
-        }
         const callbackManager_ = await CallbackManager.configure(
             config.callbacks,
             this.callbacks,
@@ -283,6 +284,9 @@ class ChatflowTool extends StructuredTool {
             await runManager?.handleToolError(e)
             throw e
         }
+        if (result && typeof result !== 'string') {
+            result = JSON.stringify(result)
+        }
         await runManager?.handleToolEnd(result)
         return result
     }

diff --git a/packages/server/src/utils/buildChatflow.ts b/packages/server/src/utils/buildChatflow.ts
@@ -360,6 +360,8 @@ export const utilBuildChatflow = async (req: Request, isInternal: boolean = fals
         const nodeModule = await import(nodeInstanceFilePath)
         const nodeInstance = new nodeModule.nodeClass({ sessionId })
 
+        isStreamValid = (req.body.streaming === 'true' || req.body.streaming === true) && isStreamValid
+
         let result = isStreamValid
             ? await nodeInstance.run(nodeToExecuteData, incomingInput.question, {
                   chatId,