change agent/chain with memory to use runnable

2026-06-28 19:00:59 +03:00 · 2024-01-08 13:02:56 +00:00
parent d5b8db5599
commit 02482f1b38
38 changed files with 1752 additions and 1394 deletions
@@ -1,11 +1,14 @@
-import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
-import { initializeAgentExecutorWithOptions, AgentExecutor, InitializeAgentExecutorOptions } from 'langchain/agents'
 import { Tool } from 'langchain/tools'
-import { BaseChatMemory } from 'langchain/memory'
-import { getBaseClasses, mapChatHistory } from '../../../src/utils'
 import { BaseChatModel } from 'langchain/chat_models/base'
 import { flatten } from 'lodash'
-import { additionalCallbacks } from '../../../src/handler'
+import { AgentStep, BaseMessage, ChainValues, AIMessage, HumanMessage } from 'langchain/schema'
+import { RunnableSequence } from 'langchain/schema/runnable'
+import { getBaseClasses } from '../../../src/utils'
+import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
+import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams } from '../../../src/Interface'
+import { AgentExecutor } from '../../../src/agents'
+import { ChatConversationalAgent } from 'langchain/agents'
+import { renderTemplate } from '@langchain/core/prompts'

 const DEFAULT_PREFIX = `Assistant is a large language model trained by OpenAI.

@@ -15,6 +18,15 @@ Assistant is constantly learning and improving, and its capabilities are constan

 Overall, Assistant is a powerful system that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.`

