1
1
import { z } from 'zod'
2
- import { DynamicStructuredTool } from '@langchain/core/tools'
3
- import { CallbackManagerForToolRun } from '@langchain/core/callbacks/manager'
4
- import { DynamicTool } from '@langchain/core/tools'
2
+ import { CallbackManager , CallbackManagerForToolRun , Callbacks , parseCallbackConfigArg } from '@langchain/core/callbacks/manager'
3
+ import { BaseDynamicToolInput , DynamicTool , StructuredTool , ToolInputParsingException } from '@langchain/core/tools'
5
4
import { BaseRetriever } from '@langchain/core/retrievers'
6
- import { INode , INodeData , INodeParams } from '../../../src/Interface'
5
+ import { ICommonObject , INode , INodeData , INodeParams } from '../../../src/Interface'
7
6
import { getBaseClasses } from '../../../src/utils'
8
7
import { SOURCE_DOCUMENTS_PREFIX } from '../../../src/agents'
8
+ import { RunnableConfig } from '@langchain/core/runnables'
9
+ import { customGet } from '../../sequentialagents/commonUtils'
10
+ import { VectorStoreRetriever } from '@langchain/core/vectorstores'
11
+
12
+ const howToUse = `Add additional filters to vector store. You can also filter with flow config, including the current "state":
13
+ - \`$flow.sessionId\`
14
+ - \`$flow.chatId\`
15
+ - \`$flow.chatflowId\`
16
+ - \`$flow.input\`
17
+ - \`$flow.state\`
18
+ `
19
+
20
+ type ZodObjectAny = z . ZodObject < any , any , any , any >
21
+ type IFlowConfig = { sessionId ?: string ; chatId ?: string ; input ?: string ; state ?: ICommonObject }
22
+ interface DynamicStructuredToolInput < T extends z . ZodObject < any , any , any , any > = z . ZodObject < any , any , any , any > >
23
+ extends BaseDynamicToolInput {
24
+ func ?: ( input : z . infer < T > , runManager ?: CallbackManagerForToolRun , flowConfig ?: IFlowConfig ) => Promise < string >
25
+ schema : T
26
+ }
27
+
28
+ class DynamicStructuredTool < T extends z . ZodObject < any , any , any , any > = z . ZodObject < any , any , any , any > > extends StructuredTool <
29
+ T extends ZodObjectAny ? T : ZodObjectAny
30
+ > {
31
+ static lc_name ( ) {
32
+ return 'DynamicStructuredTool'
33
+ }
34
+
35
+ name : string
36
+
37
+ description : string
38
+
39
+ func : DynamicStructuredToolInput [ 'func' ]
40
+
41
+ // @ts -ignore
42
+ schema : T
43
+
44
+ private flowObj : any
45
+
46
+ constructor ( fields : DynamicStructuredToolInput < T > ) {
47
+ super ( fields )
48
+ this . name = fields . name
49
+ this . description = fields . description
50
+ this . func = fields . func
51
+ this . returnDirect = fields . returnDirect ?? this . returnDirect
52
+ this . schema = fields . schema
53
+ }
54
+
55
+ async call ( arg : any , configArg ?: RunnableConfig | Callbacks , tags ?: string [ ] , flowConfig ?: IFlowConfig ) : Promise < string > {
56
+ const config = parseCallbackConfigArg ( configArg )
57
+ if ( config . runName === undefined ) {
58
+ config . runName = this . name
59
+ }
60
+ let parsed
61
+ try {
62
+ parsed = await this . schema . parseAsync ( arg )
63
+ } catch ( e ) {
64
+ throw new ToolInputParsingException ( `Received tool input did not match expected schema` , JSON . stringify ( arg ) )
65
+ }
66
+ const callbackManager_ = await CallbackManager . configure (
67
+ config . callbacks ,
68
+ this . callbacks ,
69
+ config . tags || tags ,
70
+ this . tags ,
71
+ config . metadata ,
72
+ this . metadata ,
73
+ { verbose : this . verbose }
74
+ )
75
+ const runManager = await callbackManager_ ?. handleToolStart (
76
+ this . toJSON ( ) ,
77
+ typeof parsed === 'string' ? parsed : JSON . stringify ( parsed ) ,
78
+ undefined ,
79
+ undefined ,
80
+ undefined ,
81
+ undefined ,
82
+ config . runName
83
+ )
84
+ let result
85
+ try {
86
+ result = await this . _call ( parsed , runManager , flowConfig )
87
+ } catch ( e ) {
88
+ await runManager ?. handleToolError ( e )
89
+ throw e
90
+ }
91
+ if ( result && typeof result !== 'string' ) {
92
+ result = JSON . stringify ( result )
93
+ }
94
+ await runManager ?. handleToolEnd ( result )
95
+ return result
96
+ }
97
+
98
+ // @ts -ignore
99
+ protected _call ( arg : any , runManager ?: CallbackManagerForToolRun , flowConfig ?: IFlowConfig ) : Promise < string > {
100
+ let flowConfiguration : ICommonObject = { }
101
+ if ( typeof arg === 'object' && Object . keys ( arg ) . length ) {
102
+ for ( const item in arg ) {
103
+ flowConfiguration [ `$${ item } ` ] = arg [ item ]
104
+ }
105
+ }
106
+
107
+ // inject flow properties
108
+ if ( this . flowObj ) {
109
+ flowConfiguration [ '$flow' ] = { ...this . flowObj , ...flowConfig }
110
+ }
111
+
112
+ return this . func ! ( arg as any , runManager , flowConfiguration )
113
+ }
114
+
115
+ setFlowObject ( flow : any ) {
116
+ this . flowObj = flow
117
+ }
118
+ }
9
119
10
120
class Retriever_Tools implements INode {
11
121
label : string
@@ -22,7 +132,7 @@ class Retriever_Tools implements INode {
22
132
constructor ( ) {
23
133
this . label = 'Retriever Tool'
24
134
this . name = 'retrieverTool'
25
- this . version = 2 .0
135
+ this . version = 3 .0
26
136
this . type = 'RetrieverTool'
27
137
this . icon = 'retrievertool.svg'
28
138
this . category = 'Tools'
@@ -53,23 +163,55 @@ class Retriever_Tools implements INode {
53
163
name : 'returnSourceDocuments' ,
54
164
type : 'boolean' ,
55
165
optional : true
166
+ } ,
167
+ {
168
+ label : 'Additional Metadata Filter' ,
169
+ name : 'retrieverToolMetadataFilter' ,
170
+ type : 'json' ,
171
+ description : 'Add additional metadata filter on top of the existing filter from vector store' ,
172
+ optional : true ,
173
+ additionalParams : true ,
174
+ hint : {
175
+ label : 'What can you filter?' ,
176
+ value : howToUse
177
+ }
56
178
}
57
179
]
58
180
}
59
181
60
- async init ( nodeData : INodeData ) : Promise < any > {
182
+ async init ( nodeData : INodeData , _ : string , options : ICommonObject ) : Promise < any > {
61
183
const name = nodeData . inputs ?. name as string
62
184
const description = nodeData . inputs ?. description as string
63
185
const retriever = nodeData . inputs ?. retriever as BaseRetriever
64
186
const returnSourceDocuments = nodeData . inputs ?. returnSourceDocuments as boolean
187
+ const retrieverToolMetadataFilter = nodeData . inputs ?. retrieverToolMetadataFilter
65
188
66
189
const input = {
67
190
name,
68
191
description
69
192
}
70
193
71
- const func = async ( { input } : { input : string } , runManager ?: CallbackManagerForToolRun ) => {
72
- const docs = await retriever . getRelevantDocuments ( input , runManager ?. getChild ( 'retriever' ) )
194
+ const flow = { chatflowId : options . chatflowid }
195
+
196
+ const func = async ( { input } : { input : string } , _ ?: CallbackManagerForToolRun , flowConfig ?: IFlowConfig ) => {
197
+ if ( retrieverToolMetadataFilter ) {
198
+ const flowObj = flowConfig
199
+
200
+ const metadatafilter =
201
+ typeof retrieverToolMetadataFilter === 'object' ? retrieverToolMetadataFilter : JSON . parse ( retrieverToolMetadataFilter )
202
+ const newMetadataFilter : any = { }
203
+ for ( const key in metadatafilter ) {
204
+ let value = metadatafilter [ key ]
205
+ if ( value . startsWith ( '$flow' ) ) {
206
+ value = customGet ( flowObj , value )
207
+ }
208
+ newMetadataFilter [ key ] = value
209
+ }
210
+
211
+ const vectorStore = ( retriever as VectorStoreRetriever < any > ) . vectorStore
212
+ vectorStore . filter = newMetadataFilter
213
+ }
214
+ const docs = await retriever . invoke ( input )
73
215
const content = docs . map ( ( doc ) => doc . pageContent ) . join ( '\n\n' )
74
216
const sourceDocuments = JSON . stringify ( docs )
75
217
return returnSourceDocuments ? content + SOURCE_DOCUMENTS_PREFIX + sourceDocuments : content
@@ -80,6 +222,7 @@ class Retriever_Tools implements INode {
80
222
} ) as any
81
223
82
224
const tool = new DynamicStructuredTool ( { ...input , func, schema } )
225
+ tool . setFlowObject ( flow )
83
226
return tool
84
227
}
85
228
}
0 commit comments