Feature/sse (#3125)

* Base changes for ServerSide Events (instead of socket.io)

* lint fixes

* adding of interface and separate methods for streaming events

* lint

* first draft, handles both internal and external prediction end points.

* lint fixes

* additional internal end point for streaming and associated changes

* return streamresponse as true to build agent flow

* 1) JSON formatting for internal events
2) other fixes

* 1) convert internal event to metadata to maintain consistency with external response

* fix action and metadata streaming

* fix for error when agent flow is aborted

* prevent subflows from streaming and other code cleanup

* prevent streaming from enclosed tools

* add fix for preventing chaintool streaming

* update lock file

* add open when hidden to sse

* Streaming errors

* Streaming errors

* add fix for showing error message

---------

Co-authored-by: Henry <hzj94@hotmail.com>

packages/components/nodes/agents/AirtableAgent/AirtableAgent.ts

@@ -2,7 +2,7 @@ import axios from 'axios'
 import { BaseLanguageModel } from '@langchain/core/language_models/base'
 import { AgentExecutor } from 'langchain/agents'
 import { LLMChain } from 'langchain/chains'
-import { ICommonObject, INode, INodeData, INodeParams, PromptTemplate } from '../../../src/Interface'
+import { ICommonObject, INode, INodeData, INodeParams, IServerSideEventStreamer, PromptTemplate } from '../../../src/Interface'
 import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
 import { LoadPyodide, finalSystemPrompt, systemPrompt } from './core'
@@ -104,11 +104,17 @@ class Airtable_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }
 
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
+
         const credentialData = await getCredentialData(nodeData.credential ?? '', options)
         const accessToken = getCredentialParam('accessToken', credentialData, nodeData)
 
@@ -123,7 +129,6 @@ class Airtable_Agents implements INode {
         let base64String = Buffer.from(JSON.stringify(airtableData)).toString('base64')
 
         const loggerHandler = new ConsoleCallbackHandler(options.logger)
-        const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
         const callbacks = await additionalCallbacks(nodeData, options)
 
         const pyodide = await LoadPyodide()
@@ -194,7 +199,8 @@ json.dumps(my_dict)`
                 answer: finalResult
             }
 
-            if (options.socketIO && options.socketIOClientId) {
+            if (options.shouldStreamResponse) {
+                const handler = new CustomChainHandler(shouldStreamResponse ? sseStreamer : undefined, chatId)
                 const result = await chain.call(inputs, [loggerHandler, handler, ...callbacks])
                 return result?.text
             } else {

packages/components/nodes/agents/AutoGPT/AutoGPT.ts

@@ -113,7 +113,9 @@ class AutoGPT_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }

packages/components/nodes/agents/BabyAGI/BabyAGI.ts

@@ -73,7 +73,9 @@ class BabyAGI_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }

packages/components/nodes/agents/CSVAgent/CSVAgent.ts

@@ -2,7 +2,7 @@ import { BaseLanguageModel } from '@langchain/core/language_models/base'
 import { AgentExecutor } from 'langchain/agents'
 import { LLMChain } from 'langchain/chains'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
-import { ICommonObject, INode, INodeData, INodeParams, PromptTemplate } from '../../../src/Interface'
+import { ICommonObject, INode, INodeData, INodeParams, IServerSideEventStreamer, PromptTemplate } from '../../../src/Interface'
 import { getBaseClasses } from '../../../src/utils'
 import { LoadPyodide, finalSystemPrompt, systemPrompt } from './core'
 import { checkInputs, Moderation } from '../../moderation/Moderation'
@@ -90,13 +90,18 @@ class CSV_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }
 
         const loggerHandler = new ConsoleCallbackHandler(options.logger)
-        const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
+
         const callbacks = await additionalCallbacks(nodeData, options)
 
         let files: string[] = []
@@ -203,7 +208,8 @@ json.dumps(my_dict)`
                 answer: finalResult
             }
 
