116 lines
2.7 KiB
TypeScript
116 lines
2.7 KiB
TypeScript
import type { NextApiRequest, NextApiResponse } from 'next';
|
|
import { jsonRes } from '@/service/response';
|
|
import { PgClient } from '@/service/pg';
|
|
import { withNextCors } from '@/service/utils/tools';
|
|
import type { ChatItemType } from '@/types/chat';
|
|
import { ChatRoleEnum } from '@/constants/chat';
|
|
import { openaiEmbedding_system } from '../../plugin/openaiEmbedding';
|
|
import { modelToolMap } from '@/utils/plugin';
|
|
|
|
export type QuoteItemType = {
|
|
id: string;
|
|
q: string;
|
|
a: string;
|
|
source?: string;
|
|
};
|
|
type Props = {
|
|
kb_ids: string[];
|
|
history: ChatItemType[];
|
|
similarity: number;
|
|
limit: number;
|
|
maxToken: number;
|
|
userChatInput: string;
|
|
stream?: boolean;
|
|
};
|
|
type Response = {
|
|
rawSearch: QuoteItemType[];
|
|
isEmpty?: boolean;
|
|
quotePrompt: string;
|
|
};
|
|
|
|
export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
|
|
try {
|
|
const {
|
|
kb_ids = [],
|
|
history = [],
|
|
similarity,
|
|
limit,
|
|
maxToken,
|
|
userChatInput
|
|
} = req.body as Props;
|
|
|
|
if (!similarity || !Array.isArray(kb_ids)) {
|
|
throw new Error('params is error');
|
|
}
|
|
|
|
const result = await appKbSearch({
|
|
kb_ids,
|
|
history,
|
|
similarity,
|
|
limit,
|
|
maxToken,
|
|
userChatInput
|
|
});
|
|
|
|
jsonRes<Response>(res, {
|
|
data: result
|
|
});
|
|
} catch (err) {
|
|
console.log(err);
|
|
jsonRes(res, {
|
|
code: 500,
|
|
error: err
|
|
});
|
|
}
|
|
});
|
|
|
|
export async function appKbSearch({
|
|
kb_ids = [],
|
|
history = [],
|
|
similarity = 0.8,
|
|
limit = 5,
|
|
maxToken = 2500,
|
|
userChatInput
|
|
}: Props): Promise<Response> {
|
|
// get vector
|
|
const promptVector = await openaiEmbedding_system({
|
|
input: [userChatInput]
|
|
});
|
|
|
|
// search kb
|
|
const res: any = await PgClient.query(
|
|
`BEGIN;
|
|
SET LOCAL ivfflat.probes = ${global.systemEnv.pgIvfflatProbe || 10};
|
|
select id,q,a,source from modelData where kb_id IN (${kb_ids
|
|
.map((item) => `'${item}'`)
|
|
.join(',')}) AND vector <#> '[${promptVector[0]}]' < -${similarity} order by vector <#> '[${
|
|
promptVector[0]
|
|
}]' limit ${limit};
|
|
COMMIT;`
|
|
);
|
|
|
|
const searchRes: QuoteItemType[] = res?.[2]?.rows || [];
|
|
|
|
// filter part quote by maxToken
|
|
const sliceResult = modelToolMap['gpt-3.5-turbo']
|
|
.tokenSlice({
|
|
maxToken,
|
|
messages: searchRes.map((item, i) => ({
|
|
obj: ChatRoleEnum.System,
|
|
value: `${i + 1}: [${item.q}\n${item.a}]`
|
|
}))
|
|
})
|
|
.map((item) => item.value)
|
|
.join('\n')
|
|
.trim();
|
|
|
|
// slice filterSearch
|
|
const rawSearch = searchRes.slice(0, sliceResult.length);
|
|
|
|
return {
|
|
isEmpty: rawSearch.length === 0,
|
|
rawSearch,
|
|
quotePrompt: sliceResult ? `知识库:\n${sliceResult}` : ''
|
|
};
|
|
}
|