From 421aa712515d9beeae7c0201b173cb7324473f69 Mon Sep 17 00:00:00 2001 From: Eugene Date: Tue, 10 Sep 2024 13:48:44 +0200 Subject: [PATCH] feat(Code Tool Node): Option to specify input schema (#10693) --- .../nodes/tools/ToolCode/ToolCode.node.ts | 140 +++++++++++++----- 1 file changed, 99 insertions(+), 41 deletions(-) diff --git a/packages/@n8n/nodes-langchain/nodes/tools/ToolCode/ToolCode.node.ts b/packages/@n8n/nodes-langchain/nodes/tools/ToolCode/ToolCode.node.ts index 6ae5a88075..df68fb0c6a 100644 --- a/packages/@n8n/nodes-langchain/nodes/tools/ToolCode/ToolCode.node.ts +++ b/packages/@n8n/nodes-langchain/nodes/tools/ToolCode/ToolCode.node.ts @@ -5,16 +5,25 @@ import type { INodeTypeDescription, SupplyData, ExecutionError, + IDataObject, } from 'n8n-workflow'; -import { NodeConnectionType, NodeOperationError } from 'n8n-workflow'; +import { jsonParse, NodeConnectionType, NodeOperationError } from 'n8n-workflow'; import type { Sandbox } from 'n8n-nodes-base/dist/nodes/Code/Sandbox'; import { getSandboxContext } from 'n8n-nodes-base/dist/nodes/Code/Sandbox'; import { JavaScriptSandbox } from 'n8n-nodes-base/dist/nodes/Code/JavaScriptSandbox'; import { PythonSandbox } from 'n8n-nodes-base/dist/nodes/Code/PythonSandbox'; -import { DynamicTool } from '@langchain/core/tools'; +import { DynamicStructuredTool, DynamicTool } from '@langchain/core/tools'; import { getConnectionHintNoticeField } from '../../../utils/sharedFields'; +import { + inputSchemaField, + jsonSchemaExampleField, + schemaTypeField, +} from '../../../utils/descriptions'; +import { generateSchema, getSandboxWithZod } from '../../../utils/schemaParsing'; +import type { JSONSchema7 } from 'json-schema'; +import type { DynamicZodObject } from '../../../types/zod.types'; export class ToolCode implements INodeType { description: INodeTypeDescription = { @@ -152,6 +161,18 @@ export class ToolCode implements INodeType { description: 'E.g. Converts any text to uppercase', noDataExpression: true, }, + { + displayName: 'Specify Input Schema', + name: 'specifyInputSchema', + type: 'boolean', + description: + 'Whether to specify the schema for the function. This would require the LLM to provide the input in the correct format and would validate it against the schema.', + noDataExpression: true, + default: false, + }, + { ...schemaTypeField, displayOptions: { show: { specifyInputSchema: [true] } } }, + jsonSchemaExampleField, + inputSchemaField, ], }; @@ -162,6 +183,8 @@ export class ToolCode implements INodeType { const name = this.getNodeParameter('name', itemIndex) as string; const description = this.getNodeParameter('description', itemIndex) as string; + const useSchema = this.getNodeParameter('specifyInputSchema', itemIndex) as boolean; + const language = this.getNodeParameter('language', itemIndex) as string; let code = ''; if (language === 'javaScript') { @@ -170,7 +193,7 @@ export class ToolCode implements INodeType { code = this.getNodeParameter('pythonCode', itemIndex) as string; } - const getSandbox = (query: string, index = 0) => { + const getSandbox = (query: string | IDataObject, index = 0) => { const context = getSandboxContext.call(this, index); context.query = query; @@ -191,49 +214,84 @@ export class ToolCode implements INodeType { return sandbox; }; - const runFunction = async (query: string): Promise => { + const runFunction = async (query: string | IDataObject): Promise => { const sandbox = getSandbox(query, itemIndex); return await (sandbox.runCode() as Promise); }; + const toolHandler = async (query: string | IDataObject): Promise => { + const { index } = this.addInputData(NodeConnectionType.AiTool, [[{ json: { query } }]]); + + let response: string = ''; + let executionError: ExecutionError | undefined; + try { + response = await runFunction(query); + } catch (error: unknown) { + executionError = new NodeOperationError(this.getNode(), error as ExecutionError); + response = `There was an error: "${executionError.message}"`; + } + + if (typeof response === 'number') { + response = (response as number).toString(); + } + + if (typeof response !== 'string') { + // TODO: Do some more testing. Issues here should actually fail the workflow + executionError = new NodeOperationError(this.getNode(), 'Wrong output type returned', { + description: `The response property should be a string, but it is an ${typeof response}`, + }); + response = `There was an error: "${executionError.message}"`; + } + + if (executionError) { + void this.addOutputData(NodeConnectionType.AiTool, index, executionError); + } else { + void this.addOutputData(NodeConnectionType.AiTool, index, [[{ json: { response } }]]); + } + + return response; + }; + + const commonToolOptions = { + name, + description, + func: toolHandler, + }; + + let tool: DynamicTool | DynamicStructuredTool | undefined = undefined; + + if (useSchema) { + try { + // We initialize these even though one of them will always be empty + // it makes it easier to navigate the ternary operator + const jsonExample = this.getNodeParameter('jsonSchemaExample', itemIndex, '') as string; + const inputSchema = this.getNodeParameter('inputSchema', itemIndex, '') as string; + + const schemaType = this.getNodeParameter('schemaType', itemIndex) as 'fromJson' | 'manual'; + const jsonSchema = + schemaType === 'fromJson' + ? generateSchema(jsonExample) + : jsonParse(inputSchema); + + const zodSchemaSandbox = getSandboxWithZod(this, jsonSchema, 0); + const zodSchema = (await zodSchemaSandbox.runCode()) as DynamicZodObject; + + tool = new DynamicStructuredTool({ + schema: zodSchema, + ...commonToolOptions, + }); + } catch (error) { + throw new NodeOperationError( + this.getNode(), + 'Error during parsing of JSON Schema. \n ' + error, + ); + } + } else { + tool = new DynamicTool(commonToolOptions); + } + return { - response: new DynamicTool({ - name, - description, - - func: async (query: string): Promise => { - const { index } = this.addInputData(NodeConnectionType.AiTool, [[{ json: { query } }]]); - - let response: string = ''; - let executionError: ExecutionError | undefined; - try { - response = await runFunction(query); - } catch (error: unknown) { - executionError = new NodeOperationError(this.getNode(), error as ExecutionError); - response = `There was an error: "${executionError.message}"`; - } - - if (typeof response === 'number') { - response = (response as number).toString(); - } - - if (typeof response !== 'string') { - // TODO: Do some more testing. Issues here should actually fail the workflow - executionError = new NodeOperationError(this.getNode(), 'Wrong output type returned', { - description: `The response property should be a string, but it is an ${typeof response}`, - }); - response = `There was an error: "${executionError.message}"`; - } - - if (executionError) { - void this.addOutputData(NodeConnectionType.AiTool, index, executionError); - } else { - void this.addOutputData(NodeConnectionType.AiTool, index, [[{ json: { response } }]]); - } - - return response; - }, - }), + response: tool, }; } }