91 lines
2.6 KiB
TypeScript
91 lines
2.6 KiB
TypeScript
import { z } from 'zod'
|
|
import { CallbackManagerForToolRun } from 'langchain/callbacks'
|
|
import { StructuredTool, ToolParams } from 'langchain/tools'
|
|
import { NodeVM } from 'vm2'
|
|
import { availableDependencies } from '../../../src/utils'
|
|
|
|
export interface BaseDynamicToolInput extends ToolParams {
|
|
name: string
|
|
description: string
|
|
code: string
|
|
returnDirect?: boolean
|
|
}
|
|
|
|
export interface DynamicStructuredToolInput<
|
|
// eslint-disable-next-line
|
|
T extends z.ZodObject<any, any, any, any> = z.ZodObject<any, any, any, any>
|
|
> extends BaseDynamicToolInput {
|
|
func?: (input: z.infer<T>, runManager?: CallbackManagerForToolRun) => Promise<string>
|
|
schema: T
|
|
}
|
|
|
|
export class DynamicStructuredTool<
|
|
// eslint-disable-next-line
|
|
T extends z.ZodObject<any, any, any, any> = z.ZodObject<any, any, any, any>
|
|
> extends StructuredTool {
|
|
name: string
|
|
|
|
description: string
|
|
|
|
code: string
|
|
|
|
func: DynamicStructuredToolInput['func']
|
|
|
|
schema: T
|
|
|
|
constructor(fields: DynamicStructuredToolInput<T>) {
|
|
super(fields)
|
|
this.name = fields.name
|
|
this.description = fields.description
|
|
this.code = fields.code
|
|
this.func = fields.func
|
|
this.returnDirect = fields.returnDirect ?? this.returnDirect
|
|
this.schema = fields.schema
|
|
}
|
|
|
|
protected async _call(arg: z.output<T>): Promise<string> {
|
|
let sandbox: any = {}
|
|
if (typeof arg === 'object' && Object.keys(arg).length) {
|
|
for (const item in arg) {
|
|
sandbox[`$${item}`] = arg[item]
|
|
}
|
|
}
|
|
|
|
const defaultAllowBuiltInDep = [
|
|
'assert',
|
|
'buffer',
|
|
'crypto',
|
|
'events',
|
|
'http',
|
|
'https',
|
|
'net',
|
|
'path',
|
|
'querystring',
|
|
'timers',
|
|
'tls',
|
|
'url',
|
|
'zlib'
|
|
]
|
|
|
|
const builtinDeps = process.env.TOOL_FUNCTION_BUILTIN_DEP
|
|
? defaultAllowBuiltInDep.concat(process.env.TOOL_FUNCTION_BUILTIN_DEP.split(','))
|
|
: defaultAllowBuiltInDep
|
|
const externalDeps = process.env.TOOL_FUNCTION_EXTERNAL_DEP ? process.env.TOOL_FUNCTION_EXTERNAL_DEP.split(',') : []
|
|
const deps = availableDependencies.concat(externalDeps)
|
|
|
|
const options = {
|
|
console: 'inherit',
|
|
sandbox,
|
|
require: {
|
|
external: { modules: deps },
|
|
builtin: builtinDeps
|
|
}
|
|
} as any
|
|
|
|
const vm = new NodeVM(options)
|
|
const response = await vm.run(`module.exports = async function() {${this.code}}()`, __dirname)
|
|
|
|
return response
|
|
}
|
|
}
|