From fef663a59d60071a1b01269cd5a3832d6f2aa432 Mon Sep 17 00:00:00 2001
From: balibabu <cike8899@users.noreply.github.com>
Date: Tue, 25 Jun 2024 19:28:24 +0800
Subject: [PATCH] feat: build categorize list from object #918 (#1276)

### What problem does this PR solve?

feat: build categorize list from object #918

### Type of change


- [x] New Feature (non-breaking change which adds functionality)
---
 web/src/components/llm-select/index.tsx       | 22 +++++-
 .../components/llm-setting-items/index.tsx    | 45 +++++++++--
 .../model-setting.tsx                         | 16 +---
 web/src/pages/flow/categorize-form/hooks.ts   | 79 +++++++++++++++++++
 web/src/pages/flow/categorize-form/index.tsx  | 19 ++++-
 web/src/pages/flow/flow-drawer/index.tsx      |  1 +
 web/src/pages/flow/generate-form/index.tsx    | 39 +--------
 web/src/pages/flow/hooks.ts                   | 29 +++++++
 web/src/pages/flow/interface.ts               | 18 ++++-
 web/src/pages/flow/store.ts                   |  4 +
 web/src/pages/flow/utils.ts                   |  5 +-
 11 files changed, 208 insertions(+), 69 deletions(-)

diff --git a/web/src/components/llm-select/index.tsx b/web/src/components/llm-select/index.tsx
index 4d730c80c..61d37b599 100644
--- a/web/src/components/llm-select/index.tsx
+++ b/web/src/components/llm-select/index.tsx
@@ -1,16 +1,30 @@
 import { Popover, Select } from 'antd';
 import LlmSettingItems from '../llm-setting-items';
 
