From 4018dc1f7f9b6e69f3ec40f3422c49e1931be01b Mon Sep 17 00:00:00 2001
From: Oleg Ivaniv <me@olegivaniv.com>
Date: Fri, 20 Dec 2024 09:28:03 +0100
Subject: [PATCH] Connect vector store tools nodes when adding from connection
 drag

---
 .../shared/createVectorStoreNode.ts           |  2 ++
 packages/editor-ui/src/Interface.ts           |  1 +
 .../Node/NodeCreator/Modes/NodesMode.vue      | 33 ++++++++++++++++---
 .../composables/useActionsGeneration.ts       |  3 +-
 .../NodeCreator/composables/useViewStacks.ts  | 11 ++++++-
 packages/workflow/src/Interfaces.ts           |  1 +
 6 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts
index 3135a74a08..f35f1bd254 100644
--- a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts
+++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts
@@ -110,12 +110,14 @@ function getOperationModeOptions(args: VectorStoreNodeConstructorArgs): INodePro
 			value: 'retrieve',
 			description: 'Retrieve documents from vector store to be used as vector store with AI nodes',
 			action: 'Retrieve documents for AI processing as Vector Store',
+			outputConnectionType: NodeConnectionType.AiVectorStore,
 		},
 		{
 			name: 'Retrieve Documents (As Tool for AI Agent)',
 			value: 'retrieve-as-tool',
 			description: 'Retrieve documents from vector store to be used as tool with AI nodes',
 			action: 'Retrieve documents for AI processing as Tool',
+			outputConnectionType: NodeConnectionType.AiTool,
 		},
 		{
 			name: 'Update Documents',
diff --git a/packages/editor-ui/src/Interface.ts b/packages/editor-ui/src/Interface.ts
index a2a1e8e065..03327a1d22 100644
--- a/packages/editor-ui/src/Interface.ts
+++ b/packages/editor-ui/src/Interface.ts
@@ -723,6 +723,7 @@ export interface ActionTypeDescription extends SimplifiedNodeType {
 	displayOptions?: IDisplayOptions;
 	values?: IDataObject;
 	actionKey: string;
+	outputConnectionType?: NodeConnectionType;
 	codex: {
 		label: string;
 		categories: string[];
diff --git a/packages/editor-ui/src/components/Node/NodeCreator/Modes/NodesMode.vue b/packages/editor-ui/src/components/Node/NodeCreator/Modes/NodesMode.vue
index b577669169..85a022805a 100644
--- a/packages/editor-ui/src/components/Node/NodeCreator/Modes/NodesMode.vue
+++ b/packages/editor-ui/src/components/Node/NodeCreator/Modes/NodesMode.vue
@@ -1,7 +1,7 @@
 <script setup lang="ts">
 import { camelCase } from 'lodash-es';
 import { computed } from 'vue';
-import type { INodeCreateElement, NodeFilterType } from '@/Interface';
+import type { INodeCreateElement, NodeCreateElement, NodeFilterType } from '@/Interface';
 import {
 	TRIGGER_NODE_CREATOR_VIEW,
 	HTTP_REQUEST_NODE_TYPE,
@@ -25,6 +25,8 @@ import NoResults from '../Panel/NoResults.vue';
 import { useI18n } from '@/composables/useI18n';
 import { getNodeIcon, getNodeIconColor, getNodeIconUrl } from '@/utils/nodeTypesUtils';
 import { useUIStore } from '@/stores/ui.store';
+import { useActions } from '../composables/useActions';
+import { INodeParameters } from 'n8n-workflow';
 
 export interface Props {
 	rootView: 'trigger' | 'action';
@@ -40,12 +42,21 @@ const rootStore = useRootStore();
 
 const { mergedNodes, actions, onSubcategorySelected } = useNodeCreatorStore();
 const { pushViewStack, popViewStack } = useViewStacks();
+const { setAddedNodeActionParameters } = useActions();
 
 const { registerKeyHook } = useKeyboardNavigation();
 
 const activeViewStack = computed(() => useViewStacks().activeViewStack);
 const globalSearchItemsDiff = computed(() => useViewStacks().globalSearchItemsDiff);
 
+function getFilteredActions(node: NodeCreateElement) {
+	const nodeActions = actions?.[node.key] || [];
+	if (activeViewStack.value.actionsFilter) {
+		return activeViewStack.value.actionsFilter(nodeActions);
+	}
+	return nodeActions;
+}
+
 function selectNodeType(nodeTypes: string[]) {
 	emit('nodeTypeSelected', nodeTypes);
 }
@@ -87,9 +98,21 @@ function onSelected(item: INodeCreateElement) {
 	}
 
 	if (item.type === 'node') {
-		const nodeActions = actions?.[item.key] || [];
+		const nodeActions = getFilteredActions(item);
+
+		// If there is only one action, use it
+		if (nodeActions.length === 1) {
+			selectNodeType([item.key]);
+			setAddedNodeActionParameters({
+				name: nodeActions[0].defaults.name ?? item.properties.displayName,
+				key: item.key,
+				value: nodeActions[0].values as INodeParameters,
+			});
+			return;
+		}
+
 		// Only show actions if there are more than one or if the view is not an AI subcategory
-		if (nodeActions.length <= 1 || activeViewStack.value.hideActions) {
+		if (nodeActions.length === 0 || activeViewStack.value.hideActions) {
 			selectNodeType([item.key]);
 			return;
 		}
@@ -158,7 +181,7 @@ function subcategoriesMapper(item: INodeCreateElement) {
 	if (item.type !== 'node') return item;
 
 	const hasTriggerGroup = item.properties.group.includes('trigger');
-	const nodeActions = actions?.[item.key] || [];
+	const nodeActions = getFilteredActions(item);
 	const hasActions = nodeActions.length > 0;
 
 	if (hasTriggerGroup && hasActions) {
@@ -179,7 +202,7 @@ function baseSubcategoriesFilter(item: INodeCreateElement): boolean {
 	if (item.type !== 'node') return false;
 
 	const hasTriggerGroup = item.properties.group.includes('trigger');
-	const nodeActions = actions?.[item.key] || [];
+	const nodeActions = getFilteredActions(item);
 	const hasActions = nodeActions.length > 0;
 
 	const isTriggerRootView = activeViewStack.value.rootView === TRIGGER_NODE_CREATOR_VIEW;
diff --git a/packages/editor-ui/src/components/Node/NodeCreator/composables/useActionsGeneration.ts b/packages/editor-ui/src/components/Node/NodeCreator/composables/useActionsGeneration.ts
index e428269dbe..8143489943 100644
--- a/packages/editor-ui/src/components/Node/NodeCreator/composables/useActionsGeneration.ts
+++ b/packages/editor-ui/src/components/Node/NodeCreator/composables/useActionsGeneration.ts
@@ -87,6 +87,7 @@ function operationsCategory(nodeTypeDescription: INodeTypeDescription): ActionTy
 		displayName: item.action ?? startCase(item.name),
 		description: item.description ?? '',
 		displayOptions: matchedProperty.displayOptions,
+		outputConnectionType: item.outputConnectionType,
 		values: {
 			[matchedProperty.name]: matchedProperty.type === 'multiOptions' ? [item.value] : item.value,
 		},
@@ -117,6 +118,7 @@ function modeCategory(nodeTypeDescription: INodeTypeDescription): ActionTypeDesc
 		displayName: item.action ?? startCase(item.name),
 		description: item.description ?? '',
 		displayOptions: matchedProperty.displayOptions,
+		outputConnectionType: item.outputConnectionType,
 		values: {
 			[matchedProperty.name]: item.value,
 		},
@@ -320,7 +322,6 @@ export function useActionsGenerator() {
 		const visibleNodeTypes = [...nodeTypes];
 		const actions: ActionsRecord<typeof mergedNodes> = {};
 		const mergedNodes: SimplifiedNodeType[] = [];
-
 		visibleNodeTypes
 			.filter((node) => !node.group.includes('trigger'))
 			.forEach((app) => {
diff --git a/packages/editor-ui/src/components/Node/NodeCreator/composables/useViewStacks.ts b/packages/editor-ui/src/components/Node/NodeCreator/composables/useViewStacks.ts
index 3584a38723..77c3330459 100644
--- a/packages/editor-ui/src/components/Node/NodeCreator/composables/useViewStacks.ts
+++ b/packages/editor-ui/src/components/Node/NodeCreator/composables/useViewStacks.ts
@@ -1,4 +1,5 @@
 import type {
+	ActionTypeDescription,
 	INodeCreateElement,
 	NodeCreateElement,
 	NodeFilterType,
@@ -38,8 +39,8 @@ import { useKeyboardNavigation } from './useKeyboardNavigation';
 import { useNodeTypesStore } from '@/stores/nodeTypes.store';
 import {
 	AI_TRANSFORM_NODE_TYPE,
+	NodeConnectionType,
 	type INodeInputFilter,
-	type NodeConnectionType,
 	type Themed,
 } from 'n8n-workflow';
 import { useCanvasStore } from '@/stores/canvas.store';
@@ -71,6 +72,7 @@ interface ViewStack {
 	hideActions?: boolean;
 	baseFilter?: (item: INodeCreateElement) => boolean;
 	itemsMapper?: (item: INodeCreateElement) => INodeCreateElement;
+	actionsFilter?: (items: ActionTypeDescription[]) => ActionTypeDescription[];
 	panelClass?: string;
 	sections?: string[] | NodeViewItemSection[];
 }
@@ -346,6 +348,13 @@ export const useViewStacks = defineStore('nodeCreatorViewStacks', () => {
 					subcategory: connectionType,
 				};
 			},
+			actionsFilter: (items: ActionTypeDescription[]) => {
+				// Filter out actions that are not compatible with the connection type
+				if (items.some((item) => item.outputConnectionType)) {
+					return items.filter((item) => item.outputConnectionType === connectionType);
+				}
+				return items;
+			},
 			hideActions: true,
 			preventBack: true,
 		});
diff --git a/packages/workflow/src/Interfaces.ts b/packages/workflow/src/Interfaces.ts
index c5fa0a938a..067711f035 100644
--- a/packages/workflow/src/Interfaces.ts
+++ b/packages/workflow/src/Interfaces.ts
@@ -1449,6 +1449,7 @@ export interface INodePropertyOptions {
 	action?: string;
 	description?: string;
 	routing?: INodePropertyRouting;
+	outputConnectionType?: NodeConnectionType;
 }
 
 export interface INodeListSearchItems extends INodePropertyOptions {