From 9660038095a403d2f0917182512ed6b6d0fa28ff Mon Sep 17 00:00:00 2001
From: wwayne <wayne.wang0821@gmail.com>
Date: Fri, 26 Apr 2024 20:20:51 +0800
Subject: [PATCH] update

---
 .../cluster/components/cluster.tsx            | 214 ++++++++++++------
 ee/tabby-ui/components/ui/icons.tsx           |   9 +-
 2 files changed, 142 insertions(+), 81 deletions(-)

diff --git a/ee/tabby-ui/app/(dashboard)/cluster/components/cluster.tsx b/ee/tabby-ui/app/(dashboard)/cluster/components/cluster.tsx
index 20b74ef79ca0..50cba0d927aa 100644
--- a/ee/tabby-ui/app/(dashboard)/cluster/components/cluster.tsx
+++ b/ee/tabby-ui/app/(dashboard)/cluster/components/cluster.tsx
@@ -4,9 +4,10 @@ import { noop } from 'lodash-es'
 import { useQuery } from 'urql'
 import bytes from 'bytes';
 import { useTheme } from 'next-themes';
+import {sum} from 'lodash-es';
 
 import { graphql } from '@/lib/gql/generates'
-import { WorkerKind } from '@/lib/gql/generates/graphql'
+import { WorkerKind, DiskUsageStats, DiskUsage } from '@/lib/gql/generates/graphql'
 import { useHealth } from '@/lib/hooks/use-health'
 import { useWorkers } from '@/lib/hooks/use-workers'
 import { useMutation } from '@/lib/tabby/gql'
@@ -22,12 +23,7 @@ import {
   Cell,
   Label,
 } from 'recharts'
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from "@/components/ui/tooltip"
-import { IconFolderOpen } from '@/components/ui/icons';
+import LoadingWrapper from '@/components/loading-wrapper';
 
 import WorkerCard from './worker-card'
 
