From bfa870e56b0371e661ef5d980ded8140b78de42a Mon Sep 17 00:00:00 2001 From: vinodkiran Date: Tue, 12 Dec 2023 12:19:35 +0530 Subject: [PATCH] Environment Variables: injection of variables ($env) into the custom tool and addition of ($flow) object --- .../nodes/tools/CustomTool/CustomTool.ts | 27 ++++++++++++++- .../components/nodes/tools/CustomTool/core.ts | 33 +++++++++++++++++-- packages/components/src/Interface.ts | 1 + packages/server/src/utils/index.ts | 4 ++- .../views/variables/AddEditVariableDialog.js | 16 ++++----- 5 files changed, 69 insertions(+), 12 deletions(-) diff --git a/packages/components/nodes/tools/CustomTool/CustomTool.ts b/packages/components/nodes/tools/CustomTool/CustomTool.ts index 541edcf0..aba803e4 100644 --- a/packages/components/nodes/tools/CustomTool/CustomTool.ts +++ b/packages/components/nodes/tools/CustomTool/CustomTool.ts @@ -80,7 +80,32 @@ class CustomTool_Tools implements INode { code: tool.func } if (customToolFunc) obj.code = customToolFunc - return new DynamicStructuredTool(obj) + + const variables = await appDataSource.getRepository(databaseEntities['Variable']).find() + + // override variables defined in overrideConfig + // nodeData.inputs.variables is an Object, check each property and override the variable + if (nodeData?.inputs?.variables) { + for (const propertyName of Object.getOwnPropertyNames(nodeData.inputs.variables)) { + const foundVar = variables.find((v) => v.name === propertyName) + if (foundVar) { + // even if the variable was defined as runtime, we override it with static value + foundVar.type = 'static' + foundVar.value = nodeData.inputs.variables[propertyName] + } else { + // add it the variables, if not found locally in the db + variables.push({ name: propertyName, type: 'static', value: nodeData.inputs.variables[propertyName] }) + } + } + } + const flow = { + chatId: options.chatId, // id is uppercase (I) + chatflowId: options.chatflowid // id is lowercase (i) + } + let dynamicStructuredTool = new DynamicStructuredTool(obj) + dynamicStructuredTool.setVariables(variables) + dynamicStructuredTool.setFlowObject(flow) + return dynamicStructuredTool } catch (e) { throw new Error(e) } diff --git a/packages/components/nodes/tools/CustomTool/core.ts b/packages/components/nodes/tools/CustomTool/core.ts index 343acafd..b7b1f6a6 100644 --- a/packages/components/nodes/tools/CustomTool/core.ts +++ b/packages/components/nodes/tools/CustomTool/core.ts @@ -2,6 +2,7 @@ import { z } from 'zod' import { CallbackManagerForToolRun } from 'langchain/callbacks' import { StructuredTool, ToolParams } from 'langchain/tools' import { NodeVM } from 'vm2' +import { logger } from "@zilliz/milvus2-sdk-node"; /* * List of dependencies allowed to be import in vm2 @@ -62,6 +63,8 @@ export class DynamicStructuredTool< func: DynamicStructuredToolInput['func'] schema: T + private variables: any[] + private flowObj: any constructor(fields: DynamicStructuredToolInput) { super(fields) @@ -80,8 +83,26 @@ export class DynamicStructuredTool< sandbox[`$${item}`] = arg[item] } } - sandbox['$env'] = { USER: 'VINOD' } - console.log('sandbox === ' + JSON.stringify(sandbox)) + //inject variables + let env = {} + if (this.variables) { + for (const item of this.variables) { + let value = item.value + if (item.type === 'runtime') { + value = process.env[item.name] + } + Object.defineProperty(env, item.name, { + enumerable: true, + configurable: true, + writable: true, + value: value + }) + } + } + sandbox['$env'] = env + if (this.flowObj) { + sandbox['$flow'] = this.flowObj + } const defaultAllowBuiltInDep = [ 'assert', 'buffer', @@ -118,4 +139,12 @@ export class DynamicStructuredTool< return response } + + setVariables(variables: any[]) { + this.variables = variables + } + + setFlowObject(flow: any) { + this.flowObj = flow + } } diff --git a/packages/components/src/Interface.ts b/packages/components/src/Interface.ts index 6752f944..bc50155c 100644 --- a/packages/components/src/Interface.ts +++ b/packages/components/src/Interface.ts @@ -73,6 +73,7 @@ export interface INodeParams { additionalParams?: boolean loadMethod?: string hidden?: boolean + variables?: ICommonObject[] } export interface INodeExecutionData { diff --git a/packages/server/src/utils/index.ts b/packages/server/src/utils/index.ts index 2bf1c04a..adec01cd 100644 --- a/packages/server/src/utils/index.ts +++ b/packages/server/src/utils/index.ts @@ -37,6 +37,7 @@ import { Tool } from '../database/entities/Tool' import { Assistant } from '../database/entities/Assistant' import { DataSource } from 'typeorm' import { CachePool } from '../CachePool' +import { Variable } from '../database/entities/Variable' const QUESTION_VAR_PREFIX = 'question' const CHAT_HISTORY_VAR_PREFIX = 'chat_history' @@ -47,7 +48,8 @@ export const databaseEntities: IDatabaseEntity = { ChatMessage: ChatMessage, Tool: Tool, Credential: Credential, - Assistant: Assistant + Assistant: Assistant, + Variable: Variable } /** diff --git a/packages/ui/src/views/variables/AddEditVariableDialog.js b/packages/ui/src/views/variables/AddEditVariableDialog.js index e32ce663..db2116b0 100644 --- a/packages/ui/src/views/variables/AddEditVariableDialog.js +++ b/packages/ui/src/views/variables/AddEditVariableDialog.js @@ -62,13 +62,13 @@ const AddEditVariableDialog = ({ show, dialogProps, onCancel, onConfirm }) => { setName(dialogProps.data.name) setValue(dialogProps.data.value) setType(dialogProps.data.type) - //setVariable(dialogProps.data) + setVariable(dialogProps.data) } else if (dialogProps.type === 'ADD' && dialogProps.data) { // When variable dialog is to add a new variable setName('') setValue('') setType('static') - //setVariable({ name: '', value: '', type: 'static' }) + setVariable({ name: '', value: '', type: 'static' }) } // eslint-disable-next-line react-hooks/exhaustive-deps @@ -83,9 +83,9 @@ const AddEditVariableDialog = ({ show, dialogProps, onCancel, onConfirm }) => { const addNewVariable = async () => { try { const obj = { - name, - value, - type + name: name, + value: value, + type: type } const createResp = await variablesApi.createVariable(obj) if (createResp.data) { @@ -125,9 +125,9 @@ const AddEditVariableDialog = ({ show, dialogProps, onCancel, onConfirm }) => { const saveVariable = async () => { try { const saveObj = { - name, - value, - type + name: name, + value: value, + type: type } const saveResp = await variablesApi.updateVariable(variable.id, saveObj)