From 7a38a81e12ed913a50402399875470e79854e2e2 Mon Sep 17 00:00:00 2001 From: archer <545436317@qq.com> Date: Tue, 6 Jun 2023 22:45:46 +0800 Subject: [PATCH] feat: task and inform --- src/hooks/usePagination.tsx | 2 +- src/pages/api/user/checkPayResult.ts | 20 +++++++- src/pages/api/user/inform/send.ts | 58 +++++++++++++++--------- src/pages/number/components/PayModal.tsx | 1 + src/service/events/generateQA.ts | 18 ++++++-- src/service/events/generateVector.ts | 18 ++++++-- src/service/mongo.ts | 12 ----- 7 files changed, 87 insertions(+), 42 deletions(-) diff --git a/src/hooks/usePagination.tsx b/src/hooks/usePagination.tsx index ec8e03130..a8a173adf 100644 --- a/src/hooks/usePagination.tsx +++ b/src/hooks/usePagination.tsx @@ -31,7 +31,7 @@ export const usePagination = ({ ...params }); setPageNum(num); - res.total && setTotal(res.total); + res.total !== undefined && setTotal(res.total); setData(res.data); } catch (error: any) { toast({ diff --git a/src/pages/api/user/checkPayResult.ts b/src/pages/api/user/checkPayResult.ts index 0d9934e7c..3b25d02d2 100644 --- a/src/pages/api/user/checkPayResult.ts +++ b/src/pages/api/user/checkPayResult.ts @@ -1,12 +1,13 @@ import type { NextApiRequest, NextApiResponse } from 'next'; import { jsonRes } from '@/service/response'; -import { connectToDatabase, User, Pay } from '@/service/mongo'; +import { connectToDatabase, User, Pay, TrainingData } from '@/service/mongo'; import { authUser } from '@/service/utils/auth'; import { PaySchema, UserModelSchema } from '@/types/mongoSchema'; import dayjs from 'dayjs'; import { getPayResult } from '@/service/utils/wxpay'; import { pushPromotionRecord } from '@/service/utils/promotion'; import { PRICE_SCALE } from '@/constants/common'; +import { startQueue } from '@/service/utils/tools'; /* 校验支付结果 */ export default async function handler(req: NextApiRequest, res: NextApiResponse) { @@ -75,6 +76,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) jsonRes(res, { data: '支付成功' }); + unlockTask(userId); } } catch (error) { await Pay.findByIdAndUpdate(payId, { @@ -101,3 +103,19 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) }); } } + +async function unlockTask(userId: string) { + try { + await TrainingData.updateMany( + { + userId + }, + { + lockTime: new Date('2000/1/1') + } + ); + startQueue(); + } catch (error) { + unlockTask(userId); + } +} diff --git a/src/pages/api/user/inform/send.ts b/src/pages/api/user/inform/send.ts index 462de1185..7ff6d91e1 100644 --- a/src/pages/api/user/inform/send.ts +++ b/src/pages/api/user/inform/send.ts @@ -32,30 +32,44 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) export async function sendInform({ type, title, content, userId }: Props) { if (!type || !title || !content) { - return Promise.reject('参数错误'); - } - - if (userId) { - await Inform.create({ - type, - title, - content, - userId - }); - return; } - // send to all user - const users = await User.find({}, '_id'); - await Inform.insertMany( - users.map(({ _id }) => ({ - type, - title, - content, - userId: _id - })) - ); + try { + if (userId) { + // skip it if have same inform within 5 minutes + const inform = await Inform.findOne({ + type, + title, + content, + userId, + read: false, + time: { $lte: new Date(Date.now() + 5 * 60 * 1000) } + }); - return; + if (inform) return; + + await Inform.create({ + type, + title, + content, + userId + }); + + return; + } + + // send to all user + const users = await User.find({}, '_id'); + await Inform.insertMany( + users.map(({ _id }) => ({ + type, + title, + content, + userId: _id + })) + ); + } catch (error) { + console.log('send inform error', error); + } } diff --git a/src/pages/number/components/PayModal.tsx b/src/pages/number/components/PayModal.tsx index 7d28360ef..988b0930b 100644 --- a/src/pages/number/components/PayModal.tsx +++ b/src/pages/number/components/PayModal.tsx @@ -57,6 +57,7 @@ const PayModal = ({ onClose }: { onClose: () => void }) => { return checkPayResult(payId); }, { + enabled: !!payId, refetchInterval: 2000, onSuccess(res) { if (!res) return; diff --git a/src/service/events/generateQA.ts b/src/service/events/generateQA.ts index 1e377dd86..b095c354c 100644 --- a/src/service/events/generateQA.ts +++ b/src/service/events/generateQA.ts @@ -9,6 +9,7 @@ import { BillTypeEnum } from '@/constants/user'; import { pushDataToKb } from '@/pages/api/openapi/kb/pushData'; import { TrainingModeEnum } from '@/constants/plugin'; import { ERROR_ENUM } from '../errorCode'; +import { sendInform } from '@/pages/api/user/inform/send'; const reduceQueue = () => { global.qaQueueLen = global.qaQueueLen > 0 ? global.qaQueueLen - 1 : 0; @@ -174,11 +175,22 @@ A2: } // 账号余额不足，删除任务 - if (err === ERROR_ENUM.insufficientQuota) { - console.log('余额不足，删除向量生成任务'); - await TrainingData.deleteMany({ + if (userId && err === ERROR_ENUM.insufficientQuota) { + sendInform({ + type: 'system', + title: 'QA 任务中止', + content: '由于账号余额不足，QA 任务中止，重新充值后将会继续。', userId }); + console.log('余额不足，暂停向量生成任务'); + await TrainingData.updateMany( + { + userId + }, + { + lockTime: new Date('2999/5/5') + } + ); return generateQA(); } diff --git a/src/service/events/generateVector.ts b/src/service/events/generateVector.ts index 5a22e7b47..66e0d4865 100644 --- a/src/service/events/generateVector.ts +++ b/src/service/events/generateVector.ts @@ -4,6 +4,7 @@ import { openaiEmbedding } from '@/pages/api/openapi/plugin/openaiEmbedding'; import { TrainingData } from '../models/trainingData'; import { ERROR_ENUM } from '../errorCode'; import { TrainingModeEnum } from '@/constants/plugin'; +import { sendInform } from '@/pages/api/user/inform/send'; const reduceQueue = () => { global.vectorQueueLen = global.vectorQueueLen > 0 ? global.vectorQueueLen - 1 : 0; @@ -125,11 +126,22 @@ export async function generateVector(): Promise { } // 账号余额不足，删除任务 - if (err === ERROR_ENUM.insufficientQuota) { - console.log('余额不足，删除向量生成任务'); - await TrainingData.deleteMany({ + if (userId && err === ERROR_ENUM.insufficientQuota) { + sendInform({ + type: 'system', + title: '索引生成任务中止', + content: '由于账号余额不足，索引生成任务中止，重新充值后将会继续。', userId }); + console.log('余额不足，暂停向量生成任务'); + await TrainingData.updateMany( + { + userId + }, + { + lockTime: new Date('2999/5/5') + } + ); return generateVector(); } diff --git a/src/service/mongo.ts b/src/service/mongo.ts index f93c1ee47..ae1c3b0a2 100644 --- a/src/service/mongo.ts +++ b/src/service/mongo.ts @@ -42,18 +42,6 @@ export async function connectToDatabase(): Promise { global.vectorQueueLen = 0; startQueue(); - // 5 分钟后解锁不正常的数据，并触发开始训练 - setTimeout(async () => { - await TrainingData.updateMany( - { - lockTime: { $lte: Date.now() - 5 * 60 * 1000 } - }, - { - lockTime: new Date('2000/1/1') - } - ); - startQueue(); - }, 5 * 60 * 1000); } export * from './models/authCode';