@@ -62,15 +58,6 @@ export default function Workers() {
 
   if (!healthInfo) return
 
-  // TODO: mock data
-  const usageData = [
-    { name: "Models", usage: 1024 * 1024 * 1024 * 1.1, color: "#0088FE", folders: ["~/.tabby/models"] },
-    { name: "Repository data", usage: 1024 * 1024 * 2.1, color: "#00C49F", folders: ["~/.tabby/dataset", "~/.tabby/index"] },
-    { name: "Database", usage: 1024 * 85, color: "#FFBB28", folders: ["~/.tabby/ee"] },
-    { name: "Events", usage: 1024 * 1200 * 11, color: "#FF8042", folders: ["~/.tabby/events"] }
-  ];
-  const totalUsage = 1024 * 1300 * 1200
-
   return (
     <div className="flex w-full flex-col gap-3">
       <h1>
@@ -137,64 +124,143 @@ export default function Workers() {
         />
       </div>
 
-      <Separator className="mt-5" />
+      <Separator className="mt-6" />
+      <Usage />
+    </div>
+  )
+}
 
-      <div>
-        <p className="font-bold">
-          Storage Usage
-        </p>
-        <p className="text-sm text-muted-foreground">Disk space used by Tabby</p>
-      </div>
-      
-      
-      <div className="flex flex-col items-center gap-x-3 md:flex-row">
-        <ResponsiveContainer width={230} height={220}>
-          <PieChart>
-            <Pie 
-              data={usageData} 
-              dataKey="usage"
-              cx={110}
-              cy={100}
-              innerRadius={70}
-              outerRadius={90}
-              stroke='none'
-            >
-              {usageData.map((entry) => (
-                <Cell key={entry.name} fill={entry.color} />
-              ))}
-              <Label  
-                content={<CustomLabel totalUsage={totalUsage} />}
-                position="center"
-              />
-            </Pie>
-          </PieChart>
-        </ResponsiveContainer>
-
-        <div className="flex w-full flex-col gap-y-2 md:ml-10 md:w-auto">
-          {usageData.map(item => (
-            <Tooltip key={item.name}>
-              <TooltipTrigger>
-                <div className="flex cursor-default items-center justify-between text-xs" >
-                  <div className="flex w-40 items-center">
-                    <div className="mr-1.5 h-3 w-3 rounded" style={{ backgroundColor: item.color }} />
-                    <p className="font-semibold">{item.name}</p>
+export const getDiskUsageStats = graphql(/* GraphQL */ `
+  query GetDiskUsageStats{
+    diskUsageStats {
+      events {
+        filePaths
+        size
+      }
+      indexedRepositories {
+        filePaths
+        size
+      }
+      database {
+        filePaths
+        size
+      }
+      models {
+        filePaths
+        size
+      }
+    }
+  }
+`)
+
+type UsageItem = {
+  label: string;
+  key: keyof DiskUsageStats;
+  color: string
+}
+
+const usageList: UsageItem[] = [
+  {
+    label: "Model Cache",
+    key: 'models',
+    color: "#0088FE"
+  },
+  {
+    label: "Repository Context",
+    key: 'indexedRepositories',
+    color: "#00C49F"
+  },
+  {
+    label: "Event Logs",
+    key: 'events',
+    color: "#FF8042"
+  },
+  {
+    label: "Other",
+    key: 'database',
+    color: "#FFBB28"
+  },
+]
+
+function Usage () {
+  const [{ data, fetching }] = useQuery({
+    query: getDiskUsageStats
+  })
+  
+  let usageData: Array<UsageItem & { size: number }> = []
+  let totalUsage: number = 0
+  if (data) {
+    usageData = usageList
+      .map(usage => {
+        if (!data.diskUsageStats[usage.key]) return null
+        const diskUsage = data.diskUsageStats[usage.key] as DiskUsage
+        return {
+          ...usage,
+          size: diskUsage.size
+        }
+      })
+      .filter(usage => usage) as Array<UsageItem & { size: number }>
+      totalUsage = sum(usageData.map(data => data.size))
+  }
+
+  return (
+      <LoadingWrapper
+        loading={fetching}
+        fallback={<></>}>
+        <>
+          <div>
+            <p className="font-bold">
+              Storage Usage
+            </p>
+            <p className="text-sm text-muted-foreground">Disk Space Utilization by Type</p>
+          </div>
+          <div className="flex flex-col items-center gap-x-3 md:flex-row">
+            <ResponsiveContainer width={230} height={220}>
+              <PieChart>
+                <Pie 
+                  data={usageData} 
+                  dataKey="size"
+                  cx={110}
+                  cy={100}
+                  innerRadius={70}
+                  outerRadius={90}
+                  stroke='none'
+                >
+                  {usageData.map((entry) => (
+                    <Cell key={entry.key} fill={entry.color} />
+                  ))}
+                  <Label  
+                    content={<CustomLabel totalUsage={totalUsage} />}
+                    position="center"
+                  />
+                </Pie>
+              </PieChart>
+            </ResponsiveContainer>
+
+            <div className="flex w-full flex-col gap-y-2 md:ml-10 md:w-auto">
+              {usageList
+                .map(usage => {
+                  if (!data?.diskUsageStats[usage.key]) return null
+                  const diskUsage = data?.diskUsageStats[usage.key] as DiskUsage
+                  return {
+                    ...usage,
+                    size: diskUsage.size
+                  }
+                })
+                .filter(usage => usage)
+                .map(usage => (
+                  <div className="flex cursor-default items-center justify-between text-xs" key={usage!.key}>
+                    <div className="flex w-40 items-center">
+                      <div className="mr-1.5 h-3 w-3 rounded" style={{ backgroundColor: usage!.color }} />
+                      <p className="font-semibold">{usage!.label}</p>
+                    </div>
+                    <p>{bytes(usage!.size * 1024)}</p>
                   </div>
-                  <p>{bytes(item.usage)}</p>
-                </div>
-              </TooltipTrigger>
-              <TooltipContent>
-                <div className="flex items-center text-muted-foreground">
-                  <IconFolderOpen className="mr-0.5" />
-                  <p>Folder</p>
-                </div>
-                
-                {item.folders.map(folder => (<p>{folder}</p>))}
-              </TooltipContent>
-            </Tooltip>
-          ))}
-        </div>
-      </div>
-    </div>
+                ))}
+            </div>
+          </div>
+        </>
+      </LoadingWrapper>
   )
 }
 
@@ -222,7 +288,7 @@ function CustomLabel ({
         fill={theme === 'dark' ? '#FDFDFD' : '#030302'}
         className="text-sm"
       >
-        Total Usuage
+        Total Usage
       </text>
       <text
         x={cx}
@@ -233,7 +299,7 @@ function CustomLabel ({
         fill={theme === 'dark' ? '#FDFDFD' : '#030302'}
         className="text-lg font-semibold"
       >
-        {bytes(totalUsage)}
+        {bytes(totalUsage * 1024)}
       </text>
     </g>
   );
diff --git a/ee/tabby-ui/components/ui/icons.tsx b/ee/tabby-ui/components/ui/icons.tsx
index e5ce500f6dfa..a636692de970 100644
--- a/ee/tabby-ui/components/ui/icons.tsx
+++ b/ee/tabby-ui/components/ui/icons.tsx
@@ -2,7 +2,7 @@
 
 import * as React from 'react'
 // FIXME(wwayne): Review each icons and consider re-export from `lucide-react`.
-import { BookOpenText, ChevronsDownUp, Mail, FolderOpen } from 'lucide-react'
+import { BookOpenText, ChevronsDownUp, Mail } from 'lucide-react'
 
 import { cn } from '@/lib/utils'
 
@@ -1374,10 +1374,6 @@ function IconChevronsDownUp({
   )
 }
 
-function IconFolderOpen({ className, ...props }: React.ComponentProps<typeof FolderOpen>) {
-  return <FolderOpen className={cn('h4 w-4', className)} {...props} />
-}
-
 export {
   IconEdit,
   IconNextChat,
@@ -1448,6 +1444,5 @@ export {
   IconActivity,
   IconBookOpenText,
   IconMail,
-  IconChevronsDownUp,
-  IconFolderOpen
+  IconChevronsDownUp
 }