enable streaming

packages/components/nodes/chains/ConversationChain/ConversationChain.ts

@@ -1,6 +1,6 @@
 import { ICommonObject, IMessage, INode, INodeData, INodeParams } from '../../../src/Interface'
 import { ConversationChain } from 'langchain/chains'
-import { getBaseClasses } from '../../../src/utils'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder, SystemMessagePromptTemplate } from 'langchain/prompts'
 import { BufferMemory, ChatMessageHistory } from 'langchain/memory'
 import { BaseChatModel } from 'langchain/chat_models/base'
@@ -89,8 +89,14 @@ class ConversationChain_Chains implements INode {
             chain.memory = memory
         }
 
-        const res = await chain.call({ input })
-        return res?.response
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            const res = await chain.call({ input }, [handler])
+            return res?.response
+        } else {
+            const res = await chain.call({ input })
+            return res?.text
+        }
     }
 }
 

packages/components/nodes/chains/ConversationalRetrievalQAChain/ConversationalRetrievalQAChain.ts

@@ -1,6 +1,6 @@
 import { BaseLanguageModel } from 'langchain/base_language'
 import { ICommonObject, IMessage, INode, INodeData, INodeParams } from '../../../src/Interface'
-import { getBaseClasses } from '../../../src/utils'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { ConversationalRetrievalQAChain } from 'langchain/chains'
 import { BaseRetriever } from 'langchain/schema'
 
@@ -48,6 +48,12 @@ class ConversationalRetrievalQAChain_Chains implements INode {
 
     async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
         const chain = nodeData.instance as ConversationalRetrievalQAChain
+        let model = nodeData.inputs?.model
+
+        // Temporary fix: https://github.com/hwchase17/langchainjs/issues/754
+        model.streaming = false
+        chain.questionGeneratorChain.llm = model
+
         let chatHistory = ''
 
         if (options && options.chatHistory) {
@@ -64,9 +70,14 @@ class ConversationalRetrievalQAChain_Chains implements INode {
             chat_history: chatHistory ? chatHistory : []
         }
 
-        const res = await chain.call(obj)
-
-        return res?.text
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            const res = await chain.call(obj, [handler])
+            return res?.text
+        } else {
+            const res = await chain.call(obj)
+            return res?.text
+        }
     }
 }
 

packages/components/nodes/chains/LLMChain/LLMChain.ts

@@ -1,5 +1,5 @@
 import { ICommonObject, INode, INodeData, INodeOutputsValue, INodeParams } from '../../../src/Interface'
-import { getBaseClasses } from '../../../src/utils'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { LLMChain } from 'langchain/chains'
 import { BaseLanguageModel } from 'langchain/base_language'
 
@@ -76,12 +76,14 @@ class LLMChain_Chains implements INode {
         }
     }
 
-    async run(nodeData: INodeData, input: string): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
         const inputVariables = nodeData.instance.prompt.inputVariables as string[] // ["product"]
         const chain = nodeData.instance as LLMChain
         const promptValues = nodeData.inputs?.prompt.promptValues as ICommonObject
 
-        const res = await runPrediction(inputVariables, chain, input, promptValues)
+        const res = options.socketIO
+            ? await runPrediction(inputVariables, chain, input, promptValues, true, options.socketIO, options.socketIOClientId)
+            : await runPrediction(inputVariables, chain, input, promptValues)
         // eslint-disable-next-line no-console
         console.log('\x1b[93m\x1b[1m\n*****FINAL RESULT*****\n\x1b[0m\x1b[0m')
         // eslint-disable-next-line no-console
@@ -90,10 +92,24 @@ class LLMChain_Chains implements INode {
     }
 }
 
