rerank api (labring#525)

liuhonghui · Nov 28, 2023 · abc1e57 · abc1e57
1 parent a74e1d7
commit abc1e57
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/packages/global/core/ai/api.d.ts b/packages/global/core/ai/api.d.ts
@@ -2,4 +2,4 @@ export type PostReRankProps = {
   query: string;
   inputs: { id: string; text: string }[];
 };
-export type PostReRankResponse = { id: string; score: number }[];
+export type PostReRankResponse = { id: string; score?: number }[];
diff --git a/projects/app/src/pages/api/v1/rerank.ts b/projects/app/src/pages/api/v1/rerank.ts
@@ -5,12 +5,12 @@ import { withNextCors } from '@fastgpt/service/common/middle/cors';
 import { pushReRankBill } from '@/service/support/wallet/bill/push';
 import { connectToDatabase } from '@/service/mongo';
 import { authTeamBalance } from '@/service/support/permission/auth/bill';
-import { PostReRankProps } from '@fastgpt/global/core/ai/api';
+import { PostReRankProps, PostReRankResponse } from '@fastgpt/global/core/ai/api';
 import { reRankRecall } from '@/service/core/ai/rerank';
 
 export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  let { query, inputs } = req.body as PostReRankProps;
   try {
-    let { query, inputs } = req.body as PostReRankProps;
     await connectToDatabase();
     const { teamId, tmbId } = await authCert({
       req,
@@ -29,14 +29,15 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
       source: 'api'
     });
 
-    jsonRes(res, {
+    jsonRes<PostReRankResponse>(res, {
       data: result
     });
   } catch (err) {
     console.log(err);
-    jsonRes(res, {
-      code: 500,
-      error: err
+    jsonRes<PostReRankResponse>(res, {
+      data: inputs.map((input) => ({
+        id: input.id
+      }))
     });
   }
 });
diff --git a/projects/app/src/service/core/dataset/data/pg.ts b/projects/app/src/service/core/dataset/data/pg.ts
@@ -398,7 +398,7 @@ export async function reRankSearchResult({
         if (!target) return null;
         return {
           ...target,
-          score: item.score || 0
+          score: item.score || target.score
         };
       })
       .filter(Boolean) as SearchDataResponseItemType[];