Skip to content

Commit

Permalink
Bugfix/Use vm2 for chatflow tool (FlowiseAI#2482)
Browse files Browse the repository at this point in the history
use vm2 for chatflow tool
  • Loading branch information
HenryHengZJ committed May 24, 2024
1 parent 82899d9 commit 22f3969
Showing 1 changed file with 42 additions and 12 deletions.
54 changes: 42 additions & 12 deletions packages/components/nodes/tools/ChatflowTool/ChatflowTool.ts
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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}` }

Expand Down Expand Up @@ -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,
Expand All @@ -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
}
}

Expand Down

0 comments on commit 22f3969

Please sign in to comment.