perf: generate queue
This commit is contained in:
archer
@@ -3,104 +3,109 @@ import { insertKbItem, PgClient } from '@/service/pg';
|
||||
import { openaiEmbedding } from '@/pages/api/openapi/plugin/openaiEmbedding';
|
||||
import { TrainingData } from '../models/trainingData';
|
||||
import { ERROR_ENUM } from '../errorCode';
|
||||
|
||||
// 每次最多选 5 组
|
||||
const listLen = 5;
|
||||
import { TrainingTypeEnum } from '@/constants/plugin';
|
||||
|
||||
/* 索引生成队列。每导入一次，就是一个单独的线程 */
|
||||
export async function generateVector(trainingId: string): Promise<any> {
|
||||
export async function generateVector(): Promise<any> {
|
||||
const maxProcess = Number(process.env.VECTOR_MAX_PROCESS || 10);
|
||||
|
||||
if (global.vectorQueueLen >= maxProcess) return;
|
||||
global.vectorQueueLen++;
|
||||
|
||||
let trainingId = '';
|
||||
let userId = '';
|
||||
|
||||
try {
|
||||
// 找出一个需要生成的 dataItem (2分钟锁)
|
||||
const data = await TrainingData.findOneAndUpdate(
|
||||
{
|
||||
_id: trainingId,
|
||||
lockTime: { $lte: Date.now() - 2 * 60 * 1000 }
|
||||
mode: TrainingTypeEnum.index,
|
||||
lockTime: { $lte: new Date(Date.now() - 2 * 60 * 1000) }
|
||||
},
|
||||
{
|
||||
lockTime: new Date()
|
||||
}
|
||||
);
|
||||
).select({
|
||||
_id: 1,
|
||||
userId: 1,
|
||||
kbId: 1,
|
||||
q: 1,
|
||||
a: 1
|
||||
});
|
||||
|
||||
/* 无待生成的任务 */
|
||||
if (!data) {
|
||||
await TrainingData.findOneAndDelete({
|
||||
_id: trainingId,
|
||||
qaList: [],
|
||||
vectorList: []
|
||||
});
|
||||
global.vectorQueueLen--;
|
||||
!global.vectorQueueLen && console.log(`没有需要【索引】的数据`);
|
||||
return;
|
||||
}
|
||||
|
||||
const userId = String(data.userId);
|
||||
trainingId = data._id;
|
||||
userId = String(data.userId);
|
||||
const kbId = String(data.kbId);
|
||||
|
||||
const dataItems: { q: string; a: string }[] = data.vectorList.slice(-listLen).map((item) => ({
|
||||
q: item.q,
|
||||
a: item.a
|
||||
}));
|
||||
const dataItems = [
|
||||
{
|
||||
q: data.q,
|
||||
a: data.a
|
||||
}
|
||||
];
|
||||
|
||||
// 过滤重复的 qa 内容
|
||||
const searchRes = await Promise.allSettled(
|
||||
dataItems.map(async ({ q, a = '' }) => {
|
||||
if (!q) {
|
||||
return Promise.reject('q为空');
|
||||
}
|
||||
// const searchRes = await Promise.allSettled(
|
||||
// dataItems.map(async ({ q, a = '' }) => {
|
||||
// if (!q) {
|
||||
// return Promise.reject('q为空');
|
||||
// }
|
||||
|
||||
q = q.replace(/\\n/g, '\n');
|
||||
a = a.replace(/\\n/g, '\n');
|
||||
// q = q.replace(/\\n/g, '\n');
|
||||
// a = a.replace(/\\n/g, '\n');
|
||||
|
||||
// Exactly the same data, not push
|
||||
try {
|
||||
const count = await PgClient.count('modelData', {
|
||||
where: [['user_id', userId], 'AND', ['kb_id', kbId], 'AND', ['q', q], 'AND', ['a', a]]
|
||||
});
|
||||
if (count > 0) {
|
||||
return Promise.reject('已经存在');
|
||||
}
|
||||
} catch (error) {
|
||||
error;
|
||||
}
|
||||
return Promise.resolve({
|
||||
q,
|
||||
a
|
||||
});
|
||||
})
|
||||
);
|
||||
const filterData = searchRes
|
||||
.filter((item) => item.status === 'fulfilled')
|
||||
.map<{ q: string; a: string }>((item: any) => item.value);
|
||||
// // Exactly the same data, not push
|
||||
// try {
|
||||
// const count = await PgClient.count('modelData', {
|
||||
// where: [['user_id', userId], 'AND', ['kb_id', kbId], 'AND', ['q', q], 'AND', ['a', a]]
|
||||
// });
|
||||
|
||||
if (filterData.length > 0) {
|
||||
// 生成词向量
|
||||
const vectors = await openaiEmbedding({
|
||||
input: filterData.map((item) => item.q),
|
||||
userId,
|
||||
type: 'training'
|
||||
});
|
||||
// if (count > 0) {
|
||||
// return Promise.reject('已经存在');
|
||||
// }
|
||||
// } catch (error) {
|
||||
// error;
|
||||
// }
|
||||
// return Promise.resolve({
|
||||
// q,
|
||||
// a
|
||||
// });
|
||||
// })
|
||||
// );
|
||||
// const filterData = searchRes
|
||||
// .filter((item) => item.status === 'fulfilled')
|
||||
// .map<{ q: string; a: string }>((item: any) => item.value);
|
||||
|
||||
// 生成结果插入到 pg
|
||||
await insertKbItem({
|
||||
userId,
|
||||
kbId,
|
||||
data: vectors.map((vector, i) => ({
|
||||
q: filterData[i].q,
|
||||
a: filterData[i].a,
|
||||
vector
|
||||
}))
|
||||
});
|
||||
}
|
||||
// 生成词向量
|
||||
const vectors = await openaiEmbedding({
|
||||
input: dataItems.map((item) => item.q),
|
||||
userId,
|
||||
type: 'training'
|
||||
});
|
||||
|
||||
// 删除 mongo 训练队列. 如果小于 n 条，整个数据删掉。 如果大于 n 条，仅删数组后 n 个
|
||||
if (data.vectorList.length <= listLen) {
|
||||
await TrainingData.findByIdAndDelete(trainingId);
|
||||
console.log(`全部向量生成完毕: ${trainingId}`);
|
||||
} else {
|
||||
await TrainingData.findByIdAndUpdate(trainingId, {
|
||||
vectorList: data.vectorList.slice(0, -listLen),
|
||||
lockTime: new Date('2000/1/1')
|
||||
});
|
||||
console.log(`生成向量成功: ${trainingId}`);
|
||||
generateVector(trainingId);
|
||||
}
|
||||
// 生成结果插入到 pg
|
||||
await insertKbItem({
|
||||
userId,
|
||||
kbId,
|
||||
data: vectors.map((vector, i) => ({
|
||||
q: dataItems[i].q,
|
||||
a: dataItems[i].a,
|
||||
vector
|
||||
}))
|
||||
});
|
||||
|
||||
// delete data from training
|
||||
await TrainingData.findByIdAndDelete(data._id);
|
||||
console.log(`生成向量成功: ${data._id}`);
|
||||
|
||||
global.vectorQueueLen--;
|
||||
generateVector();
|
||||
} catch (err: any) {
|
||||
// log
|
||||
if (err?.response) {
|
||||
@@ -113,25 +118,28 @@ export async function generateVector(trainingId: string): Promise<any> {
|
||||
// openai 账号异常或者账号余额不足，删除任务
|
||||
if (openaiError2[err?.response?.data?.error?.type] || err === ERROR_ENUM.insufficientQuota) {
|
||||
console.log('余额不足，删除向量生成任务');
|
||||
await TrainingData.findByIdAndDelete(trainingId);
|
||||
return;
|
||||
await TrainingData.deleteMany({
|
||||
userId
|
||||
});
|
||||
return generateVector();
|
||||
}
|
||||
|
||||
// unlock
|
||||
global.vectorQueueLen--;
|
||||
await TrainingData.findByIdAndUpdate(trainingId, {
|
||||
lockTime: new Date('2000/1/1')
|
||||
});
|
||||
|
||||
// 频率限制
|
||||
if (err?.response?.statusText === 'Too Many Requests') {
|
||||
console.log('生成向量次数限制，30s后尝试');
|
||||
console.log('生成向量次数限制，20s后尝试');
|
||||
return setTimeout(() => {
|
||||
generateVector(trainingId);
|
||||
}, 30000);
|
||||
generateVector();
|
||||
}, 20000);
|
||||
}
|
||||
|
||||
setTimeout(() => {
|
||||
generateVector(trainingId);
|
||||
generateVector();
|
||||
}, 1000);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user