+const TEMPLATE_TOOL_RESPONSE = `TOOL RESPONSE:
+---------------------
+{observation}
+
+USER'S INPUT
+--------------------
+
+Okay, so what is the response to my last comment? If using information obtained from the tools you must mention it explicitly without mentioning the tool names - I have forgotten all TOOL RESPONSES! Remember to respond with a markdown code snippet of a json blob with a single action, and NOTHING else.`
+
 class ConversationalAgent_Agents implements INode {
    label: string
    name: string
@@ -25,8 +37,9 @@ class ConversationalAgent_Agents implements INode {
    category: string
    baseClasses: string[]
    inputs: INodeParams[]
+    sessionId?: string

-    constructor() {
+    constructor(fields?: { sessionId?: string }) {
        this.label = 'Conversational Agent'
        this.name = 'conversationalAgent'
        this.version = 2.0
@@ -43,7 +56,7 @@ class ConversationalAgent_Agents implements INode {
                list: true
            },
            {
-                label: 'Language Model',
+                label: 'Chat Model',
                name: 'model',
                type: 'BaseChatModel'
            },
@@ -62,52 +75,114 @@ class ConversationalAgent_Agents implements INode {
                additionalParams: true
            }
        ]
+        this.sessionId = fields?.sessionId
    }

-    async init(nodeData: INodeData): Promise<any> {
-        const model = nodeData.inputs?.model as BaseChatModel
-        let tools = nodeData.inputs?.tools as Tool[]
-        tools = flatten(tools)
-        const memory = nodeData.inputs?.memory as BaseChatMemory
-        const systemMessage = nodeData.inputs?.systemMessage as string
-
-        const obj: InitializeAgentExecutorOptions = {
-            agentType: 'chat-conversational-react-description',
-            verbose: process.env.DEBUG === 'true' ? true : false
-        }
-
-        const agentArgs: any = {}
-        if (systemMessage) {
-            agentArgs.systemMessage = systemMessage
-        }
-
-        if (Object.keys(agentArgs).length) obj.agentArgs = agentArgs
-
-        const executor = await initializeAgentExecutorWithOptions(tools, model, obj)
-        executor.memory = memory
-        return executor
+    async init(nodeData: INodeData, input: string, options: ICommonObject): Promise<any> {
+        return prepareAgent(nodeData, { sessionId: this.sessionId, chatId: options.chatId, input }, options.chatHistory)
    }

    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
-        const executor = nodeData.instance as AgentExecutor
-        const memory = nodeData.inputs?.memory as BaseChatMemory
-
-        if (options && options.chatHistory) {
-            const chatHistoryClassName = memory.chatHistory.constructor.name
-            // Only replace when its In-Memory
-            if (chatHistoryClassName && chatHistoryClassName === 'ChatMessageHistory') {
-                memory.chatHistory = mapChatHistory(options)
-                executor.memory = memory
-            }
-        }
-
-        ;(executor.memory as any).returnMessages = true // Return true for BaseChatModel
+        const memory = nodeData.inputs?.memory as FlowiseMemory
+        const executor = await prepareAgent(nodeData, { sessionId: this.sessionId, chatId: options.chatId, input }, options.chatHistory)

+        const loggerHandler = new ConsoleCallbackHandler(options.logger)
        const callbacks = await additionalCallbacks(nodeData, options)

-        const result = await executor.call({ input }, [...callbacks])
-        return result?.output
+        let res: ChainValues = {}
+
+        if (options.socketIO && options.socketIOClientId) {
+            const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
+            res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
+        } else {
+            res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
+        }
+
+        await memory.addChatMessages(
+            [
+                {
+                    text: input,
+                    type: 'userMessage'
+                },
+                {
+                    text: res?.output,
+                    type: 'apiMessage'
+                }
+            ],
+            this.sessionId
+        )
+
+        return res?.output
    }
 }

+const prepareAgent = async (
+    nodeData: INodeData,
+    flowObj: { sessionId?: string; chatId?: string; input?: string },
+    chatHistory: IMessage[] = []
+) => {
+    const model = nodeData.inputs?.model as BaseChatModel
+    let tools = nodeData.inputs?.tools as Tool[]
+    tools = flatten(tools)
+    const memory = nodeData.inputs?.memory as FlowiseMemory
+    const systemMessage = nodeData.inputs?.systemMessage as string
+    const memoryKey = memory.memoryKey ? memory.memoryKey : 'chat_history'
+    const inputKey = memory.inputKey ? memory.inputKey : 'input'
+
+    /** Bind a stop token to the model */
+    const modelWithStop = model.bind({
+        stop: ['\nObservation']
+    })
+
+    const outputParser = ChatConversationalAgent.getDefaultOutputParser({
+        llm: model,
+        toolNames: tools.map((tool) => tool.name)
+    })
+
+    const prompt = ChatConversationalAgent.createPrompt(tools, {
+        systemMessage: systemMessage ? systemMessage : DEFAULT_PREFIX,
+        outputParser
+    })
+
+    const runnableAgent = RunnableSequence.from([
+        {
+            [inputKey]: (i: { input: string; steps: AgentStep[] }) => i.input,
+            agent_scratchpad: async (i: { input: string; steps: AgentStep[] }) => await constructScratchPad(i.steps),
+            [memoryKey]: async (_: { input: string; steps: AgentStep[] }) => {
+                const messages = (await memory.getChatMessages(flowObj?.sessionId, true, chatHistory)) as BaseMessage[]
+                return messages ?? []
+            }
+        },
+        prompt,
+        modelWithStop,
+        outputParser
+    ])
+
+    const executor = AgentExecutor.fromAgentAndTools({
+        agent: runnableAgent,
+        tools,
+        sessionId: flowObj?.sessionId,
+        chatId: flowObj?.chatId,
+        input: flowObj?.input,
+        verbose: process.env.DEBUG === 'true' ? true : false
+    })
+
+    return executor
+}
+
+const constructScratchPad = async (steps: AgentStep[]): Promise<BaseMessage[]> => {
+    const thoughts: BaseMessage[] = []
+    for (const step of steps) {
+        thoughts.push(new AIMessage(step.action.log))
+        thoughts.push(
+            new HumanMessage(
+                renderTemplate(TEMPLATE_TOOL_RESPONSE, 'f-string', {
+                    observation: step.observation
+                })
+            )
+        )
+    }
+    return thoughts
+}
+
 module.exports = { nodeClass: ConversationalAgent_Agents }