From 0d7594074ca3cfb8d6dd08fbcb823203a557f0e5 Mon Sep 17 00:00:00 2001
From: Aralhi <xiaoping0501@gmail.com>
Date: Thu, 1 Feb 2024 00:08:05 +0800
Subject: [PATCH] feat: add related nodes

---
 web/app/i18n.ts                      |  2 ++
 web/components/flow/node-handler.tsx | 48 +++++++++++++++++++++++++---
 web/types/flow.ts                    |  1 +
 web/utils/index.ts                   |  2 ++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/web/app/i18n.ts b/web/app/i18n.ts
index a3b065d31..359857800 100644
--- a/web/app/i18n.ts
+++ b/web/app/i18n.ts
@@ -192,6 +192,7 @@ const en = {
   flow_name_required: 'Please enter the flow name',
   save_flow_success: 'Save flow success',
   delete_flow_confirm: 'Are you sure you want to delete this flow?',
+  related_nodes: 'Related Nodes',
 } as const;
 
 export type I18nKeys = keyof typeof en;
@@ -390,6 +391,7 @@ const zh: Resources['translation'] = {
   flow_name_required: '请输入工作流名称',
   save_flow_success: '保存工作流成功',
   delete_flow_confirm: '确定删除该工作流吗？',
+  related_nodes: '关联节点',
 } as const;
 
 i18n.use(initReactI18next).init({
diff --git a/web/components/flow/node-handler.tsx b/web/components/flow/node-handler.tsx
index 2dc6e1397..412a40011 100644
--- a/web/components/flow/node-handler.tsx
+++ b/web/components/flow/node-handler.tsx
@@ -1,11 +1,13 @@
-import { Tooltip, Typography, message } from 'antd';
+import { Popconfirm, Tooltip, Typography, message } from 'antd';
 import React from 'react';
 import { Connection, Handle, Position, useReactFlow } from 'reactflow';
 import RequiredIcon from './required-icon';
-import { InfoCircleOutlined } from '@ant-design/icons';
+import { InfoCircleOutlined, PlusOutlined } from '@ant-design/icons';
 import { IFlowNode, IFlowNodeInput, IFlowNodeOutput, IFlowNodeParameter } from '@/types/flow';
 import { useTranslation } from 'react-i18next';
 import classNames from 'classnames';
+import { FLOW_NODES_KEY } from '@/utils';
+import StaticNodes from './static-nodes';
 
 interface NodeHandlerProps {
   node: IFlowNode;
@@ -19,6 +21,7 @@ interface NodeHandlerProps {
 const NodeHandler: React.FC<NodeHandlerProps> = ({ node, data, type, label, index }) => {
   const { t } = useTranslation();
   const reactflow = useReactFlow();
+  const [relatedNodes, setRelatedNodes] = React.useState<IFlowNode[]>([]);
 
   function isValidConnection(connection: Connection) {
     const { sourceHandle, targetHandle, source, target } = connection;
@@ -44,6 +47,27 @@ const NodeHandler: React.FC<NodeHandlerProps> = ({ node, data, type, label, inde
     return false;
   }
 
+  function showRelatedNodes() {
+    // find all nodes that can be connected to this node
+    const cache = localStorage.getItem(FLOW_NODES_KEY);
+    if (!cache) {
+      return;
+    }
+    const staticNodes = JSON.parse(cache);
+    const typeCls = data.type_cls;
+    let nodes: IFlowNode[] = [];
+    if (label === 'inputs') {
+      // find other operators and outputs matching this input type_cls
+      nodes = staticNodes
+        .filter((node: IFlowNode) => node.flow_type === 'operator')
+        .filter((node: IFlowNode) => node.outputs?.some((output: IFlowNodeOutput) => output.type_cls === typeCls));
+    } else if (label === 'parameters') {
+      // fint other resources and parent_cls including this parameter type_cls
+      nodes = staticNodes.filter((node: IFlowNode) => node.flow_type === 'resource').filter((node: IFlowNode) => node.parent_cls?.includes(typeCls));
+    }
+    setRelatedNodes(nodes);
+  }
+
   return (
     <div
       className={classNames('relative flex items-center', {
@@ -63,10 +87,26 @@ const NodeHandler: React.FC<NodeHandlerProps> = ({ node, data, type, label, inde
           'pr-4': label === 'outputs',
         })}
       >
-        {data.type_name}:<RequiredIcon optional={data.optional} />
+        <Popconfirm
+          placement="left"
+          icon={null}
+          showCancel={false}
+          okButtonProps={{ className: 'hidden' }}
+          title={t('related_nodes')}
+          description={
+            <div className="w-60">
+              <StaticNodes nodes={relatedNodes} />
+            </div>
+          }
+        >
+          {node.flow_type === 'operator' && ['inputs', 'parameters'].includes(label) && (
+            <PlusOutlined className="mr-2 cursor-pointer" onClick={showRelatedNodes} />
+          )}
+        </Popconfirm>
+        {data.type_name}:{label !== 'outputs' && <RequiredIcon optional={data.optional} />}
         {data.description && (
           <Tooltip title={data.description}>
-            <InfoCircleOutlined className="ml-2" />
+            <InfoCircleOutlined className="ml-2 cursor-pointer" />
           </Tooltip>
         )}
       </Typography>
diff --git a/web/types/flow.ts b/web/types/flow.ts
index a40345850..c3259fc86 100644
--- a/web/types/flow.ts
+++ b/web/types/flow.ts
@@ -59,6 +59,7 @@ export type IFlowNodeOutput = {
 export type IFlowNode = Node & {
   type_name: string;
   type_cls: string;
+  parent_cls?: string; // resource have this key
   type: string;
   label: string;
   name: string;
diff --git a/web/utils/index.ts b/web/utils/index.ts
index fb6556257..d90fad667 100644
--- a/web/utils/index.ts
+++ b/web/utils/index.ts
@@ -4,6 +4,8 @@ export const STORAGE_THEME_KEY = '__db_gpt_theme_key';
 export const STORAGE_LANG_KEY = '__db_gpt_lng_key';
 /** Init Message */
 export const STORAGE_INIT_MESSAGE_KET = '__db_gpt_im_key';
+/** Flow nodes */
+export const FLOW_NODES_KEY = '__db_gpt_static_flow_nodes_key';
 
 export * from './storage';
 export * from './constants';