From fa5695c25050ab4a7cb804d3e13112be12df4efb Mon Sep 17 00:00:00 2001
From: balibabu <cike8899@users.noreply.github.com>
Date: Thu, 27 Jun 2024 09:20:19 +0800
Subject: [PATCH] feat: add CategorizeHandle #918 (#1282)

### What problem does this PR solve?

feat: add CategorizeHandle #918

### Type of change


- [x] New Feature (non-breaking change which adds functionality)
---
 .../flow/canvas/node/categorize-handle.tsx    |  39 +++++++
 web/src/pages/flow/canvas/node/index.less     |   6 +
 web/src/pages/flow/canvas/node/index.tsx      |  37 +++---
 .../categorize-form/dynamic-categorize.tsx    | 108 +++++++++++-------
 web/src/pages/flow/categorize-form/hooks.ts   |  38 +++++-
 web/src/pages/flow/categorize-form/index.tsx  |   2 +-
 web/src/pages/flow/constant.tsx               |  15 +++
 web/src/pages/flow/store.ts                   |  18 +++
 8 files changed, 198 insertions(+), 65 deletions(-)
 create mode 100644 web/src/pages/flow/canvas/node/categorize-handle.tsx

diff --git a/web/src/pages/flow/canvas/node/categorize-handle.tsx b/web/src/pages/flow/canvas/node/categorize-handle.tsx
new file mode 100644
index 000000000..2d4329734
--- /dev/null
+++ b/web/src/pages/flow/canvas/node/categorize-handle.tsx
@@ -0,0 +1,39 @@
+import { Handle, Position } from 'reactflow';
+// import { v4 as uuid } from 'uuid';
+
+import styles from './index.less';
+
+const DEFAULT_HANDLE_STYLE = {
+  width: 6,
+  height: 6,
+  bottom: -5,
+  fontSize: 8,
+};
+
+interface IProps {
+  top: number;
+  right: number;
+  text: string;
+  idx: number;
+}
+
+const CategorizeHandle = ({ top, right, text, idx }: IProps) => {
+  return (
+    <Handle
+      type="source"
+      position={Position.Right}
+      id={`CategorizeHandle${idx}`}
+      isConnectable
+      style={{
+        ...DEFAULT_HANDLE_STYLE,
+        top: `${top}%`,
+        right: `${right}%`,
+        background: 'red',
+      }}
+    >
+      <span className={styles.categorizeAnchorPointText}>{text}</span>
+    </Handle>
+  );
+};
+
+export default CategorizeHandle;
diff --git a/web/src/pages/flow/canvas/node/index.less b/web/src/pages/flow/canvas/node/index.less
index bd42fd3ad..ffba5c07f 100644
--- a/web/src/pages/flow/canvas/node/index.less
+++ b/web/src/pages/flow/canvas/node/index.less
@@ -37,6 +37,12 @@
       -3px 0 6px -4px rgba(0, 0, 0, 0.12),
       -6px 0 16px 6px rgba(0, 0, 0, 0.05);
   }
