Skip to content

Commit

Permalink
fix(core): Wrap nodes for tool use at a suitable time (#11238)
Browse files Browse the repository at this point in the history
Co-authored-by: कारतोफ्फेलस्क्रिप्ट™ <aditya@netroy.in>
  • Loading branch information
jeanpaul and netroy authored Oct 15, 2024
1 parent 190665d commit c2fb881
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 7 deletions.
100 changes: 100 additions & 0 deletions packages/cli/src/__tests__/node-types.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
import { mock } from 'jest-mock-extended';
import type { INodeType, IVersionedNodeType } from 'n8n-workflow';

import type { LoadNodesAndCredentials } from '@/load-nodes-and-credentials';

import { NodeTypes } from '../node-types';

describe('NodeTypes', () => {
let nodeTypes: NodeTypes;
const loadNodesAndCredentials = mock<LoadNodesAndCredentials>();

beforeEach(() => {
jest.clearAllMocks();
nodeTypes = new NodeTypes(loadNodesAndCredentials);
});

describe('getByNameAndVersion', () => {
const nodeTypeName = 'n8n-nodes-base.testNode';

it('should throw an error if the node-type does not exist', () => {
const nodeTypeName = 'unknownNode';

// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.loadedNodes = {};
// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.knownNodes = {};

expect(() => nodeTypes.getByNameAndVersion(nodeTypeName)).toThrow(
'Unrecognized node type: unknownNode',
);
});

it('should return a regular node-type without version', () => {
const nodeType = mock<INodeType>();

// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.loadedNodes = {
[nodeTypeName]: { type: nodeType },
};

const result = nodeTypes.getByNameAndVersion(nodeTypeName);

expect(result).toEqual(nodeType);
});

it('should return a regular node-type with version', () => {
const nodeTypeV1 = mock<INodeType>();
const nodeType = mock<IVersionedNodeType>({
nodeVersions: { 1: nodeTypeV1 },
getNodeType: () => nodeTypeV1,
});

// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.loadedNodes = {
[nodeTypeName]: { type: nodeType },
};

const result = nodeTypes.getByNameAndVersion(nodeTypeName);

expect(result).toEqual(nodeTypeV1);
});

it('should throw when a node-type is requested as tool, but does not support being used as one', () => {
const nodeType = mock<INodeType>();

// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.loadedNodes = {
[nodeTypeName]: { type: nodeType },
};

expect(() => nodeTypes.getByNameAndVersion(`${nodeTypeName}Tool`)).toThrow(
'Node cannot be used as a tool',
);
});

it('should return the tool node-type when requested as tool', () => {
const nodeType = mock<INodeType>();
// @ts-expect-error can't use a mock here
nodeType.description = {
name: nodeTypeName,
displayName: 'TestNode',
usableAsTool: true,
properties: [],
};

// @ts-expect-error overwriting a readonly property
loadNodesAndCredentials.loadedNodes = {
[nodeTypeName]: { type: nodeType },
};

const result = nodeTypes.getByNameAndVersion(`${nodeTypeName}Tool`);
expect(result).not.toEqual(nodeType);
expect(result.description.name).toEqual('n8n-nodes-base.testNodeTool');
expect(result.description.displayName).toEqual('TestNode Tool');
expect(result.description.codex?.categories).toContain('AI');
expect(result.description.inputs).toEqual([]);
expect(result.description.outputs).toEqual(['ai_tool']);
});
});
});
37 changes: 30 additions & 7 deletions packages/cli/src/node-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,15 +44,38 @@ export class NodeTypes implements INodeTypes {
}

getByNameAndVersion(nodeType: string, version?: number): INodeType {
const versionedNodeType = NodeHelpers.getVersionedNodeType(
this.getNode(nodeType).type,
version,
);
if (versionedNodeType.description.usableAsTool) {
return NodeHelpers.convertNodeToAiTool(versionedNodeType);
const origType = nodeType;
const toolRequested = nodeType.startsWith('n8n-nodes-base') && nodeType.endsWith('Tool');
// Make sure the nodeType to actually get from disk is the un-wrapped type
if (toolRequested) {
nodeType = nodeType.replace(/Tool$/, '');
}

const node = this.getNode(nodeType);
const versionedNodeType = NodeHelpers.getVersionedNodeType(node.type, version);
if (!toolRequested) return versionedNodeType;

if (!versionedNodeType.description.usableAsTool)
throw new ApplicationError('Node cannot be used as a tool', { extra: { nodeType } });

const { loadedNodes } = this.loadNodesAndCredentials;
if (origType in loadedNodes) {
return loadedNodes[origType].type as INodeType;
}

return versionedNodeType;
// Instead of modifying the existing type, we extend it into a new type object
const clonedProperties = Object.create(
versionedNodeType.description.properties,
) as INodeTypeDescription['properties'];
const clonedDescription = Object.create(versionedNodeType.description, {
properties: { value: clonedProperties },
}) as INodeTypeDescription;
const clonedNode = Object.create(versionedNodeType, {
description: { value: clonedDescription },
}) as INodeType;
const tool = NodeHelpers.convertNodeToAiTool(clonedNode);
loadedNodes[nodeType + 'Tool'] = { sourcePath: '', type: tool };
return tool;
}

/* Some nodeTypes need to get special parameters applied like the polling nodes the polling times */
Expand Down

0 comments on commit c2fb881

Please sign in to comment.