feat: self vector search

This commit is contained in:
archer
2023-08-26 18:24:16 +08:00
parent 13439c5183
commit be33794a5f
22 changed files with 151 additions and 71 deletions

View File

@@ -6,9 +6,9 @@ import { withNextCors } from '@/service/utils/tools';
import { getVector } from '../plugin/vector';
import type { KbTestItemType } from '@/types/plugin';
import { PgTrainingTableName } from '@/constants/plugin';
import { KB } from '@/service/mongo';
export type Props = {
model: string;
kbId: string;
text: string;
};
@@ -16,21 +16,24 @@ export type Response = KbTestItemType['results'];
export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
const { kbId, text, model } = req.body as Props;
const { kbId, text } = req.body as Props;
if (!kbId || !text || !model) {
if (!kbId || !text) {
throw new Error('缺少参数');
}
// 凭证校验
const { userId } = await authUser({ req });
const [{ userId }, kb] = await Promise.all([
authUser({ req }),
KB.findById(kbId, 'vectorModel')
]);
if (!userId) {
if (!userId || !kb) {
throw new Error('缺少用户ID');
}
const { vectors } = await getVector({
model,
model: kb.vectorModel,
userId,
input: [text]
});

View File

@@ -24,11 +24,8 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
await connectToDatabase();
// 凭证校验
const { userId } = await authUser({ req });
// find model
const kb = await KB.findById(kbId, 'model');
// auth user and get kb
const [{ userId }, kb] = await Promise.all([authUser({ req }), KB.findById(kbId, 'model')]);
if (!kb) {
throw new Error("Can't find database");

View File

@@ -2,7 +2,7 @@ import type { NextApiRequest, NextApiResponse } from 'next';
import { jsonRes } from '@/service/response';
import { connectToDatabase, KB } from '@/service/mongo';
import { authUser } from '@/service/utils/auth';
import { getModel } from '@/service/utils/data';
import { getVectorModel } from '@/service/utils/data';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -34,7 +34,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
avatar: data.avatar,
name: data.name,
userId: data.userId,
vectorModelName: getModel(data.vectorModel)?.name || 'Unknown',
vectorModel: getVectorModel(data.vectorModel),
tags: data.tags.join(' ')
}
});

View File

@@ -3,7 +3,7 @@ import { jsonRes } from '@/service/response';
import { connectToDatabase, KB } from '@/service/mongo';
import { authUser } from '@/service/utils/auth';
import { KbListItemType } from '@/types/plugin';
import { getModel } from '@/service/utils/data';
import { getVectorModel } from '@/service/utils/data';
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
@@ -25,7 +25,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
avatar: item.avatar,
name: item.name,
tags: item.tags,
vectorModelName: getModel(item.vectorModel)?.name || 'UnKnow'
vectorModel: getVectorModel(item.vectorModel)
}))
);