+  .categorizeAnchorPointText {
+    position: absolute;
+    top: -4px;
+    left: 8px;
+    white-space: nowrap;
+  }
 }
 .selectedNode {
   border: 1px solid rgb(59, 118, 244);
diff --git a/web/src/pages/flow/canvas/node/index.tsx b/web/src/pages/flow/canvas/node/index.tsx
index 41d121166..5d0fc1bda 100644
--- a/web/src/pages/flow/canvas/node/index.tsx
+++ b/web/src/pages/flow/canvas/node/index.tsx
@@ -4,12 +4,14 @@ import { Handle, NodeProps, Position } from 'reactflow';
 import OperateDropdown from '@/components/operate-dropdown';
 import { CopyOutlined } from '@ant-design/icons';
 import { Flex, MenuProps, Space } from 'antd';
+import get from 'lodash/get';
 import { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
-import { Operator, operatorMap } from '../../constant';
+import { CategorizeAnchorPointPositions, Operator } from '../../constant';
 import { NodeData } from '../../interface';
 import OperatorIcon from '../../operator-icon';
 import useGraphStore from '../../store';
+import CategorizeHandle from './categorize-handle';
 import styles from './index.less';
 
 export function RagNode({
@@ -30,7 +32,8 @@ export function RagNode({
     duplicateNodeById(id);
   }, [id, duplicateNodeById]);
 
-  const description = operatorMap[data.label as Operator].description;
+  const isCategorize = data.label === Operator.Categorize;
+  const categoryData = get(data, 'form.category_description') ?? {};
 
   const items: MenuProps['items'] = [
     {
@@ -57,9 +60,7 @@ export function RagNode({
         position={Position.Left}
         isConnectable={isConnectable}
         className={styles.handle}
-      >
-        {/* <PlusCircleOutlined style={{ fontSize: 10 }} /> */}
-      </Handle>
+      ></Handle>
       <Handle type="source" position={Position.Top} id="d" isConnectable />
       <Handle
         type="source"
@@ -67,34 +68,32 @@ export function RagNode({
         isConnectable={isConnectable}
         className={styles.handle}
         id="b"
-      >
-        {/* <PlusCircleOutlined style={{ fontSize: 10 }} /> */}
-      </Handle>
+      ></Handle>
       <Handle type="source" position={Position.Bottom} id="a" isConnectable />
+      {isCategorize &&
+        Object.keys(categoryData).map((x, idx) => (
+          <CategorizeHandle
+            top={CategorizeAnchorPointPositions[idx].top}
+            right={CategorizeAnchorPointPositions[idx].right}
+            key={idx}
+            text={x}
+            idx={idx}
+          ></CategorizeHandle>
+        ))}
       <Flex vertical align="center" justify="center">
         <Space size={6}>
           <OperatorIcon
             name={data.label as Operator}
             fontSize={16}
           ></OperatorIcon>
-          {/* {data.label} */}
           <OperateDropdown
             iconFontSize={14}
             deleteItem={deleteNode}
             items={items}
           ></OperateDropdown>
         </Space>
-        {/* <div className={styles.nodeName}>{id}</div> */}
       </Flex>
-      {/* <div>
-        <Text
-          ellipsis={{ tooltip: description }}
-          style={{ width: 130 }}
-          className={styles.description}
-        >
-          {description}
-        </Text>
-      </div> */}
+
       <section className={styles.bottomBox}>
         <div className={styles.nodeName}>{id}</div>
       </section>
diff --git a/web/src/pages/flow/categorize-form/dynamic-categorize.tsx b/web/src/pages/flow/categorize-form/dynamic-categorize.tsx
index 3bba4f41f..d02a0bce8 100644
--- a/web/src/pages/flow/categorize-form/dynamic-categorize.tsx
+++ b/web/src/pages/flow/categorize-form/dynamic-categorize.tsx
@@ -1,58 +1,78 @@
 import { CloseOutlined } from '@ant-design/icons';
 import { Button, Card, Form, Input, Select, Typography } from 'antd';
-import { useBuildCategorizeToOptions } from './hooks';
+import { useBuildCategorizeToOptions, useHandleToSelectChange } from './hooks';
 
-const DynamicCategorize = () => {
+interface IProps {
+  nodeId?: string;
+}
+
+const DynamicCategorize = ({ nodeId }: IProps) => {
   const form = Form.useFormInstance();
   const options = useBuildCategorizeToOptions();
+  const { handleSelectChange } = useHandleToSelectChange(
+    options.map((x) => x.value),
+    nodeId,
+  );
 
   return (
     <>
       <Form.List name="items">
-        {(fields, { add, remove }) => (
-          <div style={{ display: 'flex', rowGap: 16, flexDirection: 'column' }}>
-            {fields.map((field) => (
-              <Card
-                size="small"
-                key={field.key}
-                extra={
-                  <CloseOutlined
-                    onClick={() => {
-                      remove(field.name);
-                    }}
-                  />
-                }
-              >
-                <Form.Item
-                  label="name"
-                  name={[field.name, 'name']}
-                  initialValue={`Categorize ${field.name + 1}`}
-                  rules={[
-                    { required: true, message: 'Please input your name!' },
-                  ]}
+        {(fields, { add, remove }) => {
+          const handleAdd = () => {
+            const idx = fields.length;
+            add({ name: `Categorize ${idx + 1}` });
+          };
+          return (
+            <div
+              style={{ display: 'flex', rowGap: 10, flexDirection: 'column' }}
+            >
+              {fields.map((field) => (
+                <Card
+                  size="small"
+                  key={field.key}
+                  extra={
+                    <CloseOutlined
+                      onClick={() => {
+                        remove(field.name);
+                      }}
+                    />
+                  }
                 >
-                  <Input />
-                </Form.Item>
-                <Form.Item
-                  label="description"
-                  name={[field.name, 'description']}
-                >
-                  <Input.TextArea rows={3} />
-                </Form.Item>
-                <Form.Item label="examples" name={[field.name, 'examples']}>
-                  <Input.TextArea rows={3} />
-                </Form.Item>
-                <Form.Item label="to" name={[field.name, 'to']}>
-                  <Select options={options} />
-                </Form.Item>
-              </Card>
-            ))}
+                  <Form.Item
+                    label="name"
+                    name={[field.name, 'name']}
+                    // initialValue={`Categorize ${field.name + 1}`}
+                    rules={[
+                      { required: true, message: 'Please input your name!' },
+                    ]}
+                  >
+                    <Input />
+                  </Form.Item>
+                  <Form.Item
+                    label="description"
+                    name={[field.name, 'description']}
+                  >
+                    <Input.TextArea rows={3} />
+                  </Form.Item>
+                  <Form.Item label="examples" name={[field.name, 'examples']}>
+                    <Input.TextArea rows={3} />
+                  </Form.Item>
+                  <Form.Item label="to" name={[field.name, 'to']}>
+                    <Select
+                      allowClear
+                      options={options}
+                      onChange={handleSelectChange}
+                    />
+                  </Form.Item>
+                </Card>
+              ))}
 
-            <Button type="dashed" onClick={() => add()} block>
-              + Add Item
-            </Button>
-          </div>
-        )}
+              <Button type="dashed" onClick={handleAdd} block>
+                + Add Item
+              </Button>
+            </div>
+          );
+        }}
       </Form.List>
 
       <Form.Item noStyle shouldUpdate>
diff --git a/web/src/pages/flow/categorize-form/hooks.ts b/web/src/pages/flow/categorize-form/hooks.ts
index 615ea8ce0..cac5a5236 100644
--- a/web/src/pages/flow/categorize-form/hooks.ts
+++ b/web/src/pages/flow/categorize-form/hooks.ts
@@ -1,6 +1,6 @@
 import get from 'lodash/get';
 import omit from 'lodash/omit';
-import { useCallback, useEffect } from 'react';
+import { useCallback, useEffect, useRef } from 'react';
 import { Operator } from '../constant';
 import {
   ICategorizeItem,
@@ -72,6 +72,7 @@ export const useHandleFormValuesChange = ({
 }: IOperatorForm) => {
   const handleValuesChange = useCallback(
     (changedValues: any, values: any) => {
+      console.info(changedValues, values);
       onValuesChange?.(changedValues, {
         ...omit(values, 'items'),
         category_description: buildCategorizeObjectFromList(values.items),
@@ -90,3 +91,38 @@ export const useHandleFormValuesChange = ({
 
   return { handleValuesChange };
 };
+
+export const useHandleToSelectChange = (
+  opstionIds: string[],
+  nodeId?: string,
+) => {
+  // const [previousTarget, setPreviousTarget] = useState('');
+  const previousTarget = useRef('');
+  const { addEdge, deleteEdgeBySourceAndTarget } = useGraphStore(
+    (state) => state,
+  );
+  const handleSelectChange = useCallback(
+    (value?: string) => {
+      if (nodeId) {
+        if (previousTarget.current) {
+          // delete previous edge
+          deleteEdgeBySourceAndTarget(nodeId, previousTarget.current);
+        }
+        if (value) {
+          addEdge({
+            source: nodeId,
+            target: value,
+            sourceHandle: 'b',
+            targetHandle: 'd',
+          });
+        } else {
+          // if the value is empty, delete the edges between the current node and all nodes in the drop-down box.
+        }
+        previousTarget.current = value;
+      }
+    },
+    [addEdge, nodeId, deleteEdgeBySourceAndTarget],
+  );
+
+  return { handleSelectChange };
+};
diff --git a/web/src/pages/flow/categorize-form/index.tsx b/web/src/pages/flow/categorize-form/index.tsx
index 1526e5a00..5f1192a69 100644
--- a/web/src/pages/flow/categorize-form/index.tsx
+++ b/web/src/pages/flow/categorize-form/index.tsx
@@ -32,7 +32,7 @@ const CategorizeForm = ({ form, onValuesChange, node }: IOperatorForm) => {
       >
         <LLMSelect></LLMSelect>
       </Form.Item>
-      <DynamicCategorize></DynamicCategorize>
+      <DynamicCategorize nodeId={node?.id}></DynamicCategorize>
     </Form>
   );
 };
diff --git a/web/src/pages/flow/constant.tsx b/web/src/pages/flow/constant.tsx
index 866eb4ec8..fc10b1345 100644
--- a/web/src/pages/flow/constant.tsx
+++ b/web/src/pages/flow/constant.tsx
@@ -82,3 +82,18 @@ export const initialFormValuesMap = {
   [Operator.Answer]: {},
   [Operator.Categorize]: {},
 };
+
+export const CategorizeAnchorPointPositions = [
+  { top: 1, right: 34 },
+  { top: 8, right: 18 },
+  { top: 15, right: 10 },
+  { top: 24, right: 4 },
+  { top: 31, right: 1 },
+  { top: 38, right: -2 },
+  { top: 62, right: -2 }, //bottom
+  { top: 71, right: 1 },
+  { top: 79, right: 6 },
+  { top: 86, right: 12 },
+  { top: 91, right: 20 },
+  { top: 98, right: 34 },
+];
diff --git a/web/src/pages/flow/store.ts b/web/src/pages/flow/store.ts
index dac10f2cc..459d60a0d 100644
--- a/web/src/pages/flow/store.ts
+++ b/web/src/pages/flow/store.ts
@@ -34,10 +34,12 @@ export type RFState = {
   onSelectionChange: OnSelectionChangeFunc;
   addNode: (nodes: Node) => void;
   getNode: (id: string) => Node | undefined;
+  addEdge: (connection: Connection) => void;
   duplicateNode: (id: string) => void;
   deleteEdge: () => void;
   deleteEdgeById: (id: string) => void;
   deleteNodeById: (id: string) => void;
+  deleteEdgeBySourceAndTarget: (source: string, target: string) => void;
   findNodeByName: (operatorName: Operator) => Node | undefined;
   findNodeById: (id: string) => Node | undefined;
 };
@@ -83,6 +85,14 @@ const useGraphStore = create<RFState>()(
       getNode: (id: string) => {
         return get().nodes.find((x) => x.id === id);
       },
+      addEdge: (connection: Connection) => {
+        set({
+          edges: addEdge(connection, get().edges),
+        });
+      },
+      // addOnlyOneEdgeBetweenTwoNodes: (connection: Connection) => {
+
+      // },
       duplicateNode: (id: string) => {
         const { getNode, addNode } = get();
         const node = getNode(id);
@@ -114,6 +124,14 @@ const useGraphStore = create<RFState>()(
           edges: edges.filter((edge) => edge.id !== id),
         });
       },
+      deleteEdgeBySourceAndTarget: (source: string, target: string) => {
+        const { edges } = get();
+        set({
+          edges: edges.filter(
+            (edge) => edge.target !== target && edge.source !== source,
+          ),
+        });
+      },
       deleteNodeById: (id: string) => {
         const { nodes, edges } = get();
         set({