perf: 优化tokens计算

src/constants/common.ts

@@ -43,7 +43,7 @@ wx号: YNyiqi
 | --- | --- | 
 | chatgpt - 对话 | 0.03 |
 | 知识库 - 对话 | 0.03 |
-| 知识库 - 索引 | 0.01 |
+| 知识库 - 索引 | 0.004 |
 | 文件拆分 | 0.03 |
 `;
 

src/service/events/generateQA.ts

@@ -40,7 +40,8 @@ export async function generateQA(next = false): Promise<any> {
     const textList: string[] = dataItem.textList.slice(-5);
 
     // 获取 openapi Key
-    let userApiKey, systemKey;
+    let userApiKey = '',
+      systemKey = '';
     try {
       const key = await getOpenApiKey(dataItem.userId);
       userApiKey = key.userApiKey;
@@ -93,10 +94,21 @@ export async function generateQA(next = false): Promise<any> {
               httpsAgent
             }
           )
-          .then((res) => ({
-            rawContent: res?.data.choices[0].message?.content || '', // chatgpt原本的回复
-            result: splitText(res?.data.choices[0].message?.content || '') // 格式化后的QA对
-          }))
+          .then((res) => {
+            const rawContent = res?.data.choices[0].message?.content || '';
+            // 计费
+            pushSplitDataBill({
+              isPay: !userApiKey,
+              userId: dataItem.userId,
+              type: 'QA',
+              text: systemPrompt.content + text + rawContent,
+              tokenLen: res.data.usage?.total_tokens || 0
+            });
+            return {
+              rawContent, // chatgpt 原本的回复
+              result: splitText(res?.data.choices[0].message?.content || '') // 格式化后的QA对
+            };
+          })
       )
     );
 
@@ -141,17 +153,6 @@ export async function generateQA(next = false): Promise<any> {
       resultList.length
     );
 
-    // 计费
-    pushSplitDataBill({
-      isPay: !userApiKey && resultList.length > 0,
-      userId: dataItem.userId,
-      type: 'QA',
-      text:
-        systemPrompt.content +
-        textList.join('') +
-        successResponse.map((item) => item.rawContent).join('')
-    });
-
     generateQA(true);
     generateVector();
   } catch (error: any) {

src/service/events/pushBill.ts

@@ -22,9 +22,9 @@ export const pushChatBill = async ({
 
   try {
     // 计算 token 数量
-    const tokens = encode(text);
+    const tokens = encode(text).length;
 
-    console.log(`chat generate success. text len: ${text.length}. token len: ${tokens.length}`);
+    console.log(`chat generate success. text len: ${text.length}. token len: ${tokens}`);
 
     if (isPay) {
       await connectToDatabase();
@@ -33,7 +33,7 @@ export const pushChatBill = async ({
       const modelItem = modelList.find((item) => item.model === modelName);
       // 计算价格
       const unitPrice = modelItem?.price || 5;
-      const price = unitPrice * tokens.length;
+      const price = unitPrice * tokens;
       console.log(`unit price: ${unitPrice}, price: ${formatPrice(price)}元`);
 
       try {
@@ -44,7 +44,7 @@ export const pushChatBill = async ({
           modelName,
           chatId,
           textLen: text.length,
-          tokenLen: tokens.length,
+          tokenLen: tokens,
           price
         });
         billId = res._id;
@@ -66,11 +66,13 @@ export const pushChatBill = async ({
 export const pushSplitDataBill = async ({
   isPay,
   userId,
+  tokenLen,
   text,
   type
 }: {
   isPay: boolean;
   userId: string;
+  tokenLen: number;
   text: string;
   type: DataType;
 }) => {
@@ -79,12 +81,7 @@ export const pushSplitDataBill = async ({
   let billId;
 
   try {
-    // 计算 token 数量
-    const tokens = encode(text);
-
-    console.log(
-      `splitData generate success. text len: ${text.length}. token len: ${tokens.length}`
-    );
+    console.log(`splitData generate success. text len: ${text.length}. token len: ${tokenLen}`);
 
     if (isPay) {
       try {
@@ -92,7 +89,7 @@ export const pushSplitDataBill = async ({
         const modelItem = modelList.find((item) => item.model === ChatModelNameEnum.GPT35);
         const unitPrice = modelItem?.price || 3;
         // 计算价格
-        const price = unitPrice * tokens.length;
+        const price = unitPrice * tokenLen;
 
         console.log(`price: ${formatPrice(price)}元`);
 
@@ -102,7 +99,7 @@ export const pushSplitDataBill = async ({
           type,
           modelName: ChatModelNameEnum.GPT35,
           textLen: text.length,
-          tokenLen: tokens.length,
+          tokenLen,
           price
         });
         billId = res._id;
@@ -124,27 +121,27 @@ export const pushSplitDataBill = async ({
 export const pushGenerateVectorBill = async ({
   isPay,
   userId,
-  text
+  text,
+  tokenLen
 }: {
   isPay: boolean;
   userId: string;
   text: string;
+  tokenLen: number;
 }) => {
   await connectToDatabase();
 
   let billId;
 
   try {
-    // 计算 token 数量
-    const tokens = encode(text);
-
-    console.log(`vector generate success. text len: ${text.length}. token len: ${tokens.length}`);
+    console.log(`vector generate success. text len: ${text.length}. token len: ${tokenLen}`);
 
     if (isPay) {
       try {
-        const unitPrice = 1;
-        // 计算价格
-        const price = unitPrice * tokens.length;
+        const unitPrice = 0.4;
+        // 计算价格. 至少为1
+        let price = unitPrice * tokenLen;
+        price = price > 1 ? price : 1;
 
         console.log(`price: ${formatPrice(price)}元`);
 
@@ -154,7 +151,7 @@ export const pushGenerateVectorBill = async ({
           type: BillTypeEnum.vector,
           modelName: ChatModelNameEnum.VECTOR,
           textLen: text.length,
-          tokenLen: tokens.length,
+          tokenLen,
           price
         });
         billId = res._id;

src/service/utils/openai.ts

@@ -1,4 +1,3 @@
-import axios from 'axios';
 import { getOpenAIApi } from '@/service/utils/chat';
 import { httpsAgent } from './tools';
 import { User } from '../models/user';
@@ -75,7 +74,7 @@ export const openaiCreateEmbedding = async ({
   const chatAPI = getOpenAIApi(apiKey);
 
   // 把输入的内容转成向量
-  const vector = await chatAPI
+  const res = await chatAPI
     .createEmbedding(
       {
         model: ChatModelNameEnum.VECTOR,
@@ -86,16 +85,20 @@ export const openaiCreateEmbedding = async ({
         httpsAgent
       }
     )
-    .then((res) => res?.data?.data?.[0]?.embedding || []);
+    .then((res) => ({
+      tokenLen: res.data.usage.total_tokens || 0,
+      vector: res?.data?.data?.[0]?.embedding || []
+    }));
 
   pushGenerateVectorBill({
     isPay,
     userId,
-    text
+    text,
+    tokenLen: res.tokenLen
   });
 
   return {
-    vector,
+    vector: res.vector,
     chatAPI
   };
 };