add return source documens to retriever tool

packages/components/nodes/agents/OpenAIFunctionAgent/OpenAIFunctionAgent.ts

@@ -64,7 +64,7 @@ class OpenAIFunctionAgent_Agents implements INode {
         return prepareAgent(nodeData, { sessionId: this.sessionId, chatId: options.chatId, input }, options.chatHistory)
     }
 
-    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
+    async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string | ICommonObject> {
         const memory = nodeData.inputs?.memory as FlowiseMemory
         const executor = prepareAgent(nodeData, { sessionId: this.sessionId, chatId: options.chatId, input }, options.chatHistory)
 
@@ -72,12 +72,20 @@ class OpenAIFunctionAgent_Agents implements INode {
         const callbacks = await additionalCallbacks(nodeData, options)
 
         let res: ChainValues = {}
+        let sourceDocuments: ICommonObject[] = []
 
         if (options.socketIO && options.socketIOClientId) {
             const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, ...callbacks] })
+            if (res.sourceDocuments) {
+                options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))
+                sourceDocuments = res.sourceDocuments
+            }
         } else {
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, ...callbacks] })
+            if (res.sourceDocuments) {
+                sourceDocuments = res.sourceDocuments
+            }
         }
 
         await memory.addChatMessages(
@@ -94,7 +102,7 @@ class OpenAIFunctionAgent_Agents implements INode {
             this.sessionId
         )
 
-        return res?.output
+        return sourceDocuments.length ? { text: res?.output, sourceDocuments: flatten(sourceDocuments) } : res?.output
     }
 }
 

packages/components/nodes/outputparsers/CustomListOutputParser/CustomListOutputParser.ts

@@ -29,16 +29,17 @@ class CustomListOutputParser implements INode {
                 label: 'Length',
                 name: 'length',
                 type: 'number',
-                default: 5,
                 step: 1,
-                description: 'Number of values to return'
+                description: 'Number of values to return',
+                optional: true
             },
             {
                 label: 'Separator',
                 name: 'separator',
                 type: 'string',
                 description: 'Separator between values',
-                default: ','
+                default: ',',
+                optional: true
             },
             {
                 label: 'Autofix',
@@ -54,10 +55,11 @@ class CustomListOutputParser implements INode {
         const separator = nodeData.inputs?.separator as string
         const lengthStr = nodeData.inputs?.length as string
         const autoFix = nodeData.inputs?.autofixParser as boolean
-        let length = 5
-        if (lengthStr) length = parseInt(lengthStr, 10)
 
-        const parser = new LangchainCustomListOutputParser({ length: length, separator: separator })
+        const parser = new LangchainCustomListOutputParser({
+            length: lengthStr ? parseInt(lengthStr, 10) : undefined,
+            separator: separator
+        })
         Object.defineProperty(parser, 'autoFix', {
             enumerable: true,
             configurable: true,

packages/components/nodes/tools/RetrieverTool/RetrieverTool.ts

@@ -1,8 +1,11 @@
 import { INode, INodeData, INodeParams } from '../../../src/Interface'
 import { getBaseClasses } from '../../../src/utils'
 import { DynamicTool } from 'langchain/tools'
-import { createRetrieverTool } from 'langchain/agents/toolkits'
+import { DynamicStructuredTool } from '@langchain/core/tools'
+import { CallbackManagerForToolRun } from '@langchain/core/callbacks/manager'
 import { BaseRetriever } from 'langchain/schema/retriever'
+import { z } from 'zod'
+import { SOURCE_DOCUMENTS_PREFIX } from '../../../src/agents'
 
 class Retriever_Tools implements INode {
     label: string
@@ -19,7 +22,7 @@ class Retriever_Tools implements INode {
     constructor() {
         this.label = 'Retriever Tool'
         this.name = 'retrieverTool'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'RetrieverTool'
         this.icon = 'retrievertool.svg'
         this.category = 'Tools'
@@ -44,6 +47,12 @@ class Retriever_Tools implements INode {
                 label: 'Retriever',
                 name: 'retriever',
                 type: 'BaseRetriever'
+            },
+            {
+                label: 'Return Source Documents',
+                name: 'returnSourceDocuments',
+                type: 'boolean',
+                optional: true
             }
         ]
     }
@@ -52,12 +61,25 @@ class Retriever_Tools implements INode {
         const name = nodeData.inputs?.name as string
         const description = nodeData.inputs?.description as string
         const retriever = nodeData.inputs?.retriever as BaseRetriever
+        const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
 
-        const tool = createRetrieverTool(retriever, {
+        const input = {
             name,
             description
+        }
+
+        const func = async ({ input }: { input: string }, runManager?: CallbackManagerForToolRun) => {
+            const docs = await retriever.getRelevantDocuments(input, runManager?.getChild('retriever'))
+            const content = docs.map((doc) => doc.pageContent).join('\n\n')
+            const sourceDocuments = JSON.stringify(docs)
+            return returnSourceDocuments ? content + SOURCE_DOCUMENTS_PREFIX + sourceDocuments : content
+        }
+
+        const schema = z.object({
+            input: z.string().describe('query to look up in retriever')
         })
 
+        const tool = new DynamicStructuredTool({ ...input, func, schema })
         return tool
     }
 }