From b8ced5102c291fe16b2b7d0a11b91c3cd5fa6f0e Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Wed, 11 Dec 2024 14:29:50 +0800
Subject: [PATCH] feat: new retrieval result ui

---
 .../base/file-uploader/file-type-icon.tsx     |  4 +-
 .../documents/detail/completed/index.tsx      | 37 +++++----
 .../datasets/hit-testing/assets/test-data.ts  | 12 ++-
 .../components/child-chunks-item.tsx          | 31 +++++++
 .../components/chunk-detail-modal.tsx         | 83 +++++++++++++++++++
 .../hit-testing/components/result-item.tsx    | 77 ++++++++++++++++-
 .../datasets/hit-testing/components/score.tsx | 22 +++++
 .../components/datasets/hit-testing/index.tsx | 16 +---
 web/models/datasets.ts                        |  6 ++
 9 files changed, 249 insertions(+), 39 deletions(-)
 create mode 100644 web/app/components/datasets/hit-testing/components/child-chunks-item.tsx
 create mode 100644 web/app/components/datasets/hit-testing/components/chunk-detail-modal.tsx
 create mode 100644 web/app/components/datasets/hit-testing/components/score.tsx

diff --git a/web/app/components/base/file-uploader/file-type-icon.tsx b/web/app/components/base/file-uploader/file-type-icon.tsx
index 193a630dee..4e31ab66a8 100644
--- a/web/app/components/base/file-uploader/file-type-icon.tsx
+++ b/web/app/components/base/file-uploader/file-type-icon.tsx
@@ -82,8 +82,8 @@ const FileTypeIcon = ({
   size = 'sm',
   className,
 }: FileTypeIconProps) => {
-  const Icon = FILE_TYPE_ICON_MAP[type].component || FileAppearanceTypeEnum.custom
-  const color = FILE_TYPE_ICON_MAP[type].color
+  const Icon = FILE_TYPE_ICON_MAP[type]?.component || FileAppearanceTypeEnum.custom
+  const color = FILE_TYPE_ICON_MAP[type]?.color || FILE_TYPE_ICON_MAP.custom.color
 
   return <Icon className={cn('shrink-0', SizeMap[size], color, className)} />
 }