-const LLMSelect = () => {
+interface IProps {
+  id?: string;
+  value?: string;
+  onChange?: (value: string) => void;
+}
+
+const LLMSelect = ({ id, value, onChange }: IProps) => {
   const content = (
-    <div>
-      <LlmSettingItems handleParametersChange={() => {}}></LlmSettingItems>
+    <div style={{ width: 400 }}>
+      <LlmSettingItems
+        formItemLayout={{ labelCol: { span: 10 }, wrapperCol: { span: 14 } }}
+      ></LlmSettingItems>
     </div>
   );
 
   return (
     <Popover content={content} trigger="click" placement="left" arrow={false}>
-      <Select style={{ width: '100%' }} dropdownStyle={{ display: 'none' }} />
+      <Select
+        style={{ width: '100%' }}
+        dropdownStyle={{ display: 'none' }}
+        id={id}
+        value={value}
+        onChange={onChange}
+      />
     </Popover>
   );
 };
diff --git a/web/src/components/llm-setting-items/index.tsx b/web/src/components/llm-setting-items/index.tsx
index 398d1c4ca..8b88b8703 100644
--- a/web/src/components/llm-setting-items/index.tsx
+++ b/web/src/components/llm-setting-items/index.tsx
@@ -1,24 +1,38 @@
-import { LlmModelType, ModelVariableType } from '@/constants/knowledge';
+import {
+  LlmModelType,
+  ModelVariableType,
+  settledModelVariableMap,
+} from '@/constants/knowledge';
 import { Divider, Flex, Form, InputNumber, Select, Slider, Switch } from 'antd';
 import camelCase from 'lodash/camelCase';
 
 import { useTranslate } from '@/hooks/commonHooks';
 import { useSelectLlmOptionsByModelType } from '@/hooks/llmHooks';
-import { useMemo } from 'react';
+import { useCallback, useMemo } from 'react';
 import styles from './index.less';
 
 interface IProps {
   prefix?: string;
-  handleParametersChange(value: ModelVariableType): void;
+  formItemLayout?: any;
+  handleParametersChange?(value: ModelVariableType): void;
 }
 
-const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
+const LlmSettingItems = ({ prefix, formItemLayout = {} }: IProps) => {
+  const form = Form.useFormInstance();
   const { t } = useTranslate('chat');
   const parameterOptions = Object.values(ModelVariableType).map((x) => ({
     label: t(camelCase(x)),
     value: x,
   }));
 
+  const handleParametersChange = useCallback(
+    (value: ModelVariableType) => {
+      const variable = settledModelVariableMap[value];
+      form?.setFieldsValue(variable);
+    },
+    [form],
+  );
+
   const memorizedPrefix = useMemo(() => (prefix ? [prefix] : []), [prefix]);
 
   const modelOptions = useSelectLlmOptionsByModelType();
@@ -29,6 +43,7 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
         label={t('model')}
         name="llm_id"
         tooltip={t('modelTip')}
+        {...formItemLayout}
         rules={[{ required: true, message: t('modelMessage') }]}
       >
         <Select options={modelOptions[LlmModelType.Chat]} showSearch />
@@ -38,6 +53,7 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
         label={t('freedom')}
         name="parameters"
         tooltip={t('freedomTip')}
+        {...formItemLayout}
         initialValue={ModelVariableType.Precise}
       >
         <Select<ModelVariableType>
@@ -45,7 +61,11 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
           onChange={handleParametersChange}
         />
       </Form.Item>
-      <Form.Item label={t('temperature')} tooltip={t('temperatureTip')}>
+      <Form.Item
+        label={t('temperature')}
+        tooltip={t('temperatureTip')}
+        {...formItemLayout}
+      >
         <Flex gap={20} align="center">
           <Form.Item
             name={'temperatureEnabled'}
@@ -87,7 +107,7 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
           </Form.Item>
         </Flex>
       </Form.Item>
-      <Form.Item label={t('topP')} tooltip={t('topPTip')}>
+      <Form.Item label={t('topP')} tooltip={t('topPTip')} {...formItemLayout}>
         <Flex gap={20} align="center">
           <Form.Item name={'topPEnabled'} valuePropName="checked" noStyle>
             <Switch size="small" />
@@ -122,7 +142,11 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
           </Form.Item>
         </Flex>
       </Form.Item>
-      <Form.Item label={t('presencePenalty')} tooltip={t('presencePenaltyTip')}>
+      <Form.Item
+        label={t('presencePenalty')}
+        tooltip={t('presencePenaltyTip')}
+        {...formItemLayout}
+      >
         <Flex gap={20} align="center">
           <Form.Item
             name={'presencePenaltyEnabled'}
@@ -170,6 +194,7 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
       <Form.Item
         label={t('frequencyPenalty')}
         tooltip={t('frequencyPenaltyTip')}
+        {...formItemLayout}
       >
         <Flex gap={20} align="center">
           <Form.Item
@@ -215,7 +240,11 @@ const LlmSettingItems = ({ prefix, handleParametersChange }: IProps) => {
           </Form.Item>
         </Flex>
       </Form.Item>
-      <Form.Item label={t('maxTokens')} tooltip={t('maxTokensTip')}>
+      <Form.Item
+        label={t('maxTokens')}
+        tooltip={t('maxTokensTip')}
+        {...formItemLayout}
+      >
         <Flex gap={20} align="center">
           <Form.Item name={'maxTokensEnabled'} valuePropName="checked" noStyle>
             <Switch size="small" />
diff --git a/web/src/pages/chat/chat-configuration-modal/model-setting.tsx b/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
index df457e55f..aac34cd2a 100644
--- a/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
+++ b/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
@@ -1,7 +1,3 @@
-import {
-  ModelVariableType,
-  settledModelVariableMap,
-} from '@/constants/knowledge';
 import classNames from 'classnames';
 import { useEffect } from 'react';
 import { ISegmentedContentProps } from '../interface';
@@ -20,11 +16,6 @@ const ModelSetting = ({
   initialLlmSetting?: Variable;
   visible?: boolean;
 }) => {
-  const handleParametersChange = (value: ModelVariableType) => {
-    const variable = settledModelVariableMap[value];
-    form.setFieldsValue({ llm_setting: variable });
-  };
-
   useEffect(() => {
     if (visible) {
       const values = Object.keys(variableEnabledFieldMap).reduce<
@@ -50,12 +41,7 @@ const ModelSetting = ({
         [styles.segmentedHidden]: !show,
       })}
     >
-      {visible && (
-        <LlmSettingItems
-          prefix="llm_setting"
-          handleParametersChange={handleParametersChange}
-        ></LlmSettingItems>
-      )}
+      {visible && <LlmSettingItems prefix="llm_setting"></LlmSettingItems>}
       {/* <Form.Item
         label={t('model')}
         name="llm_id"
diff --git a/web/src/pages/flow/categorize-form/hooks.ts b/web/src/pages/flow/categorize-form/hooks.ts
index d6ce5b3f2..615ea8ce0 100644
--- a/web/src/pages/flow/categorize-form/hooks.ts
+++ b/web/src/pages/flow/categorize-form/hooks.ts
@@ -1,4 +1,12 @@
+import get from 'lodash/get';
+import omit from 'lodash/omit';
+import { useCallback, useEffect } from 'react';
 import { Operator } from '../constant';
+import {
+  ICategorizeItem,
+  ICategorizeItemResult,
+  IOperatorForm,
+} from '../interface';
 import useGraphStore from '../store';
 
 // exclude some nodes downstream of the classification node
@@ -11,3 +19,74 @@ export const useBuildCategorizeToOptions = () => {
     .filter((x) => excludedNodes.every((y) => y !== x.data.label))
     .map((x) => ({ label: x.id, value: x.id }));
 };
+
+/**
+   * convert the following object into a list
+   * 
+   * {
+      "product_related": {
+      "description": "The question is about product usage, appearance and how it works.",
+      "examples": "Why it always beaming?\nHow to install it onto the wall?\nIt leaks, what to do?",
+      "to": "generate:0"
+      }
+      }
+*/
+const buildCategorizeListFromObject = (
+  categorizeItem: ICategorizeItemResult,
+) => {
+  return Object.keys(categorizeItem).reduce<Array<ICategorizeItem>>(
+    (pre, cur) => {
+      pre.push({ name: cur, ...categorizeItem[cur] });
+      return pre;
+    },
+    [],
+  );
+};
+
+/**
+   * Convert the list in the following form into an object
+   * {
+    "items": [
+      {
+        "name": "Categorize 1",
+        "description": "111",
+        "examples": "ddd",
+        "to": "Retrieval:LazyEelsStick"
+      }
+     ]
+    }
+*/
+const buildCategorizeObjectFromList = (list: Array<ICategorizeItem>) => {
+  return list.reduce<ICategorizeItemResult>((pre, cur) => {
+    if (cur?.name) {
+      pre[cur.name] = omit(cur, 'name');
+    }
+    return pre;
+  }, {});
+};
+
+export const useHandleFormValuesChange = ({
+  onValuesChange,
+  form,
+  node,
+}: IOperatorForm) => {
+  const handleValuesChange = useCallback(
+    (changedValues: any, values: any) => {
+      onValuesChange?.(changedValues, {
+        ...omit(values, 'items'),
+        category_description: buildCategorizeObjectFromList(values.items),
+      });
+    },
+    [onValuesChange],
+  );
+
+  useEffect(() => {
+    form?.setFieldsValue({
+      items: buildCategorizeListFromObject(
+        get(node, 'data.form.category_description', {}),
+      ),
+    });
+  }, [form, node]);
+
+  return { handleValuesChange };
+};
diff --git a/web/src/pages/flow/categorize-form/index.tsx b/web/src/pages/flow/categorize-form/index.tsx
index 6cdabe3e0..1526e5a00 100644
--- a/web/src/pages/flow/categorize-form/index.tsx
+++ b/web/src/pages/flow/categorize-form/index.tsx
@@ -1,11 +1,19 @@
 import LLMSelect from '@/components/llm-select';
 import { useTranslate } from '@/hooks/commonHooks';
 import { Form } from 'antd';
+import { useSetLlmSetting } from '../hooks';
 import { IOperatorForm } from '../interface';
 import DynamicCategorize from './dynamic-categorize';
+import { useHandleFormValuesChange } from './hooks';
 
-const CategorizeForm = ({ form, onValuesChange }: IOperatorForm) => {
+const CategorizeForm = ({ form, onValuesChange, node }: IOperatorForm) => {
   const { t } = useTranslate('flow');
+  const { handleValuesChange } = useHandleFormValuesChange({
+    form,
+    node,
+    onValuesChange,
+  });
+  useSetLlmSetting(form);
 
   return (
     <Form
@@ -14,11 +22,14 @@ const CategorizeForm = ({ form, onValuesChange }: IOperatorForm) => {
       wrapperCol={{ span: 15 }}
       autoComplete="off"
       form={form}
-      onValuesChange={onValuesChange}
+      onValuesChange={handleValuesChange}
       initialValues={{ items: [{}] }}
-      // layout={'vertical'}
     >
-      <Form.Item name={['cite']} label={t('cite')} tooltip={t('citeTip')}>
+      <Form.Item
+        name={'llm_id'}
+        label={t('model', { keyPrefix: 'chat' })}
+        tooltip={t('modelTip', { keyPrefix: 'chat' })}
+      >
         <LLMSelect></LLMSelect>
       </Form.Item>
       <DynamicCategorize></DynamicCategorize>
diff --git a/web/src/pages/flow/flow-drawer/index.tsx b/web/src/pages/flow/flow-drawer/index.tsx
index 310e27464..eec9beab8 100644
--- a/web/src/pages/flow/flow-drawer/index.tsx
+++ b/web/src/pages/flow/flow-drawer/index.tsx
@@ -53,6 +53,7 @@ const FlowDrawer = ({
         <OperatorForm
           onValuesChange={handleValuesChange}
           form={form}
+          node={node}
         ></OperatorForm>
       )}
     </Drawer>
diff --git a/web/src/pages/flow/generate-form/index.tsx b/web/src/pages/flow/generate-form/index.tsx
index ae4648a41..c7c0d88b9 100644
--- a/web/src/pages/flow/generate-form/index.tsx
+++ b/web/src/pages/flow/generate-form/index.tsx
@@ -1,44 +1,13 @@
 import LlmSettingItems from '@/components/llm-setting-items';
-import { variableEnabledFieldMap } from '@/constants/chat';
-import {
-  ModelVariableType,
-  settledModelVariableMap,
-} from '@/constants/knowledge';
 import { useTranslate } from '@/hooks/commonHooks';
-import { Variable } from '@/interfaces/database/chat';
 import { Form, Input, Switch } from 'antd';
-import { useCallback, useEffect } from 'react';
+import { useSetLlmSetting } from '../hooks';
 import { IOperatorForm } from '../interface';
 
 const GenerateForm = ({ onValuesChange, form }: IOperatorForm) => {
   const { t } = useTranslate('flow');
-  const initialLlmSetting = undefined;
 
-  const handleParametersChange = useCallback(
-    (value: ModelVariableType) => {
-      const variable = settledModelVariableMap[value];
-      form?.setFieldsValue(variable);
-    },
-    [form],
-  );
-
-  useEffect(() => {
-    const switchBoxValues = Object.keys(variableEnabledFieldMap).reduce<
-      Record<string, boolean>
-    >((pre, field) => {
-      pre[field] =
-        initialLlmSetting === undefined
-          ? true
-          : !!initialLlmSetting[
-              variableEnabledFieldMap[
-                field as keyof typeof variableEnabledFieldMap
-              ] as keyof Variable
-            ];
-      return pre;
-    }, {});
-    const otherValues = settledModelVariableMap[ModelVariableType.Precise];
-    form?.setFieldsValue({ ...switchBoxValues, ...otherValues });
-  }, [form, initialLlmSetting]);
+  useSetLlmSetting(form);
 
   return (
     <Form
@@ -49,9 +18,7 @@ const GenerateForm = ({ onValuesChange, form }: IOperatorForm) => {
       form={form}
       onValuesChange={onValuesChange}
     >
-      <LlmSettingItems
-        handleParametersChange={handleParametersChange}
-      ></LlmSettingItems>
+      <LlmSettingItems></LlmSettingItems>
       <Form.Item
         name={['prompt']}
         label={t('prompt', { keyPrefix: 'knowledgeConfiguration' })}
diff --git a/web/src/pages/flow/hooks.ts b/web/src/pages/flow/hooks.ts
index 68d29dbbe..bddac12ef 100644
--- a/web/src/pages/flow/hooks.ts
+++ b/web/src/pages/flow/hooks.ts
@@ -15,7 +15,14 @@ import React, {
 } from 'react';
 import { Node, Position, ReactFlowInstance } from 'reactflow';
 // import { shallow } from 'zustand/shallow';
+import { variableEnabledFieldMap } from '@/constants/chat';
+import {
+  ModelVariableType,
+  settledModelVariableMap,
+} from '@/constants/knowledge';
+import { Variable } from '@/interfaces/database/chat';
 import { useDebounceEffect } from 'ahooks';
+import { FormInstance } from 'antd';
 import { humanId } from 'human-id';
 import { useParams } from 'umi';
 import { Operator } from './constant';
@@ -218,3 +225,25 @@ export const useFetchDataOnMount = () => {
 export const useFlowIsFetching = () => {
   return useIsFetching({ queryKey: ['flowDetail'] }) > 0;
 };
+
+export const useSetLlmSetting = (form?: FormInstance) => {
+  const initialLlmSetting = undefined;
+
+  useEffect(() => {
+    const switchBoxValues = Object.keys(variableEnabledFieldMap).reduce<
+      Record<string, boolean>
+    >((pre, field) => {
+      pre[field] =
+        initialLlmSetting === undefined
+          ? true
+          : !!initialLlmSetting[
+              variableEnabledFieldMap[
+                field as keyof typeof variableEnabledFieldMap
+              ] as keyof Variable
+            ];
+      return pre;
+    }, {});
+    const otherValues = settledModelVariableMap[ModelVariableType.Precise];
+    form?.setFieldsValue({ ...switchBoxValues, ...otherValues });
+  }, [form, initialLlmSetting]);
+};
diff --git a/web/src/pages/flow/interface.ts b/web/src/pages/flow/interface.ts
index 0d2b8096b..73eca14d1 100644
--- a/web/src/pages/flow/interface.ts
+++ b/web/src/pages/flow/interface.ts
@@ -1,4 +1,5 @@
 import { FormInstance } from 'antd';
+import { Node } from 'reactflow';
 
 export interface DSLComponentList {
   id: string;
@@ -8,6 +9,7 @@ export interface DSLComponentList {
 export interface IOperatorForm {
   onValuesChange?(changedValues: any, values: any): void;
   form?: FormInstance;
+  node?: Node<NodeData>;
 }
 
 export interface IBeginForm {
@@ -35,9 +37,23 @@ export interface IGenerateForm {
   llm_id: string;
   parameters: { key: string; component_id: string };
 }
+export interface ICategorizeItem {
+  name: string;
+  description?: string;
+  examples?: string;
+  to?: string;
+}
+
+export type ICategorizeItemResult = Record<
+  string,
+  Omit<ICategorizeItem, 'name'>
+>;
+export interface ICategorizeForm extends IGenerateForm {
+  category_description: ICategorizeItemResult;
+}
 
 export type NodeData = {
   label: string;
   color: string;
-  form: IBeginForm | IRetrievalForm | IGenerateForm;
+  form: IBeginForm | IRetrievalForm | IGenerateForm | ICategorizeForm;
 };
diff --git a/web/src/pages/flow/store.ts b/web/src/pages/flow/store.ts
index 4ffd1cb6b..dac10f2cc 100644
--- a/web/src/pages/flow/store.ts
+++ b/web/src/pages/flow/store.ts
@@ -39,6 +39,7 @@ export type RFState = {
   deleteEdgeById: (id: string) => void;
   deleteNodeById: (id: string) => void;
   findNodeByName: (operatorName: Operator) => Node | undefined;
+  findNodeById: (id: string) => Node | undefined;
 };
 
 // this is our useStore hook that we can use in our components to get parts of the store and call actions
@@ -125,6 +126,9 @@ const useGraphStore = create<RFState>()(
       findNodeByName: (name: Operator) => {
         return get().nodes.find((x) => x.data.label === name);
       },
+      findNodeById: (id: string) => {
+        return get().nodes.find((x) => x.id === id);
+      },
       updateNodeForm: (nodeId: string, values: any) => {
         set({
           nodes: get().nodes.map((node) => {
diff --git a/web/src/pages/flow/utils.ts b/web/src/pages/flow/utils.ts
index 576ef48b8..d84360bd5 100644
--- a/web/src/pages/flow/utils.ts
+++ b/web/src/pages/flow/utils.ts
@@ -114,7 +114,10 @@ const buildComponentDownstreamOrUpstream = (
 
 const removeUselessDataInTheOperator = curry(
   (operatorName: string, params: Record<string, unknown>) => {
-    if (operatorName === Operator.Generate) {
+    if (
+      operatorName === Operator.Generate ||
+      operatorName === Operator.Categorize
+    ) {
       return removeUselessFieldsFromValues(params, '');
     }
     return params;