From f096c44de09e06e25476b316c3d9c8192dc9b023 Mon Sep 17 00:00:00 2001
From: williamhliu <137068196+williamhliu@users.noreply.github.com>
Date: Sat, 29 Jun 2024 16:47:14 +0800
Subject: [PATCH] (feature)(supersonic-fe) add memory manage to agent (#1284)

---
 .../src/components/ChatItem/SqlItem.tsx       |  4 +-
 .../src/pages/Agent/MemorySection.tsx         | 24 +++++-
 .../src/pages/ChatPlugin/DetailModal.tsx      | 73 +++++++++++++------
 .../src/pages/ChatPlugin/service.ts           | 14 ++--
 .../src/pages/ChatPlugin/type.ts              |  1 +
 5 files changed, 79 insertions(+), 37 deletions(-)

diff --git a/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx b/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx
index 7c898a610..21cbfd9ba 100644
--- a/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx
+++ b/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx
@@ -187,9 +187,7 @@ const SqlItem: React.FC<Props> = ({
                   <div className={`${prefixCls}-few-shot-content`}>
                     <div className={`${prefixCls}-few-shot-content-item`}>
                       <div className={`${prefixCls}-few-shot-content-title`}>问题：</div>
-                      <div className={`${prefixCls}-few-shot-content-text`}>
-                        {item.questionAugmented}
-                      </div>
+                      <div className={`${prefixCls}-few-shot-content-text`}>{item.question}</div>
                     </div>
                     <div className={`${prefixCls}-few-shot-content-item`}>
                       <div className={`${prefixCls}-few-shot-content-title`}>SQL：</div>
diff --git a/webapp/packages/supersonic-fe/src/pages/Agent/MemorySection.tsx b/webapp/packages/supersonic-fe/src/pages/Agent/MemorySection.tsx
index f01249415..a9379ecae 100644
--- a/webapp/packages/supersonic-fe/src/pages/Agent/MemorySection.tsx
+++ b/webapp/packages/supersonic-fe/src/pages/Agent/MemorySection.tsx
@@ -3,10 +3,11 @@ import { EditableProTable } from '@ant-design/pro-components';
 import React, { useState } from 'react';
 import { MemoryType, ReviewEnum, StatusEnum } from './type';
 import { getMemeoryList, saveMemory } from './service';
-import { Popover, Input, Badge } from 'antd';
+import { Popover, Input, Badge, Radio } from 'antd';
 import styles from './style.less';
 
 const { TextArea } = Input;
+const RadioGroup = Radio.Group;
 
 const MemorySection = () => {
   const [editableKeys, setEditableRowKeys] = useState<React.Key[]>([]);
@@ -24,12 +25,18 @@ const MemorySection = () => {
       dataIndex: 'dbSchema',
       width: 300,
       valueType: 'textarea',
+      renderFormItem: (_, { record }) => (
+        <TextArea rows={3} disabled={record?.status === StatusEnum.ENABLED} />
+      ),
     },
     {
       title: '大模型解析SQL',
       dataIndex: 's2sql',
       width: 300,
       valueType: 'textarea',
+      renderFormItem: (_, { record }) => (
+        <TextArea rows={3} disabled={record?.status === StatusEnum.ENABLED} />
+      ),
     },
     {
       title: '大模型评估意见',
@@ -64,7 +71,9 @@ const MemorySection = () => {
       title: '管理员评估意见',
       dataIndex: 'humanReviewCmt',
       valueType: 'textarea',
-      renderFormItem: () => <TextArea rows={12} />,
+      renderFormItem: (_, { record }) => (
+        <TextArea rows={12} disabled={record?.status === StatusEnum.ENABLED} />
+      ),
       render: (value) => {
         return value === '-' ? (
           '-'
@@ -81,6 +90,15 @@ const MemorySection = () => {
       dataIndex: 'humanReviewRet',
       width: 150,
       valueType: 'radio',
+      renderFormItem: (_, { record }) => (
+        <RadioGroup
+          disabled={record?.status === StatusEnum.ENABLED}
+          options={[
+            { label: '正确', value: ReviewEnum.POSITIVE },
+            { label: '错误', value: ReviewEnum.NEGATIVE },
+          ]}
+        />
+      ),
       valueEnum: {
         [ReviewEnum.POSITIVE]: {
           text: '正确',
@@ -98,6 +116,8 @@ const MemorySection = () => {
       dataIndex: 'status',
       valueType: 'radio',
       width: 120,
+      tooltip:
+        '若启用，将会把这条记录加入到向量库中作为样例召回供大模型参考以及作为相似问题推荐给用户',
       valueEnum: {
         [StatusEnum.PENDING]: { text: '待定' },
         [StatusEnum.ENABLED]: {
diff --git a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/DetailModal.tsx b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/DetailModal.tsx
index 4394514b7..6a167d906 100644
--- a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/DetailModal.tsx
+++ b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/DetailModal.tsx
@@ -1,6 +1,6 @@
 import React, { useEffect, useState } from 'react';
 import { Modal, Select, Form, Input, InputNumber, message, Button, Radio, TreeSelect } from 'antd';
-import { getDimensionList, getModelList, savePlugin } from './service';
+import { getDataSetSchema, getModelList, savePlugin } from './service';
 import {
   DimensionType,
   ModelType,
@@ -27,13 +27,16 @@ type Props = {
 
 const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
   const [modelList, setModelList] = useState<ModelType[]>([]);
-  const [modelDimensionList, setModelDimensionList] = useState<Record<number, DimensionType[]>>({});
+  const [dataSetDimensionList, setDataSetDimensionList] = useState<Record<number, DimensionType[]>>(
+    {},
+  );
   const [confirmLoading, setConfirmLoading] = useState(false);
   const [pluginType, setPluginType] = useState<PluginTypeEnum>();
   const [functionName, setFunctionName] = useState<string>();
   const [functionParams, setFunctionParams] = useState<FunctionParamFormItemType[]>([]);
   const [examples, setExamples] = useState<{ id: string; question?: string }[]>([]);
   const [filters, setFilters] = useState<any[]>([]);
+  const [dataSetList, setDataSetList] = useState<number[]>([]);
   const [form] = Form.useForm();
 
   const initModelList = async () => {
@@ -52,15 +55,20 @@ const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
   }, []);
 
   const initModelDimensions = async (params: any) => {
-    const modelIds = params
-      .filter((param: any) => !!param.modelId)
-      .map((param: any) => param.modelId);
-    const res = await Promise.all(modelIds.map((modelId: number) => getDimensionList(modelId)));
-    setModelDimensionList(
-      modelIds.reduce((result: Record<number, DimensionType[]>, modelId: number, index: number) => {
-        result[modelId] = res[index].data.list;
-        return result;
-      }, {}),
+    const dataSetIds = params
+      .filter((param: any) => !!param.dataSetId)
+      .map((param: any) => param.dataSetId);
+    const res = await Promise.all(
+      dataSetIds.map((dataSetId: number) => getDataSetSchema(dataSetId)),
+    );
+    setDataSetDimensionList(
+      dataSetIds.reduce(
+        (result: Record<number, DimensionType[]>, dataSetId: number, index: number) => {
+          result[dataSetId] = res[index].data.dimensions;
+          return result;
+        },
+        {},
+      ),
     );
   };
 
@@ -75,6 +83,7 @@ const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
         url: detail.config?.url,
         height,
       });
+      setDataSetList(detail.dataSetList || []);
       if (paramOptions?.length > 0) {
         const params = paramOptions.filter(
           (option: any) => option.paramType !== ParamTypeEnum.FORWARD,
@@ -170,23 +179,39 @@ const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
   };
 
   const updateDimensionList = async (value: number) => {
-    if (modelDimensionList[value]) {
+    if (dataSetDimensionList[value]) {
       return;
     }
-    const res = await getDimensionList(value);
-    setModelDimensionList({ ...modelDimensionList, [value]: res.data.list });
+    const res = await getDataSetSchema(value);
+    setDataSetDimensionList({ ...dataSetDimensionList, [value]: res.data.dimensions });
+  };
+
+  const getDataSetList = () => {
+    const list: any[] = [];
+    traverseTree(modelList, (node: any) => {
+      if (node.type === 'DATASET' && dataSetList.includes(node.id)) {
+        list.push(node);
+      }
+    });
+    return list;
+  };
+
+  const onValuesChange = (value: PluginType) => {
+    if (value.dataSetList) {
+      setDataSetList(value.dataSetList);
+    }
   };
 
   return (
     <Modal
       open
       title={detail ? '编辑插件' : '新建插件'}
-      width={900}
+      width={920}
       confirmLoading={confirmLoading}
       onOk={onOk}
       onCancel={onCancel}
     >
-      <Form {...layout} form={form} style={{ maxWidth: 820 }}>
+      <Form {...layout} form={form} style={{ maxWidth: 820 }} onValuesChange={onValuesChange}>
         <FormItem name="dataSetList" label="数据集">
           <TreeSelect
             treeData={modelList}
@@ -322,10 +347,10 @@ const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
                       {filter.paramType === ParamTypeEnum.SEMANTIC && (
                         <>
                           <Select
-                            placeholder="主题域"
-                            options={modelList.map((model) => ({
-                              label: model.name,
-                              value: model.id,
+                            placeholder="数据集"
+                            options={getDataSetList().map((dataSet) => ({
+                              label: dataSet.name,
+                              value: dataSet.id,
                             }))}
                             showSearch
                             filterOption={(input, option) =>
@@ -335,16 +360,16 @@ const DetailModal: React.FC<Props> = ({ detail, onSubmit, onCancel }) => {
                             }
                             className={styles.filterParamName}
                             allowClear
-                            value={filter.modelId}
+                            value={filter.dataSetId}
                             onChange={(value) => {
-                              filter.modelId = value;
+                              filter.dataSetId = value;
                               setFilters([...filters]);
                               updateDimensionList(value);
                             }}
                           />
                           <Select
-                            placeholder="请选择维度，需先选择主题域"
-                            options={(modelDimensionList[filter.modelId] || []).map(
+                            placeholder="请选择维度，需先选择数据集"
+                            options={(dataSetDimensionList[filter.dataSetId] || []).map(
                               (dimension) => ({
                                 label: dimension.name,
                                 value: `${dimension.id}`,
diff --git a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/service.ts b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/service.ts
index f5018adfd..c8f0f47df 100644
--- a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/service.ts
+++ b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/service.ts
@@ -27,13 +27,11 @@ export function getModelList() {
   });
 }
 
-export function getDimensionList(modelId: number) {
-  return request<Result<{ list: DimensionType[] }>>('/api/semantic/dimension/queryDimension', {
-    method: 'POST',
-    data: {
-      modelIds: [modelId],
-      current: 1,
-      pageSize: 2000,
+export function getDataSetSchema(dataSetId: number) {
+  return request<Result<{ list: DimensionType[] }>>(
+    `/api/chat/conf/getDataSetSchema/${dataSetId}`,
+    {
+      method: 'GET',
     },
-  });
+  );
 }
diff --git a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/type.ts b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/type.ts
index 2cbb18a41..a79ec2d75 100644
--- a/webapp/packages/supersonic-fe/src/pages/ChatPlugin/type.ts
+++ b/webapp/packages/supersonic-fe/src/pages/ChatPlugin/type.ts
@@ -26,6 +26,7 @@ export enum ParamTypeEnum {
 export type PluginType = {
   id: number;
   type: PluginTypeEnum;
+  dataSetList: number[];
   modelList: number[];
   pattern: string;
   parseMode: ParseModeEnum;