-const runPrediction = async (inputVariables: string[], chain: LLMChain, input: string, promptValues: ICommonObject) => {
+const runPrediction = async (
+    inputVariables: string[],
+    chain: LLMChain,
+    input: string,
+    promptValues: ICommonObject,
+    isStreaming?: boolean,
+    socketIO?: any,
+    socketIOClientId = ''
+) => {
     if (inputVariables.length === 1) {
-        const res = await chain.run(input)
-        return res
+        if (isStreaming) {
+            const handler = new CustomChainHandler(socketIO, socketIOClientId)
+            const res = await chain.run(input, [handler])
+            return res
+        } else {
+            const res = await chain.run(input)
+            return res
+        }
     } else if (inputVariables.length > 1) {
         let seen: string[] = []
 
@@ -109,8 +125,14 @@ const runPrediction = async (inputVariables: string[], chain: LLMChain, input: s
             const options = {
                 ...promptValues
             }
-            const res = await chain.call(options)
-            return res?.text
+            if (isStreaming) {
+                const handler = new CustomChainHandler(socketIO, socketIOClientId)
+                const res = await chain.call(options, [handler])
+                return res?.text
+            } else {
+                const res = await chain.call(options)
+                return res?.text
+            }
         } else if (seen.length === 1) {
             // If one inputVariable is not specify, use input (user's question) as value
             const lastValue = seen.pop()
@@ -119,14 +141,26 @@ const runPrediction = async (inputVariables: string[], chain: LLMChain, input: s
                 ...promptValues,
                 [lastValue]: input
             }
-            const res = await chain.call(options)
-            return res?.text
+            if (isStreaming) {
+                const handler = new CustomChainHandler(socketIO, socketIOClientId)
+                const res = await chain.call(options, [handler])
+                return res?.text
+            } else {
+                const res = await chain.call(options)
+                return res?.text
+            }
         } else {
             throw new Error(`Please provide Prompt Values for: ${seen.join(', ')}`)
         }
     } else {
-        const res = await chain.run(input)
-        return res
+        if (isStreaming) {
+            const handler = new CustomChainHandler(socketIO, socketIOClientId)
+            const res = await chain.run(input, [handler])
+            return res
+        } else {
+            const res = await chain.run(input)
+            return res
+        }
     }
 }
 

packages/components/nodes/chains/RetrievalQAChain/RetrievalQAChain.ts

@@ -1,7 +1,7 @@
-import { INode, INodeData, INodeParams } from '../../../src/Interface'
+import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
 import { RetrievalQAChain } from 'langchain/chains'
 import { BaseRetriever } from 'langchain/schema'
-import { getBaseClasses } from '../../../src/utils'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { BaseLanguageModel } from 'langchain/base_language'
 
 class RetrievalQAChain_Chains implements INode {
@@ -44,13 +44,20 @@ class RetrievalQAChain_Chains implements INode {
         return chain
     }
 
-    async run(nodeData: INodeData, input: string): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
         const chain = nodeData.instance as RetrievalQAChain
         const obj = {
             query: input
         }
-        const res = await chain.call(obj)
-        return res?.text
+
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            const res = await chain.call(obj, [handler])
+            return res?.text
+        } else {
+            const res = await chain.call(obj)
+            return res?.text
+        }
     }
 }
 

packages/components/nodes/chains/SqlDatabaseChain/SqlDatabaseChain.ts

@@ -1,6 +1,6 @@
-import { INode, INodeData, INodeParams } from '../../../src/Interface'
+import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
 import { SqlDatabaseChain, SqlDatabaseChainInput } from 'langchain/chains'
-import { getBaseClasses } from '../../../src/utils'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { DataSource } from 'typeorm'
 import { SqlDatabase } from 'langchain/sql_db'
 import { BaseLanguageModel } from 'langchain/base_language'
@@ -59,14 +59,20 @@ class SqlDatabaseChain_Chains implements INode {
         return chain
     }
 
-    async run(nodeData: INodeData, input: string): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
         const databaseType = nodeData.inputs?.database as 'sqlite'
         const model = nodeData.inputs?.model as BaseLanguageModel
         const dbFilePath = nodeData.inputs?.dbFilePath
 
         const chain = await getSQLDBChain(databaseType, dbFilePath, model)
-        const res = await chain.run(input)
-        return res
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            const res = await chain.run(input, [handler])
+            return res
+        } else {
+            const res = await chain.run(input)
+            return res
+        }
     }
 }
 

packages/components/nodes/chains/VectorDBQAChain/VectorDBQAChain.ts

@@ -1,5 +1,5 @@
-import { INode, INodeData, INodeParams } from '../../../src/Interface'
-import { getBaseClasses } from '../../../src/utils'
+import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
+import { CustomChainHandler, getBaseClasses } from '../../../src/utils'
 import { VectorDBQAChain } from 'langchain/chains'
 import { BaseLanguageModel } from 'langchain/base_language'
 import { VectorStore } from 'langchain/vectorstores'
@@ -44,13 +44,20 @@ class VectorDBQAChain_Chains implements INode {
         return chain
     }
 
-    async run(nodeData: INodeData, input: string): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
         const chain = nodeData.instance as VectorDBQAChain
         const obj = {
             query: input
         }
-        const res = await chain.call(obj)
-        return res?.text
+
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            const res = await chain.call(obj, [handler])
+            return res?.text
+        } else {
+            const res = await chain.call(obj)
+            return res?.text
+        }
     }
 }
 

packages/components/nodes/chatmodels/AzureChatOpenAI/AzureChatOpenAI.ts

