Bugfix/Use vm2 for chatflow tool (#2482)

use vm2 for chatflow tool
This commit is contained in:
Henry Heng 2024-05-24 18:57:49 +01:00 committed by GitHub
parent 82899d9d5d
commit 22f39692e5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 42 additions and 12 deletions

View File

@ -1,11 +1,11 @@
import { DataSource } from 'typeorm'
import { z } from 'zod'
import fetch from 'node-fetch'
import { NodeVM } from 'vm2'
import { RunnableConfig } from '@langchain/core/runnables'
import { CallbackManagerForToolRun, Callbacks, CallbackManager, parseCallbackConfigArg } from '@langchain/core/callbacks/manager'
import { StructuredTool } from '@langchain/core/tools'
import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
import { getCredentialData, getCredentialParam } from '../../../src/utils'
import { availableDependencies, defaultAllowBuiltInDep, getCredentialData, getCredentialParam } from '../../../src/utils'
class ChatflowTool_Tools implements INode {
label: string
@ -112,6 +112,8 @@ class ChatflowTool_Tools implements INode {
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
const chatflowApiKey = getCredentialParam('chatflowApiKey', credentialData, nodeData)
if (selectedChatflowId === options.chatflowid) throw new Error('Cannot call the same chatflow!')
let headers = {}
if (chatflowApiKey) headers = { Authorization: `Bearer ${chatflowApiKey}` }
@ -226,8 +228,6 @@ class ChatflowTool extends StructuredTool {
): Promise<string> {
const inputQuestion = this.input || arg.input
const url = `${this.baseURL}/api/v1/prediction/${this.chatflowid}`
const body = {
question: inputQuestion,
chatId: flowConfig?.chatId,
@ -245,14 +245,44 @@ class ChatflowTool extends StructuredTool {
body: JSON.stringify(body)
}
try {
const response = await fetch(url, options)
const resp = await response.json()
return resp.text || ''
} catch (error) {
console.error(error)
return ''
}
let sandbox = { $callOptions: options, $callBody: body }
const code = `
const fetch = require('node-fetch');
const url = "${this.baseURL}/api/v1/prediction/${this.chatflowid}";
const body = $callBody;
const options = $callOptions;
try {
const response = await fetch(url, options);
const resp = await response.json();
return resp.text;
} catch (error) {
console.error(error);
return '';
}
`
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 vmOptions = {
console: 'inherit',
sandbox,
require: {
external: { modules: deps },
builtin: builtinDeps
}
} as any
const vm = new NodeVM(vmOptions)
const response = await vm.run(`module.exports = async function() {${code}}()`, __dirname)
return response
}
}