diff --git a/web/app/components/datasets/documents/detail/completed/index.tsx b/web/app/components/datasets/documents/detail/completed/index.tsx
index 0b5414a816..5cc8b3258d 100644
--- a/web/app/components/datasets/documents/detail/completed/index.tsx
+++ b/web/app/components/datasets/documents/detail/completed/index.tsx
@@ -42,21 +42,22 @@ type SegmentListContextValue = {
 
 const SegmentListContext = createContext({
   isCollapsed: true,
-  toggleCollapsed: () => {},
+  toggleCollapsed: () => { },
   fullScreen: false,
-  toggleFullScreen: () => {},
+  toggleFullScreen: () => { },
 })
 
 export const useSegmentListContext = (selector: (value: SegmentListContextValue) => any) => {
   return useContextSelector(SegmentListContext, selector)
 }
 
-export const SegmentIndexTag: FC<{ positionId?: string | number; label?: string; className?: string }> = React.memo(({ positionId, label, className }) => {
+export const SegmentIndexTag: FC<{ positionId?: string | number; label?: string; className?: string; isParentChildRetrieval?: boolean }> = React.memo(({ positionId, label, className, isParentChildRetrieval }) => {
+  const prefix = `${isParentChildRetrieval ? 'Parent-' : ''}Chunk`
   const localPositionId = useMemo(() => {
     const positionIdStr = String(positionId)
     if (positionIdStr.length >= 3)
-      return `Chunk-${positionId}`
-    return `Chunk-${positionIdStr.padStart(2, '0')}`
+      return `${prefix}-${positionId}`
+    return `${prefix}-${positionIdStr.padStart(2, '0')}`
   }, [positionId])
   return (
     <div className={cn('flex items-center', className)}>
@@ -179,7 +180,7 @@ const Completed: FC<ICompletedProps> = ({
     setSegments([])
     setSelectedSegmentIds([])
     invalidSegmentList()
-  // eslint-disable-next-line react-hooks/exhaustive-deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
   const onClickCard = (detail: SegmentDetailModel, isEditMode = false) => {
@@ -210,7 +211,7 @@ const Completed: FC<ICompletedProps> = ({
         notify({ type: 'error', message: t('common.actionMsg.modifiedUnsuccessfully') })
       },
     })
-  // eslint-disable-next-line react-hooks/exhaustive-deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [datasetId, documentId, selectedSegmentIds, segments])
 
   const { mutateAsync: deleteSegment } = useDeleteSegment()
@@ -226,7 +227,7 @@ const Completed: FC<ICompletedProps> = ({
         notify({ type: 'error', message: t('common.actionMsg.modifiedUnsuccessfully') })
       },
     })
-  // eslint-disable-next-line react-hooks/exhaustive-deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [datasetId, documentId, selectedSegmentIds])
 
   const onCancelBatchOperation = useCallback(() => {
@@ -337,7 +338,7 @@ const Completed: FC<ICompletedProps> = ({
       resetList()
       currentPage !== totalPages && setCurrentPage(totalPages)
     }
-  // eslint-disable-next-line react-hooks/exhaustive-deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [segmentListData, limit, currentPage])
 
   return (
@@ -388,7 +389,7 @@ const Completed: FC<ICompletedProps> = ({
             />
             <ChildSegmentList
               childChunks={childSegments}
-              handleInputChange={() => {}}
+              handleInputChange={() => { }}
               enabled={!archived}
             />
           </div>
@@ -443,14 +444,14 @@ const Completed: FC<ICompletedProps> = ({
       </FullScreenDrawer>
       {/* Batch Action Buttons */}
       {selectedSegmentIds.length > 0
-      && <BatchAction
-        className='absolute left-0 bottom-16 z-20'
-        selectedIds={selectedSegmentIds}
-        onBatchEnable={onChangeSwitch.bind(null, true, '')}
-        onBatchDisable={onChangeSwitch.bind(null, false, '')}
-        onBatchDelete={onDelete.bind(null, '')}
-        onCancel={onCancelBatchOperation}
-      />}
+        && <BatchAction
+          className='absolute left-0 bottom-16 z-20'
+          selectedIds={selectedSegmentIds}
+          onBatchEnable={onChangeSwitch.bind(null, true, '')}
+          onBatchDisable={onChangeSwitch.bind(null, false, '')}
+          onBatchDelete={onDelete.bind(null, '')}
+          onCancel={onCancelBatchOperation}
+        />}
     </SegmentListContext.Provider>
   )
 }
diff --git a/web/app/components/datasets/hit-testing/assets/test-data.ts b/web/app/components/datasets/hit-testing/assets/test-data.ts
index 39a9788afa..623f7e587c 100644
--- a/web/app/components/datasets/hit-testing/assets/test-data.ts
+++ b/web/app/components/datasets/hit-testing/assets/test-data.ts
@@ -1,4 +1,6 @@
-export const generalResultData = [
+import type { HitTesting } from '@/models/datasets'
+
+export const generalResultData: HitTesting[] = [
   {
     segment: {
       id: 'b621b153-f8a7-4e85-bd3d-07feaf61bd9e',
@@ -40,7 +42,13 @@ export const generalResultData = [
         doc_type: null,
       },
     },
-    child_chunks: null,
+    child_chunks: [
+      {
+        id: '1',
+        score: 0.8771945,
+        content: 'It is quite natural for academics who are continuously told to “publish or perish” to want to always create something from scratch that is their own fresh creation.',
+      },
+    ],
     score: 0.8771945,
     tsne_position: null,
   },
diff --git a/web/app/components/datasets/hit-testing/components/child-chunks-item.tsx b/web/app/components/datasets/hit-testing/components/child-chunks-item.tsx
new file mode 100644
index 0000000000..b685689b2e
--- /dev/null
+++ b/web/app/components/datasets/hit-testing/components/child-chunks-item.tsx
@@ -0,0 +1,31 @@
+'use client'
+import type { FC } from 'react'
+import React from 'react'
+import { useTranslation } from 'react-i18next'
+import { SliceContent, SliceLabel } from '../../formatted-text/flavours/shared'
+import cn from '@/utils/classnames'
+import type { HitTestingChildChunk } from '@/models/datasets'
+
+type Props = {
+  payload: HitTestingChildChunk
+  isShowAll: boolean
+}
+
+const ChildChunks: FC<Props> = ({
+  payload,
+  isShowAll,
+}) => {
+  const { t } = useTranslation()
+  const { id, score, content } = payload
+  return (
+    <div className='flex items-center space-x-2'>
+      <SliceLabel>
+        {id} {score}
+      </SliceLabel>
+      <SliceContent className={cn(!isShowAll && 'line-clamp-2')}>
+        {content}
+      </SliceContent>
+    </div>
+  )
+}
+export default React.memo(ChildChunks)
diff --git a/web/app/components/datasets/hit-testing/components/chunk-detail-modal.tsx b/web/app/components/datasets/hit-testing/components/chunk-detail-modal.tsx
new file mode 100644
index 0000000000..d7c5264c03
--- /dev/null
+++ b/web/app/components/datasets/hit-testing/components/chunk-detail-modal.tsx
@@ -0,0 +1,83 @@
+'use client'
+import type { FC } from 'react'
+import React from 'react'
+import { useTranslation } from 'react-i18next'
+import { SegmentIndexTag } from '../../documents/detail/completed'
+import Score from './score'
+import ChildChunksItem from './child-chunks-item'
+import Modal from '@/app/components/base/modal'
+import type { HitTesting } from '@/models/datasets'
+import FileIcon from '@/app/components/base/file-uploader/file-type-icon'
+import type { FileAppearanceTypeEnum } from '@/app/components/base/file-uploader/types'
+import cn from '@/utils/classnames'
+
+type Props = {
+  payload: HitTesting
+  onHide: () => void
+}
+
+const ChunkDetailModal: FC<Props> = ({
+  payload,
+  onHide,
+}) => {
+  const { t } = useTranslation()
+  const { segment, score, child_chunks } = payload
+  const { position, word_count, content, keywords, document } = segment
+  const isParentChildRetrieval = !!(child_chunks && child_chunks.length > 0)
+  const extension = document.name.split('.').slice(0, -1)[0] as FileAppearanceTypeEnum
+
+  return (
+    <Modal
+      title={t('dataset.chunkDetail')}
+      isShow
+      closable
+      onClose={onHide}
+      className={cn(isParentChildRetrieval ? '!min-w-[1200px]' : '!min-w-[720px]')}
+    >
+      <div className='flex h-'>
+        <div>
+          {/* Meta info */}
+          <div className='flex justify-between items-center'>
+            <div className='grow flex items-center space-x-2'>
+              <SegmentIndexTag
+                isParentChildRetrieval={isParentChildRetrieval}
+                positionId={position}
+                className={cn('w-fit group-hover:opacity-100')}
+              />
+              <div className='text-xs font-medium text-text-quaternary'>·</div>
+              <div className='flex items-center space-x-1'>
+                <FileIcon type={extension} size='sm' />
+                <span className='grow w-0 truncate text-text-secondary text-[13px] font-normal'>{document.name}</span>
+              </div>
+            </div>
+            <Score value={score} />
+          </div>
+          <div className=' max-h-[752px] overflow-y-auto'>
+            {content}
+          </div>
+          {!isParentChildRetrieval && keywords && keywords.length > 0 && (
+            <div>
+              <div>{t('dataset.keywords')}</div>
+              {keywords.map(keyword => (
+                <div key={keyword} className='inline-block px-1 py-0.5 bg-components-tag-bg text-components-tag-text text-xs rounded-md mr-1'>{keyword}</div>
+              ))}
+            </div>
+          )}
+        </div>
+
+        {isParentChildRetrieval && (
+          <div className='shrink-0 w-[552px] px-6'>
+            <div>{t('dataset.hitChunks', { num: child_chunks.length })}</div>
+            <div className='space-y-2'>
+              {child_chunks.map(item => (
+                <ChildChunksItem key={item.id} payload={item} isShowAll />
+              ))}
+            </div>
+          </div>
+        )}
+      </div>
+    </Modal>
+  )
+}
+
+export default React.memo(ChunkDetailModal)
diff --git a/web/app/components/datasets/hit-testing/components/result-item.tsx b/web/app/components/datasets/hit-testing/components/result-item.tsx
index 35d9d1bdf9..29c61e9de7 100644
--- a/web/app/components/datasets/hit-testing/components/result-item.tsx
+++ b/web/app/components/datasets/hit-testing/components/result-item.tsx
@@ -2,9 +2,17 @@
 import type { FC } from 'react'
 import React from 'react'
 import { useTranslation } from 'react-i18next'
+import { RiArrowDownSLine, RiArrowRightSLine, RiArrowRightUpLine } from '@remixicon/react'
+import { useBoolean } from 'ahooks'
 import { SegmentIndexTag } from '../../documents/detail/completed'
+import Score from './score'
+import ChildChunkItem from './child-chunks-item'
+import ChunkDetailModal from './chunk-detail-modal'
 import type { HitTesting } from '@/models/datasets'
 import cn from '@/utils/classnames'
+import FileIcon from '@/app/components/base/file-uploader/file-type-icon'
+import type { FileAppearanceTypeEnum } from '@/app/components/base/file-uploader/types'
+
 type Props = {
   payload: HitTesting
 }
@@ -13,20 +21,81 @@ const ResultItem: FC<Props> = ({
   payload,
 }) => {
   const { t } = useTranslation()
-  const { segment } = payload
-  const { position, word_count } = segment
+  const { segment, score, child_chunks } = payload
+  const { position, word_count, content, keywords, document } = segment
+  const isParentChildRetrieval = !!(child_chunks && child_chunks.length > 0)
+  const extension = document.name.split('.').slice(0, -1)[0] as FileAppearanceTypeEnum
+  const [isFold, {
+    toggle: toggleFold,
+  }] = useBoolean(false)
+  const Icon = isFold ? RiArrowRightSLine : RiArrowDownSLine
+
+  const [isShowDetailModal, {
+    setTrue: showDetailModal,
+    setFalse: hideDetailModal,
+  }] = useBoolean(false)
 
   return (
     <div>
+      {/* Meta info */}
       <div className='flex justify-between items-center'>
         <div className='flex items-center space-x-2'>
-          <SegmentIndexTag positionId={position} className={cn('w-fit group-hover:opacity-100')} />
+          <SegmentIndexTag
+            isParentChildRetrieval={isParentChildRetrieval}
+            positionId={position}
+            className={cn('w-fit group-hover:opacity-100')}
+          />
           <div className='text-xs font-medium text-text-quaternary'>·</div>
           <div className='system-xs-medium text-text-tertiary'>{word_count} {t('datasetDocuments.segment.characters')}</div>
         </div>
-        {/* Score */}
+        <Score value={score} />
       </div>
 
+      {/* Main */}
+      <div>
+        <div className='line-clamp-2'>{content}</div>
+        {isParentChildRetrieval && (
+          <div>
+            <div className='flex items-center space-x-0.5 text-text-secondary' onClick={toggleFold}>
+              <Icon className={cn('w-4 h-4', isFold && 'opacity-50')} />
+              <div>{t('dataset.hitChunks', { num: child_chunks.length })}</div>
+            </div>
+            {child_chunks.map(item => (
+              <ChildChunkItem key={item.id} payload={item} isShowAll={false} />
+            ))}
+          </div>
+        )}
+        {!isParentChildRetrieval && keywords && keywords.length > 0 && (
+          <div>
+            {keywords.map(keyword => (
+              <div key={keyword} className='inline-block px-1 py-0.5 bg-components-tag-bg text-components-tag-text text-xs rounded-md mr-1'>{keyword}</div>
+            ))}
+          </div>
+        )}
+      </div>
+      {/* Foot */}
+      <div className='flex justify-between items-center h-10 pl-3 pr-2 border-t border-divider-subtle'>
+        <div className='flex items-center space-x-1'>
+          <FileIcon type={extension} size='sm' />
+          <span className='grow w-0 truncate text-text-secondary text-[13px] font-normal'>{document.name}</span>
+        </div>
+        <div
+          className='flex items-center space-x-1 cursor-pointer text-text-tertiary'
+          onClick={showDetailModal}
+        >
+          <div className='text-xs uppercase'>{t('dataset.open')}</div>
+          <RiArrowRightUpLine className='size-3.5' />
+        </div>
+      </div>
+
+      {
+        isShowDetailModal && (
+          <ChunkDetailModal
+            payload={payload}
+            onHide={hideDetailModal}
+          />
+        )
+      }
     </div>
   )
 }
diff --git a/web/app/components/datasets/hit-testing/components/score.tsx b/web/app/components/datasets/hit-testing/components/score.tsx
new file mode 100644
index 0000000000..650ecd497f
--- /dev/null
+++ b/web/app/components/datasets/hit-testing/components/score.tsx
@@ -0,0 +1,22 @@
+'use client'
+import type { FC } from 'react'
+import React from 'react'
+
+type Props = {
+  value: number
+}
+
+const Score: FC<Props> = ({
+  value,
+}) => {
+  return (
+    <div className='relative items-center px-[5px] rounded-md border border-components-progress-bar-border overflow-hidden'>
+      <div className='absolute top-0 left-0 h-full bg-util-colors-blue-brand-blue-brand-100 border-r-[1.5px] border-components-progress-brand-progress ' style={{ width: `${value * 100}%` }} />
+      <div className='relative flex items-center h-4 space-x-0.5 text-util-colors-blue-brand-blue-brand-700'>
+        <div className='system-2xs-medium-uppercase'>score</div>
+        <div className='system-xs-semibold'>{value.toFixed(2)}</div>
+      </div>
+    </div>
+  )
+}
+export default React.memo(Score)
diff --git a/web/app/components/datasets/hit-testing/index.tsx b/web/app/components/datasets/hit-testing/index.tsx
index bf3c02a84f..0aef53cd6d 100644
--- a/web/app/components/datasets/hit-testing/index.tsx
+++ b/web/app/components/datasets/hit-testing/index.tsx
@@ -12,6 +12,7 @@ import s from './style.module.css'
 import HitDetail from './hit-detail'
 import ModifyRetrievalModal from './modify-retrieval-modal'
 import { generalResultData } from './assets/test-data'
+import ResultItem from './components/result-item'
 import cn from '@/utils/classnames'
 import type { ExternalKnowledgeBaseHitTestingResponse, ExternalKnowledgeBaseHitTesting as ExternalKnowledgeBaseHitTestingType, HitTestingResponse, HitTesting as HitTestingType } from '@/models/datasets'
 import Loading from '@/app/components/base/loading'
@@ -83,20 +84,9 @@ const HitTesting: FC<Props> = ({ datasetId }: Props) => {
       <div className='overflow-auto flex-1'>
         <div className={s.cardWrapper}>
           {results.map((record, idx) => (
-            <SegmentCard
+            <ResultItem
               key={idx}
-              loading={false}
-              refSource={{
-                title: record.title,
-                uri: record.metadata ? record.metadata['x-amz-bedrock-kb-source-uri'] : '',
-              }}
-              isExternal={isExternal}
-              detail={record.segment}
-              contentExternal={record.content}
-              score={record.score}
-              scene='hitTesting'
-              className='h-[216px] mb-4'
-              onClick={() => onClickCard(record)}
+              payload={record}
             />
           ))}
         </div>
diff --git a/web/models/datasets.ts b/web/models/datasets.ts
index 10495f19e7..72e7d3751d 100644
--- a/web/models/datasets.ts
+++ b/web/models/datasets.ts
@@ -479,10 +479,16 @@ export type HitTestingRecord = {
   created_at: number
 }
 
+export type HitTestingChildChunk = {
+  id: string
+  content: string
+  score: number
+}
 export type HitTesting = {
   segment: Segment
   score: number
   tsne_position: TsnePosition
+  child_chunks?: HitTestingChildChunk[] | null
 }
 
 export type ExternalKnowledgeBaseHitTesting = {