From c2a7079392a8379f228cacaa6ce3486cde05be43 Mon Sep 17 00:00:00 2001
From: ArnoChen <arnochen101@gmail.com>
Date: Mon, 17 Feb 2025 01:05:31 +0800
Subject: [PATCH] finish document manager

---
 lightrag_webui/src/api/lightrag.ts            |  32 +++-
 lightrag_webui/src/components/StatusCard.tsx  |   2 -
 .../documents/UploadDocumentsDialog.tsx       |   6 +-
 .../src/features/DocumentManager.tsx          | 154 +++++++++---------
 4 files changed, 104 insertions(+), 90 deletions(-)

diff --git a/lightrag_webui/src/api/lightrag.ts b/lightrag_webui/src/api/lightrag.ts
index 0f077368..6308ef62 100644
--- a/lightrag_webui/src/api/lightrag.ts
+++ b/lightrag_webui/src/api/lightrag.ts
@@ -27,8 +27,6 @@ export type LightragStatus = {
   status: 'healthy'
   working_directory: string
   input_directory: string
-  indexed_files: string[]
-  indexed_files_count: number
   configuration: {
     llm_binding: string
     llm_binding_host: string
@@ -104,11 +102,29 @@ export type QueryResponse = {
   response: string
 }
 
-export type DocumentActionResponse = {
+export type DocActionResponse = {
   status: 'success' | 'partial_success' | 'failure'
   message: string
 }
 
+export type DocStatus = 'pending' | 'processing' | 'processed' | 'failed'
+
+export type DocStatusResponse = {
+  id: string
+  content_summary: string
+  content_length: number
+  status: DocStatus
+  created_at: string
+  updated_at: string
+  chunks_count?: number
+  error?: string
+  metadata?: Record<string, any>
+}
+
+export type DocsStatusesResponse = {
+  statuses: Record<DocStatus, DocStatusResponse[]>
+}
+
 export const InvalidApiKeyError = 'Invalid API Key'
 export const RequireApiKeError = 'API Key required'
 
@@ -169,7 +185,7 @@ export const checkHealth = async (): Promise<
   }
 }
 
-export const getDocuments = async (): Promise<string[]> => {
+export const getDocuments = async (): Promise<DocsStatusesResponse> => {
   const response = await axiosInstance.get('/documents')
   return response.data
 }
@@ -253,7 +269,7 @@ export const queryTextStream = async (
 export const insertText = async (
   text: string,
   description?: string
-): Promise<DocumentActionResponse> => {
+): Promise<DocActionResponse> => {
   const response = await axiosInstance.post('/documents/text', { text, description })
   return response.data
 }
@@ -261,7 +277,7 @@ export const insertText = async (
 export const uploadDocument = async (
   file: File,
   onUploadProgress?: (percentCompleted: number) => void
-): Promise<DocumentActionResponse> => {
+): Promise<DocActionResponse> => {
   const formData = new FormData()
   formData.append('file', file)
 
@@ -284,7 +300,7 @@ export const uploadDocument = async (
 export const batchUploadDocuments = async (
   files: File[],
   onUploadProgress?: (fileName: string, percentCompleted: number) => void
-): Promise<DocumentActionResponse[]> => {
+): Promise<DocActionResponse[]> => {
   return await Promise.all(
     files.map(async (file) => {
       return await uploadDocument(file, (percentCompleted) => {
@@ -294,7 +310,7 @@ export const batchUploadDocuments = async (
   )
 }
 
-export const clearDocuments = async (): Promise<DocumentActionResponse> => {
+export const clearDocuments = async (): Promise<DocActionResponse> => {
   const response = await axiosInstance.delete('/documents')
   return response.data
 }
diff --git a/lightrag_webui/src/components/StatusCard.tsx b/lightrag_webui/src/components/StatusCard.tsx
index ea038c74..3084d103 100644
--- a/lightrag_webui/src/components/StatusCard.tsx
+++ b/lightrag_webui/src/components/StatusCard.tsx
@@ -14,8 +14,6 @@ const StatusCard = ({ status }: { status: LightragStatus | null }) => {
           <span className="truncate">{status.working_directory}</span>
           <span>Input Directory:</span>
           <span className="truncate">{status.input_directory}</span>
-          <span>Indexed Files:</span>
-          <span>{status.indexed_files_count}</span>
         </div>
       </div>
 
diff --git a/lightrag_webui/src/components/documents/UploadDocumentsDialog.tsx b/lightrag_webui/src/components/documents/UploadDocumentsDialog.tsx
index 48edfd31..181b52f1 100644
--- a/lightrag_webui/src/components/documents/UploadDocumentsDialog.tsx
+++ b/lightrag_webui/src/components/documents/UploadDocumentsDialog.tsx
@@ -49,10 +49,10 @@ export default function UploadDocumentsDialog() {
         toast.error('Upload Failed\n' + errorMessage(err))
       } finally {
         setIsUploading(false)
-        setOpen(false)
+        // setOpen(false)
       }
     },
-    [setIsUploading, setProgresses, setOpen]
+    [setIsUploading, setProgresses]
   )
 
   return (
@@ -66,7 +66,7 @@ export default function UploadDocumentsDialog() {
       }}
     >
       <DialogTrigger asChild>
-        <Button variant="default" side="bottom" tooltip='Upload documents' size="sm">
+        <Button variant="default" side="bottom" tooltip="Upload documents" size="sm">
           <UploadIcon /> Upload
         </Button>
       </DialogTrigger>
diff --git a/lightrag_webui/src/features/DocumentManager.tsx b/lightrag_webui/src/features/DocumentManager.tsx
index 40185a37..0aa068d4 100644
--- a/lightrag_webui/src/features/DocumentManager.tsx
+++ b/lightrag_webui/src/features/DocumentManager.tsx
@@ -9,39 +9,35 @@ import {
   TableRow
 } from '@/components/ui/Table'
 import { Card, CardHeader, CardTitle, CardContent, CardDescription } from '@/components/ui/Card'
-import Progress from '@/components/ui/Progress'
 import EmptyCard from '@/components/ui/EmptyCard'
+import Text from '@/components/ui/Text'
 import UploadDocumentsDialog from '@/components/documents/UploadDocumentsDialog'
 import ClearDocumentsDialog from '@/components/documents/ClearDocumentsDialog'
 
-import {
-  getDocuments,
-  // getDocumentsScanProgress,
-  scanNewDocuments
-  // LightragDocumentsScanProgress
-} from '@/api/lightrag'
+import { getDocuments, scanNewDocuments, DocsStatusesResponse } from '@/api/lightrag'
 import { errorMessage } from '@/lib/utils'
 import { toast } from 'sonner'
-// import { useBackendState } from '@/stores/state'
+import { useBackendState } from '@/stores/state'
 
-import { RefreshCwIcon, TrashIcon } from 'lucide-react'
-
-// type DocumentStatus = 'indexed' | 'pending' | 'indexing' | 'error'
+import { RefreshCwIcon } from 'lucide-react'
 
 export default function DocumentManager() {
-  // const health = useBackendState.use.health()
-  const [files, setFiles] = useState<string[]>([])
-  const [indexedFiles, setIndexedFiles] = useState<string[]>([])
-  // const [scanProgress, setScanProgress] = useState<LightragDocumentsScanProgress | null>(null)
+  const health = useBackendState.use.health()
+  const [docs, setDocs] = useState<DocsStatusesResponse | null>(null)
 
   const fetchDocuments = useCallback(async () => {
     try {
       const docs = await getDocuments()
-      setFiles(docs)
+      if (docs && docs.statuses) {
+        setDocs(docs)
+        // console.log(docs)
+      } else {
+        setDocs(null)
+      }
     } catch (err) {
       toast.error('Failed to load documents\n' + errorMessage(err))
     }
-  }, [setFiles])
+  }, [setDocs])
 
   useEffect(() => {
     fetchDocuments()
@@ -56,28 +52,19 @@ export default function DocumentManager() {
     }
   }, [])
 
-  // useEffect(() => {
-  //   const interval = setInterval(async () => {
-  //     try {
-  //       if (!health) return
-  //       const progress = await getDocumentsScanProgress()
-  //       setScanProgress((pre) => {
-  //         if (pre?.is_scanning === progress.is_scanning && progress.is_scanning === false) {
-  //           return pre
-  //         }
-  //         return progress
-  //       })
-  //       console.log(progress)
-  //     } catch (err) {
-  //       toast.error('Failed to get scan progress\n' + errorMessage(err))
-  //     }
-  //   }, 2000)
-  //   return () => clearInterval(interval)
-  // }, [health])
-
-  const handleDelete = async (fileName: string) => {
-    console.log(`deleting ${fileName}`)
-  }
+  useEffect(() => {
+    const interval = setInterval(async () => {
+      if (!health) {
+        return
+      }
+      try {
+        await fetchDocuments()
+      } catch (err) {
+        toast.error('Failed to get scan progress\n' + errorMessage(err))
+      }
+    }, 5000)
+    return () => clearInterval(interval)
+  }, [health, fetchDocuments])
 
   return (
     <Card className="!size-full !rounded-none !border-none">
@@ -100,16 +87,6 @@ export default function DocumentManager() {
           <UploadDocumentsDialog />
         </div>
 
-        {/* {scanProgress?.is_scanning && (
-          <div className="space-y-2">
-            <div className="flex justify-between text-sm">
-              <span>Indexing {scanProgress.current_file}</span>
-              <span>{scanProgress.progress}%</span>
-            </div>
-            <Progress value={scanProgress.progress} />
-          </div>
-        )} */}
-
         <Card>
           <CardHeader>
             <CardTitle>Uploaded documents</CardTitle>
@@ -117,44 +94,67 @@ export default function DocumentManager() {
           </CardHeader>
 
           <CardContent>
-            {files.length == 0 && (
+            {!docs && (
               <EmptyCard
-                title="No documents uploades"
+                title="No documents uploaded"
                 description="upload documents to see them here"
               />
             )}
-            {files.length > 0 && (
+            {docs && (
               <Table>
                 <TableHeader>
                   <TableRow>
-                    <TableHead>Filename</TableHead>
+                    <TableHead>ID</TableHead>
+                    <TableHead>Summary</TableHead>
                     <TableHead>Status</TableHead>
-                    <TableHead>Actions</TableHead>
+                    <TableHead>Length</TableHead>
+                    <TableHead>Chunks</TableHead>
+                    <TableHead>Created</TableHead>
+                    <TableHead>Updated</TableHead>
+                    <TableHead>Metadata</TableHead>
                   </TableRow>
                 </TableHeader>
-                <TableBody>
-                  {files.map((file) => (
-                    <TableRow key={file}>
-                      <TableCell>{file}</TableCell>
-                      <TableCell>
-                        {indexedFiles.includes(file) ? (
-                          <span className="text-green-600">Indexed</span>
-                        ) : (
-                          <span className="text-yellow-600">Pending</span>
-                        )}
-                      </TableCell>
-                      <TableCell>
-                        <Button
-                          variant="ghost"
-                          size="sm"
-                          onClick={() => handleDelete(file)}
-                          // disabled={isUploading}
-                        >
-                          <TrashIcon />
-                        </Button>
-                      </TableCell>
-                    </TableRow>
-                  ))}
+                <TableBody className="text-sm">
+                  {Object.entries(docs.statuses).map(([status, documents]) =>
+                    documents.map((doc) => (
+                      <TableRow key={doc.id}>
+                        <TableCell className="truncate font-mono">{doc.id}</TableCell>
+                        <TableCell className="max-w-xs min-w-24 truncate">
+                          <Text
+                            text={doc.content_summary}
+                            tooltip={doc.content_summary}
+                            tooltipClassName="max-w-none overflow-visible block"
+                          />
+                        </TableCell>
+                        <TableCell>
+                          {status === 'processed' && (
+                            <span className="text-green-600">Completed</span>
+                          )}
+                          {status === 'processing' && (
+                            <span className="text-blue-600">Processing</span>
+                          )}
+                          {status === 'pending' && <span className="text-yellow-600">Pending</span>}
+                          {status === 'failed' && <span className="text-red-600">Failed</span>}
+                          {doc.error && (
+                            <span className="ml-2 text-red-500" title={doc.error}>
+                              ⚠️
+                            </span>
+                          )}
+                        </TableCell>
+                        <TableCell>{doc.content_length ?? '-'}</TableCell>
+                        <TableCell>{doc.chunks_count ?? '-'}</TableCell>
+                        <TableCell className="truncate">
+                          {new Date(doc.created_at).toLocaleString()}
+                        </TableCell>
+                        <TableCell className="truncate">
+                          {new Date(doc.updated_at).toLocaleString()}
+                        </TableCell>
+                        <TableCell className="max-w-xs truncate">
+                          {doc.metadata ? JSON.stringify(doc.metadata) : '-'}
+                        </TableCell>
+                      </TableRow>
+                    ))
+                  )}
                 </TableBody>
               </Table>
             )}