From 5128bbcce464acd170079dd6503642c0a26f050e Mon Sep 17 00:00:00 2001
From: archer <545436317@qq.com>
Date: Mon, 22 May 2023 13:16:34 +0800
Subject: [PATCH] perf: insert kb data

---
 src/pages/api/openapi/kb/pushData.ts    | 14 +++++--------
 src/pages/api/openapi/text/splitText.ts | 17 ++++++++--------
 src/service/events/generateQA.ts        | 14 +++++--------
 src/service/pg.ts                       | 26 +++++++++++++++++++++++++
 4 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/src/pages/api/openapi/kb/pushData.ts b/src/pages/api/openapi/kb/pushData.ts
index df97d1a10..bc8f2087e 100644
--- a/src/pages/api/openapi/kb/pushData.ts
+++ b/src/pages/api/openapi/kb/pushData.ts
@@ -4,7 +4,7 @@ import { jsonRes } from '@/service/response';
 import { connectToDatabase } from '@/service/mongo';
 import { authUser } from '@/service/utils/auth';
 import { generateVector } from '@/service/events/generateVector';
-import { PgClient } from '@/service/pg';
+import { PgClient, insertKbItem } from '@/service/pg';
 import { authKb } from '@/service/utils/auth';
 import { withNextCors } from '@/service/utils/tools';
 
@@ -68,14 +68,10 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
       .map<{ q: string; a: string }>((item: any) => item.value);
 
     // 插入记录
-    const insertRes = await PgClient.insert('modelData', {
-      values: filterData.map((item) => [
-        { key: 'user_id', value: userId },
-        { key: 'kb_id', value: kbId },
-        { key: 'q', value: item.q },
-        { key: 'a', value: item.a },
-        { key: 'status', value: 'waiting' }
-      ])
+    const insertRes = await insertKbItem({
+      userId,
+      kbId,
+      data: filterData
     });
 
     generateVector();
diff --git a/src/pages/api/openapi/text/splitText.ts b/src/pages/api/openapi/text/splitText.ts
index 8767bf440..489cfad77 100644
--- a/src/pages/api/openapi/text/splitText.ts
+++ b/src/pages/api/openapi/text/splitText.ts
@@ -4,7 +4,7 @@ import { connectToDatabase, SplitData } from '@/service/mongo';
 import { authKb, authUser } from '@/service/utils/auth';
 import { generateVector } from '@/service/events/generateVector';
 import { generateQA } from '@/service/events/generateQA';
-import { PgClient } from '@/service/pg';
+import { insertKbItem } from '@/service/pg';
 import { SplitTextTypEnum } from '@/constants/plugin';
 import { withNextCors } from '@/service/utils/tools';
 
@@ -43,14 +43,13 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
     } else if (mode === SplitTextTypEnum.subsection) {
       // 待优化，直接调用另一个接口
       // 插入记录
-      await PgClient.insert('modelData', {
-        values: chunks.map((item) => [
-          { key: 'user_id', value: userId },
-          { key: 'kb_id', value: kbId },
-          { key: 'q', value: item },
-          { key: 'a', value: '' },
-          { key: 'status', value: 'waiting' }
-        ])
+      await insertKbItem({
+        userId,
+        kbId,
+        data: chunks.map((item) => ({
+          q: item,
+          a: ''
+        }))
       });
 
       generateVector();
diff --git a/src/service/events/generateQA.ts b/src/service/events/generateQA.ts
index f59b2bea0..7e3cfa9f3 100644
--- a/src/service/events/generateQA.ts
+++ b/src/service/events/generateQA.ts
@@ -4,7 +4,7 @@ import { OpenAiChatEnum } from '@/constants/model';
 import { pushSplitDataBill } from '@/service/events/pushBill';
 import { generateVector } from './generateVector';
 import { openaiError2 } from '../errorCode';
-import { PgClient } from '@/service/pg';
+import { insertKbItem } from '@/service/pg';
 import { SplitDataSchema } from '@/types/mongoSchema';
 import { modelServiceToolMap } from '../utils/chat';
 import { ChatRoleEnum } from '@/constants/chat';
@@ -133,14 +133,10 @@ A2:
         textList: dataItem.textList.slice(0, -5)
       }),
       // 生成的内容插入 pg
-      PgClient.insert('modelData', {
-        values: resultList.map((item) => [
-          { key: 'user_id', value: dataItem.userId },
-          { key: 'kb_id', value: dataItem.kbId },
-          { key: 'q', value: item.q },
-          { key: 'a', value: item.a },
-          { key: 'status', value: 'waiting' }
-        ])
+      insertKbItem({
+        userId: dataItem.userId,
+        kbId: dataItem.kbId,
+        data: resultList
       })
     ]);
     console.log('生成QA成功，time:', `${(Date.now() - startTime) / 1000}s`);
diff --git a/src/service/pg.ts b/src/service/pg.ts
index 342a09d3f..b0239d920 100644
--- a/src/service/pg.ts
+++ b/src/service/pg.ts
@@ -156,3 +156,29 @@ class Pg {
 }
 
 export const PgClient = new Pg();
+
+/**
+ * data insert kb
+ */
+export const insertKbItem = ({
+  userId,
+  kbId,
+  data
+}: {
+  userId: string;
+  kbId: string;
+  data: {
+    q: string;
+    a: string;
+  }[];
+}) => {
+  return PgClient.insert('modelData', {
+    values: data.map((item) => [
+      { key: 'user_id', value: userId },
+      { key: 'kb_id', value: kbId },
+      { key: 'q', value: item.q },
+      { key: 'a', value: item.a },
+      { key: 'status', value: 'waiting' }
+    ])
+  });
+};