gitea/FastGPT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 替换redis搜索

Browse Source
This commit is contained in:
archer
2023-04-19 12:00:28 +08:00
parent 867d69659f
commit 1e5714da1b
12 changed files with 147 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
src/pages/api/model/data/exportModelData.ts
View File

@@ -2,8 +2,7 @@ import type { NextApiRequest, NextApiResponse } from 'next';
import { jsonRes } from '@/service/response';
import { connectToDatabase } from '@/service/mongo';
import { authToken } from '@/service/utils/tools';
import { connectRedis } from '@/service/redis';
import { VecModelDataIdx } from '@/constants/redis';
import { PgClient } from '@/service/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -25,28 +24,23 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
const userId = await authToken(authorization);
await connectToDatabase();
const redis = await connectRedis();
// 从 redis 中获取数据
const searchRes = await redis.ft.search(
VecModelDataIdx,
`@modelId:{${modelId}} @userId:{${userId}}`,
{
RETURN: ['q', 'text'],
LIMIT: {
from: 0,
size: 10000
}
}
);
const data: [string, string][] = [];
searchRes.documents.forEach((item: any) => {
if (item.value.q && item.value.text) {
data.push([item.value.q.replace(/\n/g, '\\n'), item.value.text.replace(/\n/g, '\\n')]);
}
// 统计数据
const count = await PgClient.count('modelData', {
where: [['model_id', modelId], 'AND', ['user_id', userId]]
});
// 从 pg 中获取所有数据
const pgData = await PgClient.select<{ q: string; a: string }>('modelData', {
where: [['model_id', modelId], 'AND', ['user_id', userId]],
fields: ['q', 'a'],
order: [{ field: 'id', mode: 'DESC' }],
limit: count
});
const data: [string, string][] = pgData.rows.map((item) => [
item.q.replace(/\n/g, '\\n'),
item.a.replace(/\n/g, '\\n')
]);
jsonRes(res, {
data

2
src/pages/api/model/data/getModelData.ts
View File

@@ -37,7 +37,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
await connectToDatabase();
const searchRes = await PgClient.select<PgModelDataItemType>('modelData', {
field: ['id', 'q', 'a', 'status'],
fields: ['id', 'q', 'a', 'status'],
where: [['user_id', userId], 'AND', ['model_id', modelId]],
order: [{ field: 'id', mode: 'DESC' }],
limit: pageSize,

55
src/pages/api/model/data/pushModelDataCsv.ts
View File

@@ -3,11 +3,8 @@ import { jsonRes } from '@/service/response';
import { connectToDatabase, Model } from '@/service/mongo';
import { authToken } from '@/service/utils/tools';
import { generateVector } from '@/service/events/generateVector';
import { connectRedis } from '@/service/redis';
import { VecModelDataPrefix, ModelDataStatusEnum } from '@/constants/redis';
import { VecModelDataIdx } from '@/constants/redis';
import { customAlphabet } from 'nanoid';
const nanoid = customAlphabet('abcdefghijklmnopqrstuvwxyz1234567890', 12);
import { ModelDataStatusEnum } from '@/constants/model';
import { PgClient } from '@/service/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -29,7 +26,6 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
const userId = await authToken(authorization);
await connectToDatabase();
const redis = await connectRedis();
// 验证是否是该用户的 model
const model = await Model.findOne({
@@ -47,10 +43,18 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
try {
q = q.replace(/\\n/g, '\n');
a = a.replace(/\\n/g, '\n');
const redisSearch = await redis.ft.search(VecModelDataIdx, `@q:${q} @text:${a}`, {
RETURN: ['q', 'text']
const count = await PgClient.count('modelData', {
where: [
['user_id', userId],
'AND',
['model_id', modelId],
'AND',
['q', q],
'AND',
['a', a]
]
});
if (redisSearch.total > 0) {
if (count > 0) {
return Promise.reject('已经存在');
}
} catch (error) {
@@ -62,35 +66,26 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
});
})
);
// 过滤重复的内容
const filterData = searchRes
.filter((item) => item.status === 'fulfilled')
.map<{ q: string; a: string }>((item: any) => item.value);
// 插入 redis
const insertRedisRes = await Promise.allSettled(
filterData.map((item) => {
return redis.sendCommand([
'HMSET',
`${VecModelDataPrefix}:${nanoid()}`,
'userId',
userId,
'modelId',
String(modelId),
'q',
item.q,
'text',
item.a,
'status',
ModelDataStatusEnum.waiting
]);
})
);
// 插入 pg
const insertRes = await PgClient.insert('modelData', {
values: filterData.map((item) => [
{ key: 'user_id', value: userId },
{ key: 'model_id', value: modelId },
{ key: 'q', value: item.q },
{ key: 'a', value: item.a },
{ key: 'status', value: ModelDataStatusEnum.waiting }
])
});
generateVector();
jsonRes(res, {
data: insertRedisRes.filter((item) => item.status === 'fulfilled').length
data: insertRes.rowCount
});
} catch (err) {
jsonRes(res, {

37
src/pages/api/model/data/putModelData.ts
View File

@@ -1,13 +1,13 @@
import type { NextApiRequest, NextApiResponse } from 'next';
import { jsonRes } from '@/service/response';
import { authToken } from '@/service/utils/tools';
import { connectRedis } from '@/service/redis';
import { ModelDataStatusEnum } from '@/constants/redis';
import { generateVector } from '@/service/events/generateVector';
import { PgClient } from '@/service/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
const { dataId, text, q } = req.body as { dataId: string; text: string; q?: string };
const { dataId, a, q } = req.body as { dataId: string; a: string; q?: string };
const { authorization } = req.headers;
if (!authorization) {
@@ -21,26 +21,21 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
// 凭证校验
const userId = await authToken(authorization);
const redis = await connectRedis();
// 更新 pg 内容
await PgClient.update('modelData', {
where: [['id', dataId], 'AND', ['user_id', userId]],
values: [
{ key: 'a', value: a },
...(q
? [
{ key: 'q', value: q },
{ key: 'status', value: ModelDataStatusEnum.waiting }
]
: [])
]
});
// 校验是否为该用户的数据
const dataItemUserId = await redis.hGet(dataId, 'userId');
if (dataItemUserId !== userId) {
throw new Error('无权操作');
}
// 更新
await redis.sendCommand([
'HMSET',
dataId,
...(q ? ['q', q, 'status', ModelDataStatusEnum.waiting] : []),
'text',
text
]);
if (q) {
generateVector();
}
q && generateVector();
jsonRes(res);
} catch (err) {

23
src/pages/api/model/del.ts
View File

@@ -6,13 +6,12 @@ import { getUserApiOpenai } from '@/service/utils/openai';
import { TrainingStatusEnum } from '@/constants/model';
import { TrainingItemType } from '@/types/training';
import { httpsAgent } from '@/service/utils/tools';
import { connectRedis } from '@/service/redis';
import { VecModelDataIdx } from '@/constants/redis';
import { PgClient } from '@/service/pg';
/* 获取我的模型 */
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
const { modelId } = req.query;
const { modelId } = req.query as { modelId: string };
const { authorization } = req.headers;
if (!authorization) {
@@ -37,21 +36,11 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
}
await connectToDatabase();
const redis = await connectRedis();
// 获取 redis 中模型关联的所有数据
const searchRes = await redis.ft.search(
VecModelDataIdx,
`@modelId:{${modelId}} @userId:{${userId}}`,
{
LIMIT: {
from: 0,
size: 10000
}
}
);
// 删除 redis 内容
await Promise.all(searchRes.documents.map((item) => redis.del(item.id)));
// 删除 pg 中所有该模型的数据
await PgClient.delete('modelData', {
where: [['user_id', userId], 'AND', ['model_id', modelId]]
});
// 删除对应的聊天
await Chat.deleteMany({