add return source documents functioanality

packages/components/src/Interface.ts

@@ -75,7 +75,7 @@ export interface INode extends INodeProperties {
     inputs?: INodeParams[]
     output?: INodeOutputsValue[]
     init?(nodeData: INodeData, input: string, options?: ICommonObject): Promise<any>
-    run?(nodeData: INodeData, input: string, options?: ICommonObject): Promise<string>
+    run?(nodeData: INodeData, input: string, options?: ICommonObject): Promise<string | ICommonObject>
 }
 
 export interface INodeData extends INodeProperties {

packages/components/src/utils.ts

@@ -4,6 +4,7 @@ import * as fs from 'fs'
 import * as path from 'path'
 import { BaseCallbackHandler } from 'langchain/callbacks'
 import { Server } from 'socket.io'
+import { ChainValues } from 'langchain/dist/schema'
 
 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
@@ -208,12 +209,14 @@ export class CustomChainHandler extends BaseCallbackHandler {
     socketIO: Server
     socketIOClientId = ''
     skipK = 0 // Skip streaming for first K numbers of handleLLMStart
+    returnSourceDocuments = false
 
-    constructor(socketIO: Server, socketIOClientId: string, skipK?: number) {
+    constructor(socketIO: Server, socketIOClientId: string, skipK?: number, returnSourceDocuments?: boolean) {
         super()
         this.socketIO = socketIO
         this.socketIOClientId = socketIOClientId
         this.skipK = skipK ?? this.skipK
+        this.returnSourceDocuments = returnSourceDocuments ?? this.returnSourceDocuments
     }
 
     handleLLMStart() {
@@ -233,4 +236,10 @@ export class CustomChainHandler extends BaseCallbackHandler {
     handleLLMEnd() {
         this.socketIO.to(this.socketIOClientId).emit('end')
     }
+
+    handleChainEnd(outputs: ChainValues): void | Promise<void> {
+        if (this.returnSourceDocuments) {
+            this.socketIO.to(this.socketIOClientId).emit('sourceDocuments', outputs?.sourceDocuments)
+        }
+    }
 }