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

feat: pg引入

Browse Source
This commit is contained in:
archer
2023-04-18 22:35:55 +08:00
parent a540ee944a
commit 9e951fbc15
17 changed files with 260 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/api/model.ts
View File

@@ -79,7 +79,7 @@ export const getWebContent = (url: string) => POST<string>(`/model/data/fetching
*/
export const postModelDataInput = (data: {
modelId: string;
data: { text: ModelDataSchema['text']; q: ModelDataSchema['q'] }[];
data: { a: ModelDataSchema['a']; q: ModelDataSchema['q'] }[];
}) => POST<number>(`/model/data/pushModelDataInput`, data);
/**
@@ -97,7 +97,7 @@ export const postModelDataCsvData = (modelId: string, data: string[][]) =>
/**
* 更新模型数据
*/
export const putModelDataById = (data: { dataId: string; text: string; q?: string }) =>
export const putModelDataById = (data: { dataId: string; a: string; q?: string }) =>
PUT('/model/data/putModelData', data);
/**
* 删除一条模型数据

8
src/constants/model.ts
View File

@@ -1,5 +1,9 @@
import type { ServiceName, ModelDataType, ModelSchema } from '@/types/mongoSchema';
import type { RedisModelDataItemType } from '@/types/redis';
export enum ModelDataStatusEnum {
ready = 'ready',
waiting = 'waiting'
}
export enum ChatModelNameEnum {
GPT35 = 'gpt-3.5-turbo',
@@ -80,7 +84,7 @@ export const formatModelStatus = {
}
};
export const ModelDataStatusMap: Record<RedisModelDataItemType['status'], string> = {
export const ModelDataStatusMap: Record<`${ModelDataStatusEnum}`, string> = {
ready: '训练完成',
waiting: '训练中'
};

12
src/pages/api/model/data/delModelDataById.ts
View File

@@ -1,7 +1,7 @@
import type { NextApiRequest, NextApiResponse } from 'next';
import { jsonRes } from '@/service/response';
import { authToken } from '@/service/utils/tools';
import { connectRedis } from '@/service/redis';
import { connectPg } from '@/service/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -21,15 +21,9 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
// 凭证校验
const userId = await authToken(authorization);
const redis = await connectRedis();
const pg = await connectPg();
await pg.query(`DELETE FROM modelData WHERE user_id = '${userId}' AND id = '${dataId}'`);
// 校验是否为该用户的数据
const dataItemUserId = await redis.hGet(dataId, 'userId');
if (dataItemUserId !== userId) {
throw new Error('无权操作');
}
// 删除
await redis.del(dataId);
jsonRes(res);
} catch (err) {
console.log(err);

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

@@ -2,8 +2,8 @@ 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 { connectPg } from '@/service/pg';
import type { PgModelDataItemType } from '@/types/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -35,34 +35,20 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
const userId = await authToken(authorization);
await connectToDatabase();
const redis = await connectRedis();
const pg = await connectPg();
// 从 redis 中获取数据
const searchRes = await redis.ft.search(
VecModelDataIdx,
`@modelId:{${modelId}} @userId:{${userId}} ${searchText ? `*${searchText}*` : ''}`,
{
RETURN: ['q', 'text', 'status'],
LIMIT: {
from: (pageNum - 1) * pageSize,
size: pageSize
},
SORTBY: {
BY: 'modelId',
DIRECTION: 'DESC'
}
}
);
const searchRes = await pg.query<PgModelDataItemType>(`SELECT id, q, a, status
FROM modelData
WHERE user_id='${userId}' AND model_id='${modelId}'
LIMIT ${pageSize} OFFSET ${pageSize * (pageNum - 1)}
`);
jsonRes(res, {
data: {
pageNum,
pageSize,
data: searchRes.documents.map((item) => ({
id: item.id,
...item.value
})),
total: searchRes.total
data: searchRes.rows,
total: searchRes.rowCount
}
});
} catch (err) {

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

@@ -4,14 +4,13 @@ import { connectToDatabase, Model } from '@/service/mongo';
import { authToken } from '@/service/utils/tools';
import { ModelDataSchema } from '@/types/mongoSchema';
import { generateVector } from '@/service/events/generateVector';
import { connectRedis } from '@/service/redis';
import { VecModelDataPrefix, ModelDataStatusEnum } from '@/constants/redis';
import { connectPg } from '@/service/pg';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
const { modelId, data } = req.body as {
modelId: string;
data: { text: ModelDataSchema['text']; q: ModelDataSchema['q'] }[];
data: { a: ModelDataSchema['a']; q: ModelDataSchema['q'] }[];
};
const { authorization } = req.headers;
@@ -27,7 +26,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
const userId = await authToken(authorization);
await connectToDatabase();
const redis = await connectRedis();
const pg = await connectPg();
// 验证是否是该用户的 model
const model = await Model.findOne({
@@ -39,29 +38,23 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
throw new Error('无权操作该模型');
}
const insertRes = await Promise.allSettled(
data.map((item) => {
return redis.sendCommand([
'HMSET',
`${VecModelDataPrefix}:${item.q.id}`,
'userId',
userId,
'modelId',
modelId,
'q',
item.q.text,
'text',
item.text,
'status',
ModelDataStatusEnum.waiting
]);
})
// 插入记录
await pg.query(
`INSERT INTO modelData (user_id, model_id, q, a, status) VALUES ${data
.map(
(item) =>
`('${userId}', '${modelId}', '${item.q.replace(/\'/g, '"')}', '${item.a.replace(
/\'/g,
'"'
)}', 'waiting')`
)
.join(',')}`
);
generateVector();
jsonRes(res, {
data: insertRes.filter((item) => item.status === 'rejected').length
data: 0
});
} catch (err) {
jsonRes(res, {

19
src/pages/model/detail/components/InputDataModal.tsx
View File

@@ -16,14 +16,14 @@ import { useToast } from '@/hooks/useToast';
import { customAlphabet } from 'nanoid';
const nanoid = customAlphabet('abcdefghijklmnopqrstuvwxyz1234567890', 12);
export type FormData = { dataId?: string; text: string; q: string };
export type FormData = { dataId?: string; a: string; q: string };
const InputDataModal = ({
onClose,
onSuccess,
modelId,
defaultValues = {
text: '',
a: '',
q: ''
}
}: {
@@ -51,11 +51,8 @@ const InputDataModal = ({
modelId: modelId,
data: [
{
text: e.text,
q: {
id: nanoid(),
text: e.q
}
a: e.a,
q: e.q
}
]
});
@@ -65,7 +62,7 @@ const InputDataModal = ({
status: res === 0 ? 'success' : 'warning'
});
reset({
text: '',
a: '',
q: ''
});
onSuccess();
@@ -81,10 +78,10 @@ const InputDataModal = ({
async (e: FormData) => {
if (!e.dataId) return;
if (e.text !== defaultValues.text || e.q !== defaultValues.q) {
if (e.a !== defaultValues.a || e.q !== defaultValues.q) {
await putModelDataById({
dataId: e.dataId,
text: e.text,
a: e.a,
q: e.q === defaultValues.q ? '' : e.q
});
onSuccess();
@@ -144,7 +141,7 @@ const InputDataModal = ({
maxLength={1000}
resize={'none'}
h={'calc(100% - 30px)'}
{...register(`text`, {
{...register(`a`, {
required: '知识点'
})}
/>

12
src/pages/model/detail/components/ModelDataCard.tsx
View File

@@ -19,7 +19,7 @@ import {
Input
} from '@chakra-ui/react';
import type { ModelSchema } from '@/types/mongoSchema';
import type { RedisModelDataItemType } from '@/types/redis';
import type { ModelDataItemType } from '@/types/model';
import { ModelDataStatusMap } from '@/constants/model';
import { usePagination } from '@/hooks/usePagination';
import {
@@ -53,9 +53,9 @@ const ModelDataCard = ({ model }: { model: ModelSchema }) => {
total,
getData,
pageNum
} = usePagination<RedisModelDataItemType>({
} = usePagination<ModelDataItemType>({
api: getModelDataList,
pageSize: 8,
pageSize: 10,
params: {
modelId: model._id,
searchText
@@ -149,7 +149,7 @@ const ModelDataCard = ({ model }: { model: ModelSchema }) => {
<MenuItem
onClick={() =>
setEditInputData({
text: '',
a: '',
q: ''
})
}
@@ -216,7 +216,7 @@ const ModelDataCard = ({ model }: { model: ModelSchema }) => {
maxH={'250px'}
overflowY={'auto'}
>
{item.text}
{item.a}
</Box>
</Td>
<Td>{ModelDataStatusMap[item.status]}</Td>
@@ -231,7 +231,7 @@ const ModelDataCard = ({ model }: { model: ModelSchema }) => {
setEditInputData({
dataId: item.id,
q: item.q,
text: item.text
a: item.a
})
}
/>

35
src/service/events/generateQA.ts
View File

@@ -6,12 +6,9 @@ import type { ChatCompletionRequestMessage } from 'openai';
import { ChatModelNameEnum } from '@/constants/model';
import { pushSplitDataBill } from '@/service/events/pushBill';
import { generateVector } from './generateVector';
import { connectRedis } from '../redis';
import { VecModelDataPrefix } from '@/constants/redis';
import { customAlphabet } from 'nanoid';
import { openaiError2 } from '../errorCode';
import { connectPg } from '@/service/pg';
import { ModelSplitDataSchema } from '@/types/mongoSchema';
const nanoid = customAlphabet('abcdefghijklmnopqrstuvwxyz1234567890', 12);
export async function generateQA(next = false): Promise<any> {
if (process.env.queueTask !== '1') {
@@ -25,7 +22,7 @@ export async function generateQA(next = false): Promise<any> {
let dataId = null;
try {
const redis = await connectRedis();
const pg = await connectPg();
// 找出一个需要生成的 dataItem
const data = await SplitData.aggregate([
{ $match: { textList: { $exists: true, $ne: [] } } },
@@ -139,23 +136,17 @@ export async function generateQA(next = false): Promise<any> {
SplitData.findByIdAndUpdate(dataItem._id, {
textList: dataItem.textList.slice(0, -5)
}), // 删掉后5个数据
...resultList.map((item) => {
// 插入 redis
return redis.sendCommand([
'HMSET',
`${VecModelDataPrefix}:${nanoid()}`,
'userId',
String(dataItem.userId),
'modelId',
String(dataItem.modelId),
'q',
item.q,
'text',
item.a,
'status',
'waiting'
]);
})
// 生成的内容插入 pg
pg.query(`INSERT INTO modelData (user_id, model_id, q, a, status) VALUES ${resultList
.map(
(item) =>
`('${String(dataItem.userId)}', '${String(dataItem.modelId)}', '${item.q.replace(
/\'/g,
'"'
)}', '${item.a.replace(/\'/g, '"')}', 'waiting')`
)
.join(',')}
`)
]);
console.log('生成QA成功，time:', `${(Date.now() - startTime) / 1000}s`);

47
src/service/events/generateVector.ts
View File

@@ -1,9 +1,8 @@
import { connectRedis } from '../redis';
import { VecModelDataIdx } from '@/constants/redis';
import { vectorToBuffer } from '@/utils/tools';
import { ModelDataStatusEnum } from '@/constants/redis';
import { openaiCreateEmbedding, getOpenApiKey } from '../utils/openai';
import { openaiError2 } from '../errorCode';
import { connectPg } from '@/service/pg';
import type { PgModelDataItemType } from '@/types/pg';
export async function generateVector(next = false): Promise<any> {
if (process.env.queueTask !== '1') {
@@ -15,32 +14,27 @@ export async function generateVector(next = false): Promise<any> {
global.generatingVector = true;
let dataId = null;
try {
const redis = await connectRedis();
const pg = await connectPg();
// 从找出一个 status = waiting 的数据
const searchRes = await redis.ft.search(
VecModelDataIdx,
`@status:{${ModelDataStatusEnum.waiting}}`,
{
RETURN: ['q', 'userId'],
LIMIT: {
from: 0,
size: 1
}
}
);
const searchRes = await pg.query<PgModelDataItemType>(`SELECT id, q, user_id
FROM modelData
WHERE status='waiting'
LIMIT 1
`);
if (searchRes.total === 0) {
if (searchRes.rowCount === 0) {
console.log('没有需要生成 【向量】 的数据');
global.generatingVector = false;
return;
}
const dataItem: { id: string; q: string; userId: string } = {
id: searchRes.documents[0].id,
q: String(searchRes.documents[0]?.value?.q || ''),
userId: String(searchRes.documents[0]?.value?.userId || '')
id: searchRes.rows[0].id,
q: searchRes.rows[0].q,
userId: searchRes.rows[0].user_id
};
dataId = dataItem.id;
@@ -53,7 +47,7 @@ export async function generateVector(next = false): Promise<any> {
systemKey = res.systemKey;
} catch (error: any) {
if (error?.code === 501) {
await redis.del(dataItem.id);
await pg.query(`DELETE FROM modelData WHERE id = '${dataId}'`);
generateVector(true);
return;
}
@@ -69,15 +63,10 @@ export async function generateVector(next = false): Promise<any> {
apiKey: userApiKey || systemKey
});
// 更新 redis 向量和状态数据
await redis.sendCommand([
'HMSET',
dataItem.id,
'vector',
vectorToBuffer(vector),
'status',
ModelDataStatusEnum.ready
]);
// 更新 pg 向量和状态数据
await pg.query(
`UPDATE modelData SET vector = '[${vector}]', status = 'ready' WHERE id = ${dataId}`
);
console.log(`生成向量成功: ${dataItem.id}`);

34
src/service/pg.ts Normal file
View File

@@ -0,0 +1,34 @@
import { Pool } from 'pg';
export const connectPg = async () => {
if (global.pgClient) {
return global.pgClient;
}
global.pgClient = new Pool({
host: process.env.PG_HOST,
port: process.env.PG_PORT ? +process.env.PG_PORT : 5432,
user: process.env.PG_USER,
password: process.env.PG_PASSWORD,
database: process.env.PG_DB_NAME,
max: 20,
idleTimeoutMillis: 30000,
connectionTimeoutMillis: 2000
});
global.pgClient.on('connect', () => {
console.log('pg connected');
});
global.pgClient.on('error', (err) => {
console.log(err);
global.pgClient = null;
});
try {
await global.pgClient.connect();
return global.pgClient;
} catch (error) {
global.pgClient = null;
return Promise.reject(error);
}
};

2
src/types/index.d.ts vendored
View File

@@ -1,10 +1,12 @@
import type { Mongoose } from 'mongoose';
import type { RedisClientType } from 'redis';
import type { Agent } from 'http';
import type { Pool } from 'pg';
declare global {
var mongodb: Mongoose | string | null;
var redisClient: RedisClientType | null;
var pgClient: Pool | null;
var generatingQA: boolean;
var generatingAbstract: boolean;
var generatingVector: boolean;

2
src/types/model.d.ts vendored
View File

@@ -12,7 +12,7 @@ export interface ModelUpdateParams {
export interface ModelDataItemType {
id: string;
status: 0 | 1; // 1代表向量生成完毕
status: 'waiting' | 'ready';
q: string; // 提问词
a: string; // 原文
modelId: string;

7
src/types/mongoSchema.d.ts vendored
View File

@@ -66,11 +66,8 @@ export interface ModelDataSchema {
_id: string;
modelId: string;
userId: string;
text: string;
q: {
id: string;
text: string;
};
a: string;
q: string;
status: ModelDataType;
}

10
src/types/pg.d.ts vendored Normal file
View File

@@ -0,0 +1,10 @@
import { ModelDataStatusEnum } from '@/constants/model';
export interface PgModelDataItemType {
id: string;
q: string;
a: string;
status: `${ModelDataStatusEnum}`;
model_id: string;
user_id: string;
}

7
src/types/redis.d.ts vendored
View File

@@ -1,7 +0,0 @@
import { ModelDataStatusEnum } from '@/constants/redis';
export interface RedisModelDataItemType {
id: string;
q: string;
text: string;
status: `${ModelDataStatusEnum}`;
}