@@ -121,6 +121,7 @@ class AzureChatOpenAI_ChatModels implements INode {
         const frequencyPenalty = nodeData.inputs?.frequencyPenalty as string
         const presencePenalty = nodeData.inputs?.presencePenalty as string
         const timeout = nodeData.inputs?.timeout as string
+        const streaming = nodeData.inputs?.streaming as boolean
 
         const obj: Partial<AzureOpenAIInput> & Partial<OpenAIBaseInput> = {
             temperature: parseInt(temperature, 10),
@@ -128,7 +129,8 @@ class AzureChatOpenAI_ChatModels implements INode {
             azureOpenAIApiKey,
             azureOpenAIApiInstanceName,
             azureOpenAIApiDeploymentName,
-            azureOpenAIApiVersion
+            azureOpenAIApiVersion,
+            streaming: streaming ?? true
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)

packages/components/nodes/chatmodels/ChatAnthropic/ChatAnthropic.ts

@@ -117,11 +117,13 @@ class ChatAnthropic_ChatModels implements INode {
         const maxTokensToSample = nodeData.inputs?.maxTokensToSample as string
         const topP = nodeData.inputs?.topP as string
         const topK = nodeData.inputs?.topK as string
+        const streaming = nodeData.inputs?.streaming as boolean
 
         const obj: Partial<AnthropicInput> & { anthropicApiKey?: string } = {
             temperature: parseInt(temperature, 10),
             modelName,
-            anthropicApiKey
+            anthropicApiKey,
+            streaming: streaming ?? true
         }
 
         if (maxTokensToSample) obj.maxTokensToSample = parseInt(maxTokensToSample, 10)

packages/components/nodes/chatmodels/ChatOpenAI/ChatOpenAI.ts

@@ -109,11 +109,13 @@ class ChatOpenAI_ChatModels implements INode {
         const frequencyPenalty = nodeData.inputs?.frequencyPenalty as string
         const presencePenalty = nodeData.inputs?.presencePenalty as string
         const timeout = nodeData.inputs?.timeout as string
+        const streaming = nodeData.inputs?.streaming as boolean
 
         const obj: Partial<OpenAIChatInput> & { openAIApiKey?: string } = {
             temperature: parseInt(temperature, 10),
             modelName,
-            openAIApiKey
+            openAIApiKey,
+            streaming: streaming ?? true
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)

packages/components/nodes/llms/Azure OpenAI/AzureOpenAI.ts

@@ -176,6 +176,7 @@ class AzureOpenAI_LLMs implements INode {
         const presencePenalty = nodeData.inputs?.presencePenalty as string
         const timeout = nodeData.inputs?.timeout as string
         const bestOf = nodeData.inputs?.bestOf as string
+        const streaming = nodeData.inputs?.streaming as boolean
 
         const obj: Partial<AzureOpenAIInput> & Partial<OpenAIInput> = {
             temperature: parseInt(temperature, 10),
@@ -183,7 +184,8 @@ class AzureOpenAI_LLMs implements INode {
             azureOpenAIApiKey,
             azureOpenAIApiInstanceName,
             azureOpenAIApiDeploymentName,
-            azureOpenAIApiVersion
+            azureOpenAIApiVersion,
+            streaming: streaming ?? true
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)

packages/components/nodes/llms/OpenAI/OpenAI.ts

@@ -121,11 +121,13 @@ class OpenAI_LLMs implements INode {
         const timeout = nodeData.inputs?.timeout as string
         const batchSize = nodeData.inputs?.batchSize as string
         const bestOf = nodeData.inputs?.bestOf as string
+        const streaming = nodeData.inputs?.streaming as boolean
 
         const obj: Partial<OpenAIInput> & { openAIApiKey?: string } = {
             temperature: parseInt(temperature, 10),
             modelName,
-            openAIApiKey
+            openAIApiKey,
+            streaming: streaming ?? true
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)

packages/components/src/utils.ts

@@ -2,6 +2,8 @@ import axios from 'axios'
 import { load } from 'cheerio'
 import * as fs from 'fs'
 import * as path from 'path'
+import { BaseCallbackHandler } from 'langchain/callbacks'
+import { Server } from 'socket.io'
 
 export const numberOrExpressionRegex = '^(\\d+\\.?\\d*|{{.*}})$' //return true if string consists only numbers OR expression {{}}
 export const notEmptyRegex = '(.|\\s)*\\S(.|\\s)*' //return true if string is not empty or blank
@@ -152,6 +154,12 @@ export const getInputVariables = (paramValue: string): string[] => {
     return inputVariables
 }
 
+/**
+ * Crawl all available urls given a domain url and limit
+ * @param {string} url
+ * @param {number} limit
+ * @returns {string[]}
+ */
 export const getAvailableURLs = async (url: string, limit: number) => {
     try {
         const availableUrls: string[] = []
@@ -190,3 +198,31 @@ export const getAvailableURLs = async (url: string, limit: number) => {
         throw new Error(`getAvailableURLs: ${err?.message}`)
     }
 }
+
+/**
+ * Custom chain handler class
+ */
+export class CustomChainHandler extends BaseCallbackHandler {
+    name = 'custom_chain_handler'
+    isLLMStarted = false
+    socketIO: Server
+    socketIOClientId = ''
+
+    constructor(socketIO: Server, socketIOClientId: string) {
+        super()
+        this.socketIO = socketIO
+        this.socketIOClientId = socketIOClientId
+    }
+
+    handleLLMNewToken(token: string) {
+        if (!this.isLLMStarted) {
+            this.isLLMStarted = true
+            this.socketIO.to(this.socketIOClientId).emit('start', token)
+        }
+        this.socketIO.to(this.socketIOClientId).emit('token', token)
+    }
+
+    handleLLMEnd() {
+        this.socketIO.to(this.socketIOClientId).emit('end')
+    }
+}