From 86f76df586f6d45552f09e5b47edd7488aa29a84 Mon Sep 17 00:00:00 2001
From: balibabu <cike8899@users.noreply.github.com>
Date: Thu, 17 Apr 2025 19:03:55 +0800
Subject: [PATCH] Feat: Retrieval test #3221 (#7121)

### What problem does this PR solve?

Feat: Retrieval test #3221

### Type of change


- [x] New Feature (non-breaking change which adds functionality)
---
 web/src/components/rerank.tsx                 |   4 +-
 .../components/similarity-slider/index.tsx    |  14 ++
 web/src/components/switch-fom-field.tsx       |  20 ++-
 .../components/use-knowledge-graph-item.tsx   |   1 +
 web/src/hooks/use-knowledge-request.ts        |  49 ++++++
 web/src/interfaces/request/knowledge.ts       |  10 ++
 web/src/pages/dataset/sidebar/index.tsx       |   9 +-
 .../pages/dataset/testing/testing-form.tsx    | 164 ++++++------------
 8 files changed, 147 insertions(+), 124 deletions(-)
 create mode 100644 web/src/hooks/use-knowledge-request.ts
 create mode 100644 web/src/interfaces/request/knowledge.ts

diff --git a/web/src/components/rerank.tsx b/web/src/components/rerank.tsx
index 6c7002071..232b3a275 100644
--- a/web/src/components/rerank.tsx
+++ b/web/src/components/rerank.tsx
@@ -106,7 +106,7 @@ function RerankFormField() {
       name={RerankId}
       render={({ field }) => (
         <FormItem>
-          <FormLabel>{t('rerankModel')}</FormLabel>
+          <FormLabel tooltip={t('rerankTip')}>{t('rerankModel')}</FormLabel>
           <FormControl>
             <Select onValueChange={field.onChange} {...field}>
               <SelectTrigger
@@ -156,7 +156,7 @@ export function RerankFormFields() {
           name={'top_k'}
           render={({ field }) => (
             <FormItem>
-              <FormLabel>{t('topK')}</FormLabel>
+              <FormLabel tooltip={t('topKTip')}>{t('topK')}</FormLabel>
               <FormControl>
                 <SingleFormSlider
                   {...field}
diff --git a/web/src/components/similarity-slider/index.tsx b/web/src/components/similarity-slider/index.tsx
index 1bb0a9440..37a3f0aad 100644
--- a/web/src/components/similarity-slider/index.tsx
+++ b/web/src/components/similarity-slider/index.tsx
@@ -1,6 +1,7 @@
 import { useTranslate } from '@/hooks/common-hooks';
 import { Form, Slider } from 'antd';
 import { useFormContext } from 'react-hook-form';
+import { z } from 'zod';
 import { SingleFormSlider } from '../ui/dual-range-slider';
 import {
   FormControl,
@@ -55,6 +56,19 @@ interface SimilaritySliderFormFieldProps {
   isTooltipShown?: boolean;
 }
 
+export const initialSimilarityThresholdValue = {
+  similarity_threshold: 0.2,
+};
+export const initialKeywordsSimilarityWeightValue = {
+  keywords_similarity_weight: 0.7,
+};
+
+export const similarityThresholdSchema = { similarity_threshold: z.number() };
+
+export const keywordsSimilarityWeightSchema = {
+  keywords_similarity_weight: z.number(),
+};
+
 export function SimilaritySliderFormField({
   vectorSimilarityWeightName = 'vector_similarity_weight',
   isTooltipShown,
diff --git a/web/src/components/switch-fom-field.tsx b/web/src/components/switch-fom-field.tsx
index d89b51562..308608c0b 100644
--- a/web/src/components/switch-fom-field.tsx
+++ b/web/src/components/switch-fom-field.tsx
@@ -5,15 +5,23 @@ import {
   FormLabel,
 } from '@/components/ui/form';
 import { Switch } from '@/components/ui/switch';
+import { cn } from '@/lib/utils';
 import { ReactNode } from 'react';
 import { useFormContext } from 'react-hook-form';
 
 interface SwitchFormItemProps {
   name: string;
   label: ReactNode;
+  vertical?: boolean;
+  tooltip?: ReactNode;
 }
 
-export function SwitchFormField({ label, name }: SwitchFormItemProps) {
+export function SwitchFormField({
+  label,
+  name,
+  vertical = true,
+  tooltip,
+}: SwitchFormItemProps) {
   const form = useFormContext();
 
   return (
@@ -21,8 +29,14 @@ export function SwitchFormField({ label, name }: SwitchFormItemProps) {
       control={form.control}
       name={name}
       render={({ field }) => (
-        <FormItem className="flex justify-between">
-          <FormLabel className="text-base">{label}</FormLabel>
+        <FormItem
+          className={cn('flex', {
+            'gap-2': vertical,
+            'flex-col': vertical,
+            'justify-between': !vertical,
+          })}
+        >
+          <FormLabel tooltip={tooltip}>{label}</FormLabel>
           <FormControl>
             <Switch
               checked={field.value}
diff --git a/web/src/components/use-knowledge-graph-item.tsx b/web/src/components/use-knowledge-graph-item.tsx
index cebf576ca..af05a931b 100644
--- a/web/src/components/use-knowledge-graph-item.tsx
+++ b/web/src/components/use-knowledge-graph-item.tsx
@@ -34,6 +34,7 @@ export function UseKnowledgeGraphFormField({
     <SwitchFormField
       name={name}
       label={t('chat.useKnowledgeGraph')}
+      tooltip={t('chat.useKnowledgeGraphTip')}
     ></SwitchFormField>
   );
 }
diff --git a/web/src/hooks/use-knowledge-request.ts b/web/src/hooks/use-knowledge-request.ts
new file mode 100644
index 000000000..e66368d28
--- /dev/null
+++ b/web/src/hooks/use-knowledge-request.ts
@@ -0,0 +1,49 @@
+import { ITestRetrievalRequestBody } from '@/interfaces/request/knowledge';
+import kbService from '@/services/knowledge-service';
+import { useQuery } from '@tanstack/react-query';
+import { useMemo, useState } from 'react';
+import { useParams } from 'umi';
+import { useSetPaginationParams } from './route-hook';
+
+export const enum KnowledgeApiAction {
+  TestRetrieval = 'testRetrieval',
+}
+
+export const useKnowledgeBaseId = () => {
+  const { id } = useParams();
+
+  return id;
+};
+
+export const useTestRetrieval = () => {
+  const knowledgeBaseId = useKnowledgeBaseId();
+  const { page, size: pageSize } = useSetPaginationParams();
+  const [values, setValues] = useState<ITestRetrievalRequestBody>();
+
+  const queryParams = useMemo(() => {
+    return {
+      ...values,
+      kb_id: values?.kb_id || knowledgeBaseId,
+      page,
+      size: pageSize,
+    };
+  }, [knowledgeBaseId, page, pageSize, values]);
+
+  const {
+    data,
+    isFetching: loading,
+    refetch,
+  } = useQuery<any>({
+    queryKey: [KnowledgeApiAction.TestRetrieval, queryParams],
+    initialData: {},
+    // enabled: !!values?.question && !!knowledgeBaseId,
+    enabled: false,
+    gcTime: 0,
+    queryFn: async () => {
+      const { data } = await kbService.retrieval_test(queryParams);
+      return data?.data ?? {};
+    },
+  });
+
+  return { data, loading, setValues, refetch };
+};
diff --git a/web/src/interfaces/request/knowledge.ts b/web/src/interfaces/request/knowledge.ts
new file mode 100644
index 000000000..75bacbc20
--- /dev/null
+++ b/web/src/interfaces/request/knowledge.ts
@@ -0,0 +1,10 @@
+export interface ITestRetrievalRequestBody {
+  question: string;
+  similarity_threshold: number;
+  keywords_similarity_weight: number;
+  rerank_id?: string;
+  top_k?: number;
+  use_kg?: boolean;
+  highlight?: boolean;
+  kb_id?: string[];
+}
diff --git a/web/src/pages/dataset/sidebar/index.tsx b/web/src/pages/dataset/sidebar/index.tsx
index b41546cd7..6b443a37b 100644
--- a/web/src/pages/dataset/sidebar/index.tsx
+++ b/web/src/pages/dataset/sidebar/index.tsx
@@ -2,7 +2,7 @@ import { Button } from '@/components/ui/button';
 import { useSecondPathName } from '@/hooks/route-hook';
 import { cn } from '@/lib/utils';
 import { Routes } from '@/routes';
-import { Banknote, LayoutGrid, Trash2, User } from 'lucide-react';
+import { Banknote, LayoutGrid, User } from 'lucide-react';
 import { useHandleMenuClick } from './hooks';
 
 const items = [
@@ -61,13 +61,6 @@ export function SideBar() {
           );
         })}
       </div>
-      <Button
-        variant="outline"
-        className="absolute bottom-6 left-6 right-6 text-colors-text-functional-danger border-colors-text-functional-danger"
-      >
-        <Trash2 />
-        Delete dataset
-      </Button>
     </aside>
   );
 }
diff --git a/web/src/pages/dataset/testing/testing-form.tsx b/web/src/pages/dataset/testing/testing-form.tsx
index bae983b92..b8037133a 100644
--- a/web/src/pages/dataset/testing/testing-form.tsx
+++ b/web/src/pages/dataset/testing/testing-form.tsx
@@ -1,160 +1,102 @@
 'use client';
 
 import { zodResolver } from '@hookform/resolvers/zod';
-import { useForm } from 'react-hook-form';
+import { useForm, useWatch } from 'react-hook-form';
 import { z } from 'zod';
 
-import { Button } from '@/components/ui/button';
+import { RerankFormFields } from '@/components/rerank';
+import {
+  initialKeywordsSimilarityWeightValue,
+  initialSimilarityThresholdValue,
+  keywordsSimilarityWeightSchema,
+  SimilaritySliderFormField,
+  similarityThresholdSchema,
+} from '@/components/similarity-slider';
 import {
   Form,
   FormControl,
-  FormDescription,
   FormField,
   FormItem,
   FormLabel,
   FormMessage,
 } from '@/components/ui/form';
-import { RAGFlowSelect } from '@/components/ui/select';
-import { FormSlider } from '@/components/ui/slider';
+import { LoadingButton } from '@/components/ui/loading-button';
 import { Textarea } from '@/components/ui/textarea';
-
-const options = [
-  { label: 'xx', value: 'xx' },
-  { label: 'ii', value: 'ii' },
-];
-
-const groupOptions = [
-  { label: 'scsdv', options },
-  { label: 'thtyu', options: [{ label: 'jj', value: 'jj' }] },
-];
-
-const formSchema = z.object({
-  username: z.number().min(2, {
-    message: 'Username must be at least 2 characters.',
-  }),
-  a: z.number().min(2, {
-    message: 'Username must be at least 2 characters.',
-  }),
-  b: z.string().min(2, {
-    message: 'Username must be at least 2 characters.',
-  }),
-  c: z.number().min(2, {
-    message: 'Username must be at least 2 characters.',
-  }),
-  d: z.string().min(2, {
-    message: 'Username must be at least 2 characters.',
-  }),
-});
+import { UseKnowledgeGraphFormField } from '@/components/use-knowledge-graph-item';
+import { useTestRetrieval } from '@/hooks/use-knowledge-request';
+import { trim } from 'lodash';
+import { useEffect } from 'react';
+import { useTranslation } from 'react-i18next';
 
 export default function TestingForm() {
+  const { t } = useTranslation();
+
+  const { loading, setValues, refetch } = useTestRetrieval();
+
+  const formSchema = z.object({
+    question: z.string().min(1, {
+      message: t('knowledgeDetails.testTextPlaceholder'),
+    }),
+    ...similarityThresholdSchema,
+    ...keywordsSimilarityWeightSchema,
+  });
+
   const form = useForm<z.infer<typeof formSchema>>({
     resolver: zodResolver(formSchema),
     defaultValues: {
-      username: 0,
+      ...initialSimilarityThresholdValue,
+      ...initialKeywordsSimilarityWeightValue,
     },
   });
 
-  function onSubmit(values: z.infer<typeof formSchema>) {
-    console.log(values);
+  const question = form.watch('question');
+
+  const values = useWatch({ control: form.control });
+
+  useEffect(() => {
+    setValues(values as Required<z.infer<typeof formSchema>>);
+  }, [setValues, values]);
+
+  function onSubmit() {
+    refetch();
   }
 
   return (
     <Form {...form}>
       <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-8">
+        <SimilaritySliderFormField
+          vectorSimilarityWeightName="keywords_similarity_weight"
+          isTooltipShown
+        ></SimilaritySliderFormField>
+        <RerankFormFields></RerankFormFields>
+        <UseKnowledgeGraphFormField name="use_kg"></UseKnowledgeGraphFormField>
         <FormField
           control={form.control}
-          name="username"
+          name="question"
           render={({ field }) => (
             <FormItem>
-              <FormLabel>Username</FormLabel>
-              <FormControl>
-                <FormSlider {...field}></FormSlider>
-              </FormControl>
-              <FormDescription>
-                This is your public display name.
-              </FormDescription>
-              <FormMessage />
-            </FormItem>
-          )}
-        />
-        <FormField
-          control={form.control}
-          name="a"
-          render={({ field }) => (
-            <FormItem>
-              <FormLabel>Username</FormLabel>
-              <FormControl>
-                <FormSlider {...field}></FormSlider>
-              </FormControl>
-              <FormDescription>
-                This is your public display name.
-              </FormDescription>
-              <FormMessage />
-            </FormItem>
-          )}
-        />
-        <FormField
-          control={form.control}
-          name="b"
-          render={({ field }) => (
-            <FormItem>
-              <FormLabel>Username</FormLabel>
-              <RAGFlowSelect
-                value={field.value}
-                onChange={field.onChange}
-                FormControlComponent={FormControl}
-                options={groupOptions}
-              ></RAGFlowSelect>
-              <FormDescription>
-                This is your public display name.
-              </FormDescription>
-              <FormMessage />
-            </FormItem>
-          )}
-        />
-        <FormField
-          control={form.control}
-          name="c"
-          render={({ field }) => (
-            <FormItem>
-              <FormLabel>Username</FormLabel>
-              <FormControl>
-                <FormSlider {...field}></FormSlider>
-              </FormControl>
-              <FormDescription>
-                This is your public display name.
-              </FormDescription>
-              <FormMessage />
-            </FormItem>
-          )}
-        />
-        <FormField
-          control={form.control}
-          name="d"
-          render={({ field }) => (
-            <FormItem>
-              <FormLabel>Username</FormLabel>
+              <FormLabel>{t('knowledgeDetails.testText')}</FormLabel>
               <FormControl>
                 <Textarea
                   {...field}
                   className="bg-colors-background-inverse-weak"
                 ></Textarea>
               </FormControl>
-              <FormDescription>
-                This is your public display name.
-              </FormDescription>
+
               <FormMessage />
             </FormItem>
           )}
         />
-        <Button
+        <LoadingButton
           variant={'tertiary'}
           size={'sm'}
           type="submit"
           className="w-full"
+          disabled={!!!trim(question)}
+          loading={loading}
         >
-          Test
-        </Button>
+          {t('knowledgeDetails.testingLabel')}
+        </LoadingButton>
       </form>
     </Form>
   );