From 5f69a0652ea6d5914ccf1243551cf0a346940395 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Wed, 13 Mar 2024 15:11:22 +0800
Subject: [PATCH] update xml output parser

---
 .../nodes/agents/XMLAgent/XMLAgent.ts         |  3 +-
 packages/components/src/agents.ts             | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/packages/components/nodes/agents/XMLAgent/XMLAgent.ts b/packages/components/nodes/agents/XMLAgent/XMLAgent.ts
index 175dcefd8..753f5e76b 100644
--- a/packages/components/nodes/agents/XMLAgent/XMLAgent.ts
+++ b/packages/components/nodes/agents/XMLAgent/XMLAgent.ts
@@ -5,12 +5,11 @@ import { BaseChatModel } from '@langchain/core/language_models/chat_models'
 import { RunnableSequence } from '@langchain/core/runnables'
 import { Tool } from '@langchain/core/tools'
 import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder } from '@langchain/core/prompts'
-import { XMLAgentOutputParser } from 'langchain/agents/xml/output_parser'
 import { formatLogToMessage } from 'langchain/agents/format_scratchpad/log_to_message'
 import { getBaseClasses } from '../../../src/utils'
 import { FlowiseMemory, ICommonObject, IMessage, INode, INodeData, INodeParams } from '../../../src/Interface'
 import { ConsoleCallbackHandler, CustomChainHandler, additionalCallbacks } from '../../../src/handler'
-import { AgentExecutor } from '../../../src/agents'
+import { AgentExecutor, XMLAgentOutputParser } from '../../../src/agents'
 import { Moderation, checkInputs } from '../../moderation/Moderation'
 import { formatResponse } from '../../outputparsers/OutputParserHelpers'
 
diff --git a/packages/components/src/agents.ts b/packages/components/src/agents.ts
index 41bc076d6..5e2c60fca 100644
--- a/packages/components/src/agents.ts
+++ b/packages/components/src/agents.ts
@@ -779,3 +779,36 @@ Final Answer: the final answer to the original input question`
         })
     }
 }
+
+export class XMLAgentOutputParser extends AgentActionOutputParser {
+    lc_namespace = ['langchain', 'agents', 'xml']
+
+    static lc_name() {
+        return 'XMLAgentOutputParser'
+    }
+
+    /**
+     * Parses the output text from the agent and returns an AgentAction or
+     * AgentFinish object.
+     * @param text The output text from the agent.
+     * @returns An AgentAction or AgentFinish object.
+     */
+    async parse(text: string): Promise<AgentAction | AgentFinish> {
+        if (text.includes('</tool>')) {
+            const [tool, toolInput] = text.split('</tool>')
+            const _tool = tool.split('<tool>')[1]
+            const _toolInput = toolInput.split('<tool_input>')[1]
+            return { tool: _tool, toolInput: _toolInput, log: text }
+        } else if (text.includes('<final_answer>')) {
+            const [, answer] = text.split('<final_answer>')
+            return { returnValues: { output: answer }, log: text }
+        } else {
+            // Instead of throwing Error, we return a AgentFinish object
+            return { returnValues: { output: text }, log: text }
+        }
+    }
+
+    getFormatInstructions(): string {
+        throw new Error('getFormatInstructions not implemented inside XMLAgentOutputParser.')
+    }
+}