feat: sql 封装

src/service/events/generateQA.ts

@@ -7,7 +7,7 @@ import { ChatModelNameEnum } from '@/constants/model';
 import { pushSplitDataBill } from '@/service/events/pushBill';
 import { generateVector } from './generateVector';
 import { openaiError2 } from '../errorCode';
-import { connectPg } from '@/service/pg';
+import { PgClient } from '@/service/pg';
 import { ModelSplitDataSchema } from '@/types/mongoSchema';
 
 export async function generateQA(next = false): Promise<any> {
@@ -22,7 +22,6 @@ export async function generateQA(next = false): Promise<any> {
   let dataId = null;
 
   try {
-    const pg = await connectPg();
     // 找出一个需要生成的 dataItem
     const data = await SplitData.aggregate([
       { $match: { textList: { $exists: true, $ne: [] } } },
@@ -115,7 +114,7 @@ export async function generateQA(next = false): Promise<any> {
             };
           })
           .catch((err) => {
-            console.log('QA 拆分错误');
+            console.log('QA 拆分错误', err);
             return Promise.reject(err);
           })
       )
@@ -137,18 +136,16 @@ export async function generateQA(next = false): Promise<any> {
         textList: dataItem.textList.slice(0, -5)
       }), // 删掉后5个数据
       // 生成的内容插入 pg
-      pg.query(`INSERT INTO modelData (user_id, model_id, q, a, status) VALUES ${resultList
-        .map(
-          (item) =>
-            `('${String(dataItem.userId)}', '${String(dataItem.modelId)}', '${item.q.replace(
-              /\'/g,
-              '"'
-            )}', '${item.a.replace(/\'/g, '"')}', 'waiting')`
-        )
-        .join(',')}
-            `)
+      PgClient.insert('modelData', {
+        values: resultList.map((item) => [
+          { key: 'user_id', value: dataItem.userId },
+          { key: 'model_id', value: dataItem.modelId },
+          { key: 'q', value: item.q },
+          { key: 'a', value: item.a },
+          { key: 'status', value: 'waiting' }
+        ])
+      })
     ]);
-
     console.log('生成QA成功，time:', `${(Date.now() - startTime) / 1000}s`);
 
     generateQA(true);

src/service/events/generateVector.ts

@@ -1,8 +1,7 @@
 import { connectRedis } from '../redis';
 import { openaiCreateEmbedding, getOpenApiKey } from '../utils/openai';
 import { openaiError2 } from '../errorCode';
-import { connectPg } from '@/service/pg';
-import type { PgModelDataItemType } from '@/types/pg';
+import { connectPg, PgClient } from '@/service/pg';
 
 export async function generateVector(next = false): Promise<any> {
   if (process.env.queueTask !== '1') {
@@ -16,14 +15,12 @@ export async function generateVector(next = false): Promise<any> {
   let dataId = null;
 
   try {
-    const pg = await connectPg();
-
     // 从找出一个 status = waiting 的数据
-    const searchRes = await pg.query<PgModelDataItemType>(`SELECT id, q, user_id
-                        FROM modelData
-                        WHERE status='waiting'
-                        LIMIT 1
-                      `);
+    const searchRes = await PgClient.select('modelData', {
+      field: ['id', 'q', 'user_id'],
+      where: [['status', 'waiting']],
+      limit: 1
+    });
 
     if (searchRes.rowCount === 0) {
       console.log('没有需要生成 【向量】 的数据');
@@ -47,7 +44,9 @@ export async function generateVector(next = false): Promise<any> {
       systemKey = res.systemKey;
     } catch (error: any) {
       if (error?.code === 501) {
-        await pg.query(`DELETE FROM modelData WHERE id = '${dataId}'`);
+        await PgClient.delete('modelData', {
+          where: [['id', dataId]]
+        });
         generateVector(true);
         return;
       }
@@ -64,9 +63,13 @@ export async function generateVector(next = false): Promise<any> {
     });
 
     // 更新 pg 向量和状态数据
-    await pg.query(
-      `UPDATE modelData SET vector = '[${vector}]', status = 'ready' WHERE id = ${dataId}`
-    );
+    await PgClient.update('modelData', {
+      values: [
+        { key: 'vector', value: `[${vector}]` },
+        { key: 'status', value: `ready` }
+      ],
+      where: [['id', dataId]]
+    });
 
     console.log(`生成向量成功: ${dataItem.id}`);