Skip to content

Commit 7356349

Browse files
HenryHengZJHenryHengZJ
authored
Update RetrieverTool.ts
1 parent b954d60 commit 7356349

File tree

1 file changed

+144
-47
lines changed

1 file changed

+144
-47
lines changed

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

Lines changed: 144 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,119 @@
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 { VectorStoreRetriever } from '@langchain/core/vectorstores'
7-
import { INode, INodeData, INodeParams } from '../../../src/Interface'
5+
import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
86
import { getBaseClasses } from '../../../src/utils'
97
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
1043

11-
class FlowAwareDynamicStructuredTool<T extends z.ZodObject<any, any, any, any>> extends DynamicStructuredTool<T> {
1244
private flowObj: any
1345

14-
setFlowObject(flow: any) {
15-
this.flowObj = flow
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
1653
}
1754

18-
getFlowObject() {
19-
return this.flowObj
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
20117
}
21118
}
22119

@@ -35,7 +132,7 @@ class Retriever_Tools implements INode {
35132
constructor() {
36133
this.label = 'Retriever Tool'
37134
this.name = 'retrieverTool'
38-
this.version = 2.1
135+
this.version = 3.0
39136
this.type = 'RetrieverTool'
40137
this.icon = 'retrievertool.svg'
41138
this.category = 'Tools'
@@ -56,13 +153,6 @@ class Retriever_Tools implements INode {
56153
rows: 3,
57154
placeholder: 'Searches and returns documents regarding the state-of-the-union.'
58155
},
59-
{
60-
label: 'Filter Property',
61-
name: 'filterProperty',
62-
optional: true,
63-
type: 'string',
64-
description: 'State property that will be used to filter the documents'
65-
},
66156
{
67157
label: 'Retriever',
68158
name: 'retriever',
@@ -73,49 +163,55 @@ class Retriever_Tools implements INode {
73163
name: 'returnSourceDocuments',
74164
type: 'boolean',
75165
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+
}
76178
}
77179
]
78180
}
79181

80-
private populateFilterFromState(tool: FlowAwareDynamicStructuredTool<any>, filterProperty: string, retriever: BaseRetriever) {
81-
const flowState = tool.getFlowObject().state
82-
if (!flowState[filterProperty]) {
83-
throw new Error(`Property '${filterProperty}' is not defined in the state`)
84-
}
85-
86-
const vectorStore = (retriever as VectorStoreRetriever<any>).vectorStore
87-
if (vectorStore.filter && Object.keys(vectorStore.filter).length > 0) {
88-
throw new Error(`Default VectorStore filter is not empty: ${JSON.stringify(vectorStore.filter)}`)
89-
}
90-
91-
const filter = flowState[filterProperty]
92-
if (!filter || typeof filter !== 'object') {
93-
throw new Error(`Filter property '${filterProperty}' must be an object, but it is ${typeof filter}`)
94-
}
95-
96-
vectorStore.filter = filter
97-
}
98-
99-
async init(nodeData: INodeData): Promise<any> {
182+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
100183
const name = nodeData.inputs?.name as string
101184
const description = nodeData.inputs?.description as string
102-
const filterProperty = nodeData.inputs?.filterProperty as string
103185
const retriever = nodeData.inputs?.retriever as BaseRetriever
104186
const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments as boolean
187+
const retrieverToolMetadataFilter = nodeData.inputs?.retrieverToolMetadataFilter
105188

106189
const input = {
107190
name,
108191
description
109192
}
110-
let tool: FlowAwareDynamicStructuredTool<any>
111193

112-
const func = async ({ input }: { input: string }, runManager?: CallbackManagerForToolRun) => {
113-
if (filterProperty) {
114-
// tool has a filter property set - populate the filter from the state
115-
this.populateFilterFromState(tool, filterProperty, 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
116213
}
117-
118-
const docs = await retriever.getRelevantDocuments(input, runManager?.getChild('retriever'))
214+
const docs = await retriever.invoke(input)
119215
const content = docs.map((doc) => doc.pageContent).join('\n\n')
120216
const sourceDocuments = JSON.stringify(docs)
121217
return returnSourceDocuments ? content + SOURCE_DOCUMENTS_PREFIX + sourceDocuments : content
@@ -125,7 +221,8 @@ class Retriever_Tools implements INode {
125221
input: z.string().describe('input to look up in retriever')
126222
}) as any
127223

128-
tool = new FlowAwareDynamicStructuredTool({ ...input, func, schema })
224+
const tool = new DynamicStructuredTool({ ...input, func, schema })
225+
tool.setFlowObject(flow)
129226
return tool
130227
}
131228
}

0 commit comments

Comments
 (0)