perf: 知识库数据结构

src/service/events/generateQA.ts

@@ -1,10 +1,12 @@
-import { SplitData, ModelData } from '@/service/mongo';
+import { SplitData } from '@/service/mongo';
 import { getOpenAIApi } from '@/service/utils/chat';
 import { httpsAgent, getOpenApiKey } from '@/service/utils/tools';
 import type { ChatCompletionRequestMessage } from 'openai';
 import { ChatModelNameEnum } from '@/constants/model';
 import { pushSplitDataBill } from '@/service/events/pushBill';
 import { generateVector } from './generateVector';
+import { connectRedis } from '../redis';
+import { VecModelDataPrefix } from '@/constants/redis';
 import { customAlphabet } from 'nanoid';
 const nanoid = customAlphabet('abcdefghijklmnopqrstuvwxyz1234567890', 12);
 
@@ -18,6 +20,7 @@ export async function generateQA(next = false): Promise<any> {
   };
 
   try {
+    const redis = await connectRedis();
     // 找出一个需要生成的 dataItem
     const dataItem = await SplitData.findOne({
       textList: { $exists: true, $ne: [] }
@@ -29,8 +32,10 @@ export async function generateQA(next = false): Promise<any> {
       return;
     }
 
+    // 源文本
     const text = dataItem.textList[dataItem.textList.length - 1];
     if (!text) {
+      await SplitData.findByIdAndUpdate(dataItem._id, { $pop: { textList: 1 } }); // 弹出无效文本
       throw new Error('无文本');
     }
 
@@ -63,7 +68,7 @@ export async function generateQA(next = false): Promise<any> {
       .createChatCompletion(
         {
           model: ChatModelNameEnum.GPT35,
-          temperature: 0.2,
+          temperature: 0.4,
           n: 1,
           messages: [
             systemPrompt,
@@ -79,26 +84,29 @@ export async function generateQA(next = false): Promise<any> {
         }
       )
       .then((res) => ({
-        rawContent: res?.data.choices[0].message?.content || '',
-        result: splitText(res?.data.choices[0].message?.content || '')
-      })); // 从 content 中提取 QA
+        rawContent: res?.data.choices[0].message?.content || '', // chatgpt原本的回复
+        result: splitText(res?.data.choices[0].message?.content || '') // 格式化后的QA对
+      }));
 
     await Promise.allSettled([
-      SplitData.findByIdAndUpdate(dataItem._id, { $pop: { textList: 1 } }),
-      ModelData.insertMany(
-        response.result.map((item) => ({
-          modelId: dataItem.modelId,
-          userId: dataItem.userId,
-          text: item.a,
-          q: [
-            {
-              id: nanoid(),
-              text: item.q
-            }
-          ],
-          status: 1
-        }))
-      )
+      SplitData.findByIdAndUpdate(dataItem._id, { $pop: { textList: 1 } }), // 弹出已经拆分的文本
+      ...response.result.map((item) => {
+        // 插入 redis
+        return redis.sendCommand([
+          'HMSET',
+          `${VecModelDataPrefix}:${nanoid()}`,
+          'userId',
+          String(dataItem.userId),
+          'modelId',
+          String(dataItem.modelId),
+          'q',
+          item.q,
+          'text',
+          item.a,
+          'status',
+          'waiting'
+        ]);
+      })
     ]);
 
     console.log(

src/service/events/generateVector.ts

@@ -1,9 +1,9 @@
 import { getOpenAIApi } from '@/service/utils/chat';
 import { httpsAgent } from '@/service/utils/tools';
-import { ModelData } from '../models/modelData';
 import { connectRedis } from '../redis';
-import { VecModelDataIndex } from '@/constants/redis';
+import { VecModelDataIdx } from '@/constants/redis';
 import { vectorToBuffer } from '@/utils/tools';
+import { ModelDataStatusEnum } from '@/constants/redis';
 
 export async function generateVector(next = false): Promise<any> {
   if (global.generatingVector && !next) return;
@@ -12,74 +12,71 @@ export async function generateVector(next = false): Promise<any> {
   try {
     const redis = await connectRedis();
 
-    // 找出一个需要生成的 dataItem
-    const dataItem = await ModelData.findOne({
-      status: { $ne: 0 }
-    });
+    // 从找出一个 status = waiting 的数据
+    const searchRes = await redis.ft.search(
+      VecModelDataIdx,
+      `@status:{${ModelDataStatusEnum.waiting}}`,
+      {
+        RETURN: ['q'],
+        LIMIT: {
+          from: 0,
+          size: 1
+        }
+      }
+    );
 
-    if (!dataItem) {
+    if (searchRes.total === 0) {
       console.log('没有需要生成 【向量】 的数据');
       global.generatingVector = false;
       return;
     }
 
+    const dataItem: { id: string; q: string } = {
+      id: searchRes.documents[0].id,
+      q: String(searchRes.documents[0]?.value?.q || '')
+    };
+
     // 获取 openapi Key
     const openAiKey = process.env.OPENAIKEY as string;
 
     // 获取 openai 请求实例
     const chatAPI = getOpenAIApi(openAiKey);
 
-    const dataId = String(dataItem._id);
-
     // 生成词向量
-    const response = await Promise.allSettled(
-      dataItem.q.map((item, i) =>
-        chatAPI
-          .createEmbedding(
-            {
-              model: 'text-embedding-ada-002',
-              input: item.text
-            },
-            {
-              timeout: 120000,
-              httpsAgent
-            }
-          )
-          .then((res) => res?.data?.data?.[0]?.embedding || [])
-          .then((vector) =>
-            redis.sendCommand([
-              'HMSET',
-              `${VecModelDataIndex}:${item.id}`,
-              'vector',
-              vectorToBuffer(vector),
-              'modelId',
-              String(dataItem.modelId),
-              'dataId',
-              String(dataId)
-            ])
-          )
+    const vector = await chatAPI
+      .createEmbedding(
+        {
+          model: 'text-embedding-ada-002',
+          input: dataItem.q
+        },
+        {
+          timeout: 120000,
+          httpsAgent
+        }
       )
-    );
+      .then((res) => res?.data?.data?.[0]?.embedding || []);
 
-    if (response.filter((item) => item.status === 'fulfilled').length === 0) {
-      throw new Error(JSON.stringify(response));
-    }
-    // 修改该数据状态
-    await ModelData.findByIdAndUpdate(dataItem._id, {
-      status: 0
-    });
+    // 更新 redis 向量和状态数据
+    await redis.sendCommand([
+      'HMSET',
+      dataItem.id,
+      'vector',
+      vectorToBuffer(vector),
+      'status',
+      ModelDataStatusEnum.ready
+    ]);
 
-    console.log(`生成向量成功: ${dataItem._id}`);
+    console.log(`生成向量成功: ${dataItem.id}`);
 
     setTimeout(() => {
       generateVector(true);
-    }, 3000);
+    }, 2000);
   } catch (error: any) {
-    console.log(error);
-    console.log('error: 生成向量错误', error?.response?.data);
+    console.log('error: 生成向量错误', error?.response?.statusText);
+    !error?.response && console.log(error);
 
     if (error?.response?.statusText === 'Too Many Requests') {
-      console.log('次数限制，1分钟后尝试');
+      console.log('生成向量次数限制，1分钟后尝试');
       // 限制次数，1分钟后再试
       setTimeout(() => {
         generateVector(true);

src/service/models/modelData.ts

@@ -1,37 +0,0 @@
-/* 模型的知识库 */
-import { Schema, model, models, Model as MongoModel } from 'mongoose';
-import { ModelDataSchema as ModelDataType } from '@/types/mongoSchema';
-
-const ModelDataSchema = new Schema({
-  modelId: {
-    type: Schema.Types.ObjectId,
-    ref: 'model',
-    required: true
-  },
-  userId: {
-    type: Schema.Types.ObjectId,
-    ref: 'user',
-    required: true
-  },
-  text: {
-    type: String,
-    required: true
-  },
-  q: {
-    type: [
-      {
-        id: String, // 对应redis的key
-        text: String
-      }
-    ],
-    default: []
-  },
-  status: {
-    type: Number,
-    enum: [0, 1], // 1 训练ing
-    default: 1
-  }
-});
-
-export const ModelData: MongoModel<ModelDataType> =
-  models['modelData'] || model('modelData', ModelDataSchema);

src/service/mongo.ts

@@ -35,7 +35,6 @@ export async function connectToDatabase(): Promise<void> {
 export * from './models/authCode';
 export * from './models/chat';
 export * from './models/model';
-export * from './models/modelData';
 export * from './models/user';
 export * from './models/training';
 export * from './models/bill';

src/service/redis.ts

@@ -29,8 +29,8 @@ export const connectRedis = async () => {
 
     await global.redisClient.connect();
 
-    // 0 - 测试库，1 - 正式
-    await global.redisClient.select(0);
+    // 1 - 测试库，0 - 正式
+    await global.redisClient.select(process.env.NODE_ENV === 'development' ? 0 : 0);
 
     return global.redisClient;
   } catch (error) {