feat: kb crud

src/pages/api/plugins/kb/create.ts

@@ -0,0 +1,35 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, KB } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    const { name, tags } = req.body as {
+      name: string;
+      tags: string[];
+    };
+
+    if (!name) {
+      throw new Error('缺少参数');
+    }
+
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    const { _id } = await KB.create({
+      name,
+      userId,
+      tags
+    });
+
+    jsonRes(res, { data: _id });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/data/exportModelData.ts

@@ -0,0 +1,48 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import { PgClient } from '@/service/pg';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    let { kbId } = req.query as {
+      kbId: string;
+    };
+
+    if (!kbId) {
+      throw new Error('缺少参数');
+    }
+
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    // 统计数据
+    const count = await PgClient.count('modelData', {
+      where: [['kb_id', kbId], 'AND', ['user_id', userId]]
+    });
+    // 从 pg 中获取所有数据
+    const pgData = await PgClient.select<{ q: string; a: string }>('modelData', {
+      where: [['kb_id', kbId], '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
+    });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/data/getDataList.ts

@@ -0,0 +1,61 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import { PgClient } from '@/service/pg';
+import type { PgKBDataItemType } from '@/types/pg';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    let {
+      kbId,
+      pageNum = 1,
+      pageSize = 10,
+      searchText = ''
+    } = req.body as {
+      kbId: string;
+      pageNum: number;
+      pageSize: number;
+      searchText: string;
+    };
+    if (!kbId) {
+      throw new Error('缺少参数');
+    }
+
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    const where: any = [
+      ['user_id', userId],
+      'AND',
+      ['kb_id', kbId],
+      ...(searchText ? ['AND', `(q LIKE '%${searchText}%' OR a LIKE '%${searchText}%')`] : [])
+    ];
+
+    const searchRes = await PgClient.select<PgKBDataItemType>('modelData', {
+      fields: ['id', 'q', 'a', 'status'],
+      where,
+      order: [{ field: 'id', mode: 'DESC' }],
+      limit: pageSize,
+      offset: pageSize * (pageNum - 1)
+    });
+
+    jsonRes(res, {
+      data: {
+        pageNum,
+        pageSize,
+        data: searchRes.rows,
+        total: await PgClient.count('modelData', {
+          where
+        })
+      }
+    });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/data/getTrainingData.ts

@@ -0,0 +1,49 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, SplitData, Model } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import { ModelDataStatusEnum } from '@/constants/model';
+import { PgClient } from '@/service/pg';
+
+/* 拆分数据成QA */
+export default async function handler(req: NextApiRequest, res: NextApiResponse) {
+  try {
+    const { kbId } = req.query as { kbId: string };
+    if (!kbId) {
+      throw new Error('参数错误');
+    }
+    await connectToDatabase();
+
+    const userId = await authToken(req);
+
+    // split queue data
+    const data = await SplitData.find({
+      userId,
+      kbId,
+      textList: { $exists: true, $not: { $size: 0 } }
+    });
+
+    // embedding queue data
+    const embeddingData = await PgClient.count('modelData', {
+      where: [
+        ['user_id', userId],
+        'AND',
+        ['kb_id', kbId],
+        'AND',
+        ['status', ModelDataStatusEnum.waiting]
+      ]
+    });
+
+    jsonRes(res, {
+      data: {
+        splitDataQueue: data.map((item) => item.textList).flat().length,
+        embeddingQueue: embeddingData
+      }
+    });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/delete.ts

@@ -0,0 +1,43 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, KB } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import { PgClient } from '@/service/pg';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    const { id } = req.query as {
+      id: string;
+    };
+
+    if (!id) {
+      throw new Error('缺少参数');
+    }
+
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    // delete mongo data
+    await KB.findOneAndDelete({
+      _id: id,
+      userId
+    });
+
+    // delete all pg data
+    // 删除 pg 中所有该模型的数据
+    await PgClient.delete('modelData', {
+      where: [['user_id', userId], 'AND', ['kb_id', id]]
+    });
+
+    // delete related model
+
+    jsonRes(res);
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/list.ts

@@ -0,0 +1,42 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, KB } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import { PgClient } from '@/service/pg';
+import { KbItemType } from '@/types/plugin';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    const kbList = await KB.find({
+      userId
+    }).sort({ updateTime: -1 });
+
+    const data = await Promise.all(
+      kbList.map(async (item) => ({
+        _id: item._id,
+        avatar: item.avatar,
+        name: item.name,
+        userId: item.userId,
+        updateTime: item.updateTime,
+        tags: item.tags.join(' '),
+        totalData: await PgClient.count('modelData', {
+          where: [['user_id', userId], 'AND', ['kb_id', item._id]]
+        })
+      }))
+    );
+
+    jsonRes<KbItemType[]>(res, {
+      data
+    });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}

src/pages/api/plugins/kb/update.ts

@@ -0,0 +1,39 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, KB } from '@/service/mongo';
+import { authToken } from '@/service/utils/auth';
+import type { KbUpdateParams } from '@/api/plugins/kb';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    const { id, name, tags, avatar } = req.body as KbUpdateParams;
+
+    if (!id || !name) {
+      throw new Error('缺少参数');
+    }
+
+    // 凭证校验
+    const userId = await authToken(req);
+
+    await connectToDatabase();
+
+    await KB.findOneAndUpdate(
+      {
+        _id: id,
+        userId
+      },
+      {
+        avatar,
+        name,
+        tags: tags.split(' ').filter((item) => item)
+      }
+    );
+
+    jsonRes(res);
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}