Skip to content

Commit 16ceed1

Browse files
Feature/Add support for state-based metadata filter to Retriever Tool (#3501)
* Added support for state-based metadata filter to Retriever Tool * Update RetrieverTool.ts --------- Co-authored-by: Henry Heng <henryheng@flowiseai.com>
1 parent 38ddbd8 commit 16ceed1

File tree

1 file changed

+151
-8
lines changed

1 file changed

+151
-8
lines changed

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

+151-8
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,121 @@
11
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'
54
import { BaseRetriever } from '@langchain/core/retrievers'
6-
import { INode, INodeData, INodeParams } from '../../../src/Interface'
5+
import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
76
import { getBaseClasses } from '../../../src/utils'
87
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+
}
9119

10120
class Retriever_Tools implements INode {
11121
label: string
@@ -22,7 +132,7 @@ class Retriever_Tools implements INode {
22132
constructor() {
23133
this.label = 'Retriever Tool'
24134
this.name = 'retrieverTool'
25-
this.version = 2.0
135+
this.version = 3.0
26136
this.type = 'RetrieverTool'
27137
this.icon = 'retrievertool.svg'
28138
this.category = 'Tools'
@@ -53,23 +163,55 @@ class Retriever_Tools implements INode {
53163
name: 'returnSourceDocuments',
54164
type: 'boolean',
55165
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+
}
56178
}
57179
]
58180
}
59181

60-
async init(nodeData: INodeData): Promise<any> {
182+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
61183
const name = nodeData.inputs?.name as string
62184
const description = nodeData.inputs?.description as string
63185
const retriever = nodeData.inputs?.retriever as BaseRetriever
64186
const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
187+
const retrieverToolMetadataFilter = nodeData.inputs?.retrieverToolMetadataFilter
65188

66189
const input = {
67190
name,
68191
description
69192
}
70193

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)
73215
const content = docs.map((doc) => doc.pageContent).join('\n\n')
74216
const sourceDocuments = JSON.stringify(docs)
75217
return returnSourceDocuments ? content + SOURCE_DOCUMENTS_PREFIX + sourceDocuments : content
@@ -80,6 +222,7 @@ class Retriever_Tools implements INode {
80222
}) as any
81223

82224
const tool = new DynamicStructuredTool({ ...input, func, schema })
225+
tool.setFlowObject(flow)
83226
return tool
84227
}
85228
}

0 commit comments

Comments
 (0)