perf: bill

client/src/pages/api/openapi/kb/pushData.ts

@@ -15,6 +15,7 @@ type DateItemType = { a: string; q: string; source?: string };
 export type Props = {
   kbId: string;
   data: DateItemType[];
+  model: string;
   mode: `${TrainingModeEnum}`;
   prompt?: string;
 };
@@ -25,14 +26,14 @@ export type Response = {
 
 const modeMaxToken = {
   [TrainingModeEnum.index]: 6000,
-  [TrainingModeEnum.qa]: 10000
+  [TrainingModeEnum.qa]: 12000
 };
 
 export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
   try {
-    const { kbId, data, mode, prompt } = req.body as Props;
+    const { kbId, data, mode, prompt, model } = req.body as Props;
 
-    if (!kbId || !Array.isArray(data)) {
+    if (!kbId || !Array.isArray(data) || !model) {
       throw new Error('缺少参数');
     }
     await connectToDatabase();
@@ -46,7 +47,8 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
         data,
         userId,
         mode,
-        prompt
+        prompt,
+        model
       })
     });
   } catch (err) {
@@ -62,7 +64,8 @@ export async function pushDataToKb({
   kbId,
   data,
   mode,
-  prompt
+  prompt,
+  model
 }: { userId: string } & Props): Promise<Response> {
   await authKb({
     userId,
@@ -79,7 +82,7 @@ export async function pushDataToKb({
     if (mode === TrainingModeEnum.qa) {
       // count token
       const token = modelToolMap.countTokens({
-        model: OpenAiChatEnum.GPT3516k,
+        model: 'gpt-3.5-turbo-16k',
         messages: [{ obj: 'System', value: item.q }]
       });
       if (token > modeMaxToken[TrainingModeEnum.qa]) {
@@ -144,6 +147,7 @@ export async function pushDataToKb({
     insertData.map((item) => ({
       q: item.q,
       a: item.a,
+      model,
       source: item.source,
       userId,
       kbId,