-            if (options.socketIO && options.socketIOClientId) {
+            if (options.shouldStreamResponse) {
+                const handler = new CustomChainHandler(shouldStreamResponse ? sseStreamer : undefined, chatId)
                 const result = await chain.call(inputs, [loggerHandler, handler, ...callbacks])
                 return result?.text
             } else {

packages/components/nodes/agents/ConversationalAgent/ConversationalAgent.ts

@@ -9,7 +9,16 @@ import { RunnableSequence } from '@langchain/core/runnables'
 import { ChatConversationalAgent } from 'langchain/agents'
 import { getBaseClasses } from '../../../src/utils'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
-import { IVisionChatModal, FlowiseMemory, ICommonObject, INode, INodeData, INodeParams, IUsedTool } from '../../../src/Interface'
+import {
+    IVisionChatModal,
+    FlowiseMemory,
+    ICommonObject,
+    INode,
+    INodeData,
+    INodeParams,
+    IUsedTool,
+    IServerSideEventStreamer
+} from '../../../src/Interface'
 import { AgentExecutor } from '../../../src/agents'
 import { addImagesToMessages, llmSupportsVision } from '../../../src/multiModalUtils'
 import { checkInputs, Moderation } from '../../moderation/Moderation'
@@ -106,12 +115,18 @@ class ConversationalAgent_Agents implements INode {
         const memory = nodeData.inputs?.memory as FlowiseMemory
         const moderations = nodeData.inputs?.inputModeration as Moderation[]
 
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
         if (moderations && moderations.length > 0) {
             try {
                 // Use the output of the moderation chain as input for the BabyAGI agent
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
                 return formatResponse(e.message)
             }
@@ -125,15 +140,17 @@ class ConversationalAgent_Agents implements INode {
         let sourceDocuments: ICommonObject[] = []
         let usedTools: IUsedTool[] = []
 
-        if (options.socketIO && options.socketIOClientId) {
-            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+        if (options.shouldStreamResponse) {
+            const handler = new CustomChainHandler(shouldStreamResponse ? sseStreamer : undefined, chatId)
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
             if (res.sourceDocuments) {
-                options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))
+                if (options.sseStreamer) {
+                    sseStreamer.streamSourceDocumentsEvent(options.chatId, flatten(res.sourceDocuments))
+                }
                 sourceDocuments = res.sourceDocuments
             }
             if (res.usedTools) {
-                options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
+                sseStreamer.streamUsedToolsEvent(options.chatId, res.usedTools)
                 usedTools = res.usedTools
             }
             // If the tool is set to returnDirect, stream the output to the client
@@ -142,11 +159,14 @@ class ConversationalAgent_Agents implements INode {
                 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)
+                    if (inputTool && inputTool.returnDirect && options.sseStreamer) {
+                        sseStreamer.streamTokenEvent(options.chatId, tool.toolOutput)
                     }
                 }
             }
+            if (sseStreamer) {
+                sseStreamer.streamEndEvent(options.chatId)
+            }
         } else {
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
             if (res.sourceDocuments) {

packages/components/nodes/agents/ConversationalRetrievalToolAgent/ConversationalRetrievalToolAgent.ts

@@ -7,7 +7,16 @@ import { ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate, Pr
 import { formatToOpenAIToolMessages } from 'langchain/agents/format_scratchpad/openai_tools'
 import { getBaseClasses } from '../../../src/utils'
 import { type ToolsAgentStep } from 'langchain/agents/openai/output_parser'
-import { FlowiseMemory, ICommonObject, INode, INodeData, INodeParams, IUsedTool, IVisionChatModal } from '../../../src/Interface'
+import {
+    FlowiseMemory,
+    ICommonObject,
+    INode,
+    INodeData,
+    INodeParams,
+    IServerSideEventStreamer,
+    IUsedTool,
+    IVisionChatModal
+} from '../../../src/Interface'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
 import { AgentExecutor, ToolCallingAgentOutputParser } from '../../../src/agents'
 import { Moderation, checkInputs, streamResponse } from '../../moderation/Moderation'
@@ -104,7 +113,9 @@ class ConversationalRetrievalToolAgent_Agents implements INode {
         const memory = nodeData.inputs?.memory as FlowiseMemory
         const moderations = nodeData.inputs?.inputModeration as Moderation[]
 
-        const isStreamable = options.socketIO && options.socketIOClientId
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
 
         if (moderations && moderations.length > 0) {
             try {
@@ -112,8 +123,9 @@ class ConversationalRetrievalToolAgent_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                if (isStreamable)
-                    streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                if (shouldStreamResponse) {
+                    streamResponse(sseStreamer, chatId, e.message)
+                }
                 return formatResponse(e.message)
             }
         }
@@ -127,15 +139,15 @@ class ConversationalRetrievalToolAgent_Agents implements INode {
         let sourceDocuments: ICommonObject[] = []
         let usedTools: IUsedTool[] = []
 
-        if (isStreamable) {
-            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+        if (shouldStreamResponse) {
+            const handler = new CustomChainHandler(sseStreamer, chatId)
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
             if (res.sourceDocuments) {
-                options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))
+                sseStreamer.streamSourceDocumentsEvent(chatId, flatten(res.sourceDocuments))
                 sourceDocuments = res.sourceDocuments
             }
             if (res.usedTools) {
-                options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
+                sseStreamer.streamUsedToolsEvent(chatId, res.usedTools)
                 usedTools = res.usedTools
             }
         } else {

packages/components/nodes/agents/LlamaIndexAgents/OpenAIToolAgent/OpenAIToolAgent_LlamaIndex.ts

@@ -1,7 +1,16 @@
 import { flatten } from 'lodash'
 import { ChatMessage, OpenAI, OpenAIAgent } from 'llamaindex'
 import { getBaseClasses } from '../../../../src/utils'
-import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams, IUsedTool } from '../../../../src/Interface'
+import {
+    FlowiseMemory,
+    ICommonObject,
+    IMessage,
+    INode,
+    INodeData,
+    INodeParams,
+    IServerSideEventStreamer,
+    IUsedTool
+} from '../../../../src/Interface'
 
 class OpenAIFunctionAgent_LlamaIndex_Agents implements INode {
     label: string
@@ -67,7 +76,9 @@ class OpenAIFunctionAgent_LlamaIndex_Agents implements INode {
         let tools = nodeData.inputs?.tools
         tools = flatten(tools)
 
-        const isStreamingEnabled = options.socketIO && options.socketIOClientId
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
 
         const chatHistory = [] as ChatMessage[]
 
@@ -104,7 +115,7 @@ class OpenAIFunctionAgent_LlamaIndex_Agents implements INode {
         let isStreamingStarted = false
         const usedTools: IUsedTool[] = []
 
-        if (isStreamingEnabled) {
+        if (shouldStreamResponse) {
             const stream = await agent.chat({
                 message: input,
                 chatHistory,
@@ -116,7 +127,9 @@ class OpenAIFunctionAgent_LlamaIndex_Agents implements INode {
                 text += chunk.response.delta
                 if (!isStreamingStarted) {
                     isStreamingStarted = true
-                    options.socketIO.to(options.socketIOClientId).emit('start', chunk.response.delta)
+                    if (sseStreamer) {
+                        sseStreamer.streamStartEvent(chatId, chunk.response.delta)
+                    }
                     if (chunk.sources.length) {
                         for (const sourceTool of chunk.sources) {
                             usedTools.push({
@@ -125,11 +138,14 @@ class OpenAIFunctionAgent_LlamaIndex_Agents implements INode {
                                 toolOutput: sourceTool.output as any
                             })
                         }
-                        options.socketIO.to(options.socketIOClientId).emit('usedTools', usedTools)
+                        if (sseStreamer) {
+                            sseStreamer.streamUsedToolsEvent(chatId, usedTools)
+                        }
                     }
                 }
-
-                options.socketIO.to(options.socketIOClientId).emit('token', chunk.response.delta)
+                if (sseStreamer) {
+                    sseStreamer.streamTokenEvent(chatId, chunk.response.delta)
+                }
             }
         } else {
             const response = await agent.chat({ message: input, chatHistory, verbose: process.env.DEBUG === 'true' ? true : false })

packages/components/nodes/agents/OpenAIAssistant/OpenAIAssistant.ts

@@ -1,4 +1,13 @@
-import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, INodeParams, IUsedTool } from '../../../src/Interface'
+import {
+    ICommonObject,
+    IDatabaseEntity,
+    INode,
+    INodeData,
+    INodeOptionsValue,
+    INodeParams,
+    IServerSideEventStreamer,
+    IUsedTool
+} from '../../../src/Interface'
 import OpenAI from 'openai'
 import { DataSource } from 'typeorm'
 import { getCredentialData, getCredentialParam } from '../../../src/utils'
@@ -176,16 +185,19 @@ class OpenAIAssistant_Agents implements INode {
         const moderations = nodeData.inputs?.inputModeration as Moderation[]
         const _toolChoice = nodeData.inputs?.toolChoice as string
         const parallelToolCalls = nodeData.inputs?.parallelToolCalls as boolean
-        const isStreaming = options.socketIO && options.socketIOClientId
-        const socketIO = isStreaming ? options.socketIO : undefined
-        const socketIOClientId = isStreaming ? options.socketIOClientId : ''
+
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
 
         if (moderations && moderations.length > 0) {
             try {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                streamResponse(isStreaming, e.message, socketIO, socketIOClientId)
+                if (shouldStreamResponse) {
+                    streamResponse(sseStreamer, chatId, e.message)
+                }
                 return formatResponse(e.message)
             }
         }
@@ -307,7 +319,7 @@ class OpenAIAssistant_Agents implements INode {
                 }
             }
 
-            if (isStreaming) {
+            if (shouldStreamResponse) {
                 const streamThread = await openai.beta.threads.runs.create(threadId, {
                     assistant_id: retrievedAssistant.id,
                     stream: true,
@@ -389,26 +401,37 @@ class OpenAIAssistant_Agents implements INode {
                                 if (message_content.value) {
                                     if (!isStreamingStarted) {
                                         isStreamingStarted = true
-                                        socketIO.to(socketIOClientId).emit('start', message_content.value)
+                                        if (sseStreamer) {
+                                            sseStreamer.streamStartEvent(chatId, message_content.value)
+                                        }
+                                    }
+                                    if (sseStreamer) {
+                                        sseStreamer.streamTokenEvent(chatId, message_content.value)
                                     }
-                                    socketIO.to(socketIOClientId).emit('token', message_content.value)
                                 }
 
                                 if (fileAnnotations.length) {
                                     if (!isStreamingStarted) {
                                         isStreamingStarted = true
-                                        socketIO.to(socketIOClientId).emit('start', '')
+                                        if (sseStreamer) {
+                                            sseStreamer.streamStartEvent(chatId, ' ')
+                                        }
+                                    }
+                                    if (sseStreamer) {
+                                        sseStreamer.streamFileAnnotationsEvent(chatId, fileAnnotations)
                                     }
-                                    socketIO.to(socketIOClientId).emit('fileAnnotations', fileAnnotations)
                                 }
                             } else {
                                 text += chunk.text?.value
                                 if (!isStreamingStarted) {
                                     isStreamingStarted = true
-                                    socketIO.to(socketIOClientId).emit('start', chunk.text?.value)
+                                    if (sseStreamer) {
+                                        sseStreamer.streamStartEvent(chatId, chunk.text?.value || '')
+                                    }
+                                }
+                                if (sseStreamer) {
+                                    sseStreamer.streamTokenEvent(chatId, chunk.text?.value || '')
                                 }
-
-                                socketIO.to(socketIOClientId).emit('token', chunk.text?.value)
                             }
                         }
 
@@ -425,10 +448,13 @@ class OpenAIAssistant_Agents implements INode {
 
                             if (!isStreamingStarted) {
                                 isStreamingStarted = true
-                                socketIO.to(socketIOClientId).emit('start', imgHTML)
+                                if (sseStreamer) {
+                                    sseStreamer.streamStartEvent(chatId, imgHTML)
+                                }
+                            }
+                            if (sseStreamer) {
+                                sseStreamer.streamTokenEvent(chatId, imgHTML)
                             }
-
-                            socketIO.to(socketIOClientId).emit('token', imgHTML)
                         }
                     }
 
@@ -495,15 +521,19 @@ class OpenAIAssistant_Agents implements INode {
                                             text += chunk.text.value
                                             if (!isStreamingStarted) {
                                                 isStreamingStarted = true
-                                                socketIO.to(socketIOClientId).emit('start', chunk.text.value)
+                                                if (sseStreamer) {
+                                                    sseStreamer.streamStartEvent(chatId, chunk.text.value)
+                                                }
+                                            }
+                                            if (sseStreamer) {
+                                                sseStreamer.streamTokenEvent(chatId, chunk.text.value)
                                             }
-
-                                            socketIO.to(socketIOClientId).emit('token', chunk.text.value)
                                         }
                                     }
                                 }
-
-                                socketIO.to(socketIOClientId).emit('usedTools', usedTools)
+                                if (sseStreamer) {
+                                    sseStreamer.streamUsedToolsEvent(chatId, usedTools)
+                                }
                             } catch (error) {
                                 console.error('Error submitting tool outputs:', error)
                                 await openai.beta.threads.runs.cancel(threadId, runThreadId)
@@ -574,7 +604,9 @@ class OpenAIAssistant_Agents implements INode {
 
                                     // Start tool analytics
                                     const toolIds = await analyticHandlers.onToolStart(tool.name, actions[i].toolInput, parentIds)
-                                    if (socketIO && socketIOClientId) socketIO.to(socketIOClientId).emit('tool', tool.name)
+                                    if (shouldStreamResponse && sseStreamer) {
+                                        sseStreamer.streamToolEvent(chatId, tool.name)
+                                    }
 
                                     try {
                                         const toolOutput = await tool.call(actions[i].toolInput, undefined, undefined, {

packages/components/nodes/agents/ReActAgentChat/ReActAgentChat.ts

@@ -88,7 +88,9 @@ class ReActAgentChat_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }

packages/components/nodes/agents/ReActAgentLLM/ReActAgentLLM.ts

@@ -77,7 +77,9 @@ class ReActAgentLLM_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }

packages/components/nodes/agents/ToolAgent/ToolAgent.ts

@@ -8,7 +8,16 @@ import { ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate, Pr
 import { formatToOpenAIToolMessages } from 'langchain/agents/format_scratchpad/openai_tools'
 import { type ToolsAgentStep } from 'langchain/agents/openai/output_parser'
 import { getBaseClasses, handleEscapeCharacters } from '../../../src/utils'
-import { FlowiseMemory, ICommonObject, INode, INodeData, INodeParams, IUsedTool, IVisionChatModal } from '../../../src/Interface'
+import {
+    FlowiseMemory,
+    ICommonObject,
+    INode,
+    INodeData,
+    INodeParams,
+    IServerSideEventStreamer,
+    IUsedTool,
+    IVisionChatModal
+} from '../../../src/Interface'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
 import { AgentExecutor, ToolCallingAgentOutputParser } from '../../../src/agents'
 import { Moderation, checkInputs, streamResponse } from '../../moderation/Moderation'
@@ -100,7 +109,9 @@ class ToolAgent_Agents implements INode {
         const memory = nodeData.inputs?.memory as FlowiseMemory
         const moderations = nodeData.inputs?.inputModeration as Moderation[]
 
-        const isStreamable = options.socketIO && options.socketIOClientId
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
 
         if (moderations && moderations.length > 0) {
             try {
@@ -108,8 +119,9 @@ class ToolAgent_Agents implements INode {
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                if (isStreamable)
-                    streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                if (shouldStreamResponse) {
+                    streamResponse(sseStreamer, chatId, e.message)
+                }
                 return formatResponse(e.message)
             }
         }
@@ -123,15 +135,19 @@ class ToolAgent_Agents implements INode {
         let sourceDocuments: ICommonObject[] = []
         let usedTools: IUsedTool[] = []
 
-        if (isStreamable) {
-            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+        if (shouldStreamResponse) {
+            const handler = new CustomChainHandler(sseStreamer, chatId)
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
             if (res.sourceDocuments) {
-                options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))
+                if (sseStreamer) {
+                    sseStreamer.streamSourceDocumentsEvent(chatId, flatten(res.sourceDocuments))
+                }
                 sourceDocuments = res.sourceDocuments
             }
             if (res.usedTools) {
-                options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
+                if (sseStreamer) {
+                    sseStreamer.streamUsedToolsEvent(chatId, flatten(res.usedTools))
+                }
                 usedTools = res.usedTools
             }
             // If the tool is set to returnDirect, stream the output to the client
@@ -140,8 +156,8 @@ class ToolAgent_Agents implements INode {
                 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)
+                    if (inputTool && inputTool.returnDirect && shouldStreamResponse) {
+                        sseStreamer.streamTokenEvent(chatId, tool.toolOutput)
                     }
                 }
             }

packages/components/nodes/agents/XMLAgent/XMLAgent.ts

@@ -7,7 +7,16 @@ import { Tool } from '@langchain/core/tools'
 import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder } from '@langchain/core/prompts'
 import { formatLogToMessage } from 'langchain/agents/format_scratchpad/log_to_message'
 import { getBaseClasses } from '../../../src/utils'
-import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams, IUsedTool } from '../../../src/Interface'
+import {
+    FlowiseMemory,
+    ICommonObject,
+    IMessage,
+    INode,
+    INodeData,
+    INodeParams,
+    IServerSideEventStreamer,
+    IUsedTool
+} from '../../../src/Interface'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
 import { AgentExecutor, XMLAgentOutputParser } from '../../../src/agents'
 import { Moderation, checkInputs } from '../../moderation/Moderation'
@@ -112,13 +121,19 @@ class XMLAgent_Agents implements INode {
         const memory = nodeData.inputs?.memory as FlowiseMemory
         const moderations = nodeData.inputs?.inputModeration as Moderation[]
 
+        const shouldStreamResponse = options.shouldStreamResponse
+        const sseStreamer: IServerSideEventStreamer = options.sseStreamer as IServerSideEventStreamer
+        const chatId = options.chatId
+
         if (moderations && moderations.length > 0) {
             try {
                 // Use the output of the moderation chain as input for the OpenAI Function Agent
                 input = await checkInputs(moderations, input)
             } catch (e) {
                 await new Promise((resolve) => setTimeout(resolve, 500))
-                //streamResponse(options.socketIO && options.socketIOClientId, e.message, options.socketIO, options.socketIOClientId)
+                // if (options.shouldStreamResponse) {
+                //     streamResponse(options.sseStreamer, options.chatId, e.message)
+                // }
                 return formatResponse(e.message)
             }
         }
@@ -131,15 +146,19 @@ class XMLAgent_Agents implements INode {
         let sourceDocuments: ICommonObject[] = []
         let usedTools: IUsedTool[] = []
 
-        if (options.socketIO && options.socketIOClientId) {
-            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+        if (shouldStreamResponse) {
+            const handler = new CustomChainHandler(sseStreamer, chatId)
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
             if (res.sourceDocuments) {
-                options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))
+                if (sseStreamer) {
+                    sseStreamer.streamSourceDocumentsEvent(chatId, flatten(res.sourceDocuments))
+                }
                 sourceDocuments = res.sourceDocuments
             }
             if (res.usedTools) {
-                options.socketIO.to(options.socketIOClientId).emit('usedTools', res.usedTools)
+                if (sseStreamer) {
+                    sseStreamer.streamUsedToolsEvent(chatId, flatten(res.usedTools))
+                }
                 usedTools = res.usedTools
             }
             // If the tool is set to returnDirect, stream the output to the client
@@ -149,7 +168,9 @@ class XMLAgent_Agents implements INode {
                 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)
+                        if (sseStreamer) {
+                            sseStreamer.streamTokenEvent(chatId, tool.toolOutput)
+                        }
                     }
                 }
             }