perf: chat framwork

2023-05-03 15:28:25 +08:00
parent 91decc3683
commit 00a99261ae
23 changed files with 811 additions and 1011 deletions
--- a/src/service/events/generateQA.ts
+++ b/src/service/events/generateQA.ts
@@ -1,14 +1,13 @@
 import { SplitData } from '@/service/mongo';
-import { getOpenAIApi } from '@/service/utils/auth';
-import { axiosConfig } from '@/service/utils/tools';
-import { getOpenApiKey } from '../utils/openai';
-import type { ChatCompletionRequestMessage } from 'openai';
+import { getApiKey } from '../utils/auth';
 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 { ModelSplitDataSchema } from '@/types/mongoSchema';
+import { modelServiceToolMap } from '../utils/chat';
+import { ChatRoleEnum } from '@/constants/chat';

 export async function generateQA(next = false): Promise<any> {
  if (process.env.queueTask !== '1') {
@@ -47,11 +46,11 @@ export async function generateQA(next = false): Promise<any> {

    // 获取 openapi Key
    let userApiKey = '',
-      systemKey = '';
+      systemApiKey = '';
    try {
-      const key = await getOpenApiKey(dataItem.userId);
+      const key = await getApiKey({ model: OpenAiChatEnum.GPT35, userId: dataItem.userId });
      userApiKey = key.userApiKey;
-      systemKey = key.systemKey;
+      systemApiKey = key.systemApiKey;
    } catch (error: any) {
      if (error?.code === 501) {
        // 余额不够了, 清空该记录
@@ -69,55 +68,44 @@ export async function generateQA(next = false): Promise<any> {

    const startTime = Date.now();

-    // 获取 openai 请求实例
-    const chatAPI = getOpenAIApi(userApiKey || systemKey);
-    const systemPrompt: ChatCompletionRequestMessage = {
-      role: 'system',
-      content: `你是出题人
-${dataItem.prompt || '下面是"一段长文本"'}
-从中选出5至20个题目和答案.答案详细.按格式返回: Q1:
-A1:
-Q2:
-A2:
-...`
-    };
-
    // 请求 chatgpt 获取回答
    const response = await Promise.allSettled(
      textList.map((text) =>
-        chatAPI
-          .createChatCompletion(
-            {
-              model: OpenAiChatEnum.GPT35,
-              temperature: 0.8,
-              n: 1,
-              messages: [
-                systemPrompt,
-                {
-                  role: 'user',
-                  content: text
-                }
-              ]
-            },
-            {
-              timeout: 180000,
-              ...axiosConfig()
-            }
-          )
-          .then((res) => {
-            const rawContent = res?.data.choices[0].message?.content || ''; // chatgpt 原本的回复
-            const result = formatSplitText(res?.data.choices[0].message?.content || ''); // 格式化后的QA对
+        modelServiceToolMap[OpenAiChatEnum.GPT35]
+          .chatCompletion({
+            apiKey: userApiKey || systemApiKey,
+            temperature: 0.8,
+            messages: [
+              {
+                obj: ChatRoleEnum.System,
+                value: `你是出题人
+          ${dataItem.prompt || '下面是"一段长文本"'}
+          从中选出5至20个题目和答案.答案详细.按格式返回: Q1:
+          A1:
+          Q2:
+          A2:
+          ...`
+              },
+              {
+                obj: 'Human',
+                value: text
+              }
+            ],
+            stream: false
+          })
+          .then(({ totalTokens, responseText, responseMessages }) => {
+            const result = formatSplitText(responseText); // 格式化后的QA对
            console.log(`split result length: `, result.length);
            // 计费
            pushSplitDataBill({
              isPay: !userApiKey && result.length > 0,
              userId: dataItem.userId,
              type: 'QA',
-              text: systemPrompt.content + text + rawContent,
-              tokenLen: res.data.usage?.total_tokens || 0
+              textLen: responseMessages.map((item) => item.value).join('').length,
+              totalTokens
            });
            return {
-              rawContent,
+              rawContent: responseText,
              result
            };
          })
--- a/src/service/events/generateVector.ts
+++ b/src/service/events/generateVector.ts
@@ -1,6 +1,8 @@
-import { openaiCreateEmbedding, getOpenApiKey } from '../utils/openai';
+import { openaiCreateEmbedding } from '../utils/chat/openai';
+import { getApiKey } from '../utils/auth';
 import { openaiError2 } from '../errorCode';
 import { PgClient } from '@/service/pg';
+import { embeddingModel } from '@/constants/model';

 export async function generateVector(next = false): Promise<any> {
  if (process.env.queueTask !== '1') {
@@ -40,11 +42,11 @@ export async function generateVector(next = false): Promise<any> {
    dataId = dataItem.id;

    // 获取 openapi Key
-    let userApiKey, systemKey;
+    let userApiKey, systemApiKey;
    try {
-      const res = await getOpenApiKey(dataItem.userId);
+      const res = await getApiKey({ model: embeddingModel, userId: dataItem.userId });
      userApiKey = res.userApiKey;
-      systemKey = res.systemKey;
+      systemApiKey = res.systemApiKey;
    } catch (error: any) {
      if (error?.code === 501) {
        await PgClient.delete('modelData', {
@@ -61,8 +63,8 @@ export async function generateVector(next = false): Promise<any> {
    const { vector } = await openaiCreateEmbedding({
      text: dataItem.q,
      userId: dataItem.userId,
-      isPay: !userApiKey,
-      apiKey: userApiKey || systemKey
+      userApiKey,
+      systemApiKey
    });

    // 更新 pg 向量和状态数据
--- a/src/service/events/pushBill.ts
+++ b/src/service/events/pushBill.ts
@@ -1,60 +1,54 @@
 import { connectToDatabase, Bill, User } from '../mongo';
 import { ChatModelMap, OpenAiChatEnum, ChatModelType, embeddingModel } from '@/constants/model';
 import { BillTypeEnum } from '@/constants/user';
-import { countChatTokens } from '@/utils/tools';

 export const pushChatBill = async ({
  isPay,
  chatModel,
  userId,
  chatId,
-  messages
+  textLen,
+  tokens
 }: {
  isPay: boolean;
  chatModel: ChatModelType;
  userId: string;
  chatId?: '' | string;
-  messages: { role: 'system' | 'user' | 'assistant'; content: string }[];
+  textLen: number;
+  tokens: number;
 }) => {
+  console.log(`chat generate success. text len: ${textLen}. token len: ${tokens}. pay:${isPay}`);
+  if (!isPay) return;
+
  let billId = '';

  try {
-    // 计算 token 数量
-    const tokens = countChatTokens({ model: chatModel, messages });
-    const text = messages.map((item) => item.content).join('');
+    await connectToDatabase();

-    console.log(
-      `chat generate success. text len: ${text.length}. token len: ${tokens}. pay:${isPay}`
-    );
+    // 计算价格
+    const unitPrice = ChatModelMap[chatModel]?.price || 5;
+    const price = unitPrice * tokens;

-    if (isPay) {
-      await connectToDatabase();
+    try {
+      // 插入 Bill 记录
+      const res = await Bill.create({
+        userId,
+        type: 'chat',
+        modelName: chatModel,
+        chatId: chatId ? chatId : undefined,
+        textLen,
+        tokenLen: tokens,
+        price
+      });
+      billId = res._id;

-      // 计算价格
-      const unitPrice = ChatModelMap[chatModel]?.price || 5;
-      const price = unitPrice * tokens;
-
-      try {
-        // 插入 Bill 记录
-        const res = await Bill.create({
-          userId,
-          type: 'chat',
-          modelName: chatModel,
-          chatId: chatId ? chatId : undefined,
-          textLen: text.length,
-          tokenLen: tokens,
-          price
-        });
-        billId = res._id;
-
-        // 账号扣费
-        await User.findByIdAndUpdate(userId, {
-          $inc: { balance: -price }
-        });
-      } catch (error) {
-        console.log('创建账单失败:', error);
-        billId && Bill.findByIdAndDelete(billId);
-      }
+      // 账号扣费
+      await User.findByIdAndUpdate(userId, {
+        $inc: { balance: -price }
+      });
+    } catch (error) {
+      console.log('创建账单失败:', error);
+      billId && Bill.findByIdAndDelete(billId);
    }
  } catch (error) {
    console.log(error);
@@ -64,54 +58,49 @@ export const pushChatBill = async ({
 export const pushSplitDataBill = async ({
  isPay,
  userId,
-  tokenLen,
-  text,
+  totalTokens,
+  textLen,
  type
 }: {
  isPay: boolean;
  userId: string;
-  tokenLen: number;
-  text: string;
+  totalTokens: number;
+  textLen: number;
  type: `${BillTypeEnum}`;
 }) => {
-  await connectToDatabase();
+  console.log(
+    `splitData generate success. text len: ${textLen}. token len: ${totalTokens}. pay:${isPay}`
+  );
+  if (!isPay) return;

  let billId;

  try {
-    console.log(
-      `splitData generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
-    );
+    await connectToDatabase();

-    if (isPay) {
-      try {
-        // 获取模型单价格, 都是用 gpt35 拆分
-        const unitPrice = ChatModelMap[OpenAiChatEnum.GPT35]?.price || 3;
-        // 计算价格
-        const price = unitPrice * tokenLen;
+    // 获取模型单价格, 都是用 gpt35 拆分
+    const unitPrice = ChatModelMap[OpenAiChatEnum.GPT35].price || 3;
+    // 计算价格
+    const price = unitPrice * totalTokens;

-        // 插入 Bill 记录
-        const res = await Bill.create({
-          userId,
-          type,
-          modelName: OpenAiChatEnum.GPT35,
-          textLen: text.length,
-          tokenLen,
-          price
-        });
-        billId = res._id;
+    // 插入 Bill 记录
+    const res = await Bill.create({
+      userId,
+      type,
+      modelName: OpenAiChatEnum.GPT35,
+      textLen,
+      tokenLen: totalTokens,
+      price
+    });
+    billId = res._id;

-        // 账号扣费
-        await User.findByIdAndUpdate(userId, {
-          $inc: { balance: -price }
-        });
-      } catch (error) {
-        console.log('创建账单失败:', error);
-        billId && Bill.findByIdAndDelete(billId);
-      }
-    }
+    // 账号扣费
+    await User.findByIdAndUpdate(userId, {
+      $inc: { balance: -price }
+    });
  } catch (error) {
-    console.log(error);
+    console.log('创建账单失败:', error);
+    billId && Bill.findByIdAndDelete(billId);
  }
 };

@@ -126,41 +115,40 @@ export const pushGenerateVectorBill = async ({
  text: string;
  tokenLen: number;
 }) => {
-  await connectToDatabase();
+  console.log(
+    `vector generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
+  );
+  if (!isPay) return;

  let billId;

  try {
-    console.log(
-      `vector generate success. text len: ${text.length}. token len: ${tokenLen}. pay:${isPay}`
-    );
+    await connectToDatabase();

-    if (isPay) {
-      try {
-        const unitPrice = 0.4;
-        // 计算价格. 至少为1
-        let price = unitPrice * tokenLen;
-        price = price > 1 ? price : 1;
+    try {
+      const unitPrice = 0.4;
+      // 计算价格. 至少为1
+      let price = unitPrice * tokenLen;
+      price = price > 1 ? price : 1;

-        // 插入 Bill 记录
-        const res = await Bill.create({
-          userId,
-          type: BillTypeEnum.vector,
-          modelName: embeddingModel,
-          textLen: text.length,
-          tokenLen,
-          price
-        });
-        billId = res._id;
+      // 插入 Bill 记录
+      const res = await Bill.create({
+        userId,
+        type: BillTypeEnum.vector,
+        modelName: embeddingModel,
+        textLen: text.length,
+        tokenLen,
+        price
+      });
+      billId = res._id;

-        // 账号扣费
-        await User.findByIdAndUpdate(userId, {
-          $inc: { balance: -price }
-        });
-      } catch (error) {
-        console.log('创建账单失败:', error);
-        billId && Bill.findByIdAndDelete(billId);
-      }
+      // 账号扣费
+      await User.findByIdAndUpdate(userId, {
+        $inc: { balance: -price }
+      });
+    } catch (error) {
+      console.log('创建账单失败:', error);
+      billId && Bill.findByIdAndDelete(billId);
    }
  } catch (error) {
    console.log(error);