devflowinc · skeptrunedev · Jun 28, 2024 · Jun 27, 2024 · Jun 28, 2024 · Jun 28, 2024
diff --git a/dashboard/src/components/DatasetOverview.tsx b/dashboard/src/components/DatasetOverview.tsx
@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/no-unsafe-member-access */
 import { TbDatabasePlus } from "solid-icons/tb";
 import {
   Show,
@@ -18,6 +19,8 @@ import {
 } from "solid-icons/ai";
 import { formatDate } from "../formatters";
 import { Organization } from "../types/apiTypes";
+import { TbReload } from "solid-icons/tb";
+import { createToast } from "./ShowToasts";
 
 export interface DatasetOverviewProps {
   setOpenNewDatasetModal: Setter<boolean>;
@@ -28,7 +31,9 @@ export const DatasetOverview = (props: DatasetOverviewProps) => {
   const navigate = useNavigate();
   const [page, setPage] = createSignal(0);
   const [datasetSearchQuery, setDatasetSearchQuery] = createSignal("");
-
+  const [usage, setUsage] = createSignal<
+    Record<string, { chunk_count: number }>
+  >({});
   const { datasets, maxPageDiscovered, maxDatasets, removeDataset, hasLoaded } =
     useDatasetPages({
       // eslint-disable-next-line solid/reactivity
@@ -38,6 +43,57 @@ export const DatasetOverview = (props: DatasetOverviewProps) => {
       setPage,
     });
 
+  createEffect(() => {
+    const api_host = import.meta.env.VITE_API_HOST as unknown as string;
+    const newUsage: Record<string, { chunk_count: number }> = {};
+
+    const fetchUsage = (datasetId: string) => {
+      return new Promise((resolve, reject) => {
+        fetch(`${api_host}/dataset/usage/${datasetId}`, {
+          method: "GET",
+          headers: {
+            "TR-Dataset": datasetId,
+            "Content-Type": "application/json",
+          },
+          credentials: "include",
+        })
+          .then((response) => {
+            if (!response.ok) {
+              reject(new Error("Failed to fetch dataset usage"));
+            }
+            return response.json();
+          })
+          .then((data) => {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            newUsage[datasetId] = data;
+            resolve(data);
+          })
+          .catch((error) => {
+            console.error(
+              `Failed to fetch usage for dataset ${datasetId}:`,
+              error,
+            );
+            reject(error);
+          });
+      });
+    };
+
+    const fetchInitialUsage = () => {
+      const promises = datasets().map((dataset) =>
+        fetchUsage(dataset.dataset.id),
+      );
+      Promise.all(promises)
+        .then(() => {
+          setUsage(newUsage);
+        })
+        .catch((error) => {
+          console.error("Failed to fetch initial usage: ", error);
+        });
+    };
+
+    fetchInitialUsage();
+  });
+
   createEffect(() => {
     props.selectedOrganization();
     setPage(0);
@@ -69,6 +125,64 @@ export const DatasetOverview = (props: DatasetOverviewProps) => {
     });
   };
 
+  const reloadChunkCount = (datasetId: string) => {
+    const api_host = import.meta.env.VITE_API_HOST as unknown as string;
+    if (!datasetId) {
+      console.error("Dataset ID is undefined.");
+      return;
+    }
+
+    const currentUsage = usage();
+
+    fetch(`${api_host}/dataset/usage/${datasetId}`, {
+      method: "GET",
+      headers: {
+        "TR-Dataset": datasetId,
+        "Content-Type": "application/json",
+      },
+      credentials: "include",
+    })
+      .then((response) => {
+        if (!response.ok) {
+          throw new Error("Failed to fetch dataset usage");
+        }
+        return response.json();
+      })
+      .then((newData) => {
+        const prevCount = currentUsage[datasetId]?.chunk_count || 0;
+        const newCount: number = newData.chunk_count as number;
+        const countDifference = newCount - prevCount;
+
+        setUsage((prevUsage) => ({
+          ...prevUsage,
+          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+          [datasetId]: newData,
+        }));
+
+        createToast({
+          title: "Updated",
+          type: "success",
+          message: `Successfully updated chunk count: ${countDifference} chunk${
+            Math.abs(countDifference) === 1 ? " has" : "s have"
+          } been ${
+            countDifference > 0
+              ? "added"
+              : countDifference < 0
+                ? "removed"
+                : "added or removed"
+          } since last update.`,
+          timeout: 3000,
+        });
+      })
+      .catch((error) => {
+        createToast({
+          title: "Error",
+          type: "error",
+          message: `Failed to reload chunk count: ${error}`,
+        });
+      });
+  };
+
   return (
     <>
       <div class="flex items-center">
@@ -181,7 +295,22 @@ export const DatasetOverview = (props: DatasetOverviewProps) => {
                           );
                         }}
                       >
-                        {datasetAndUsage.dataset_usage.chunk_count}
+                        <span class="inline-flex items-center">
+                          <div>
+                            {usage()[datasetAndUsage.dataset.id]?.chunk_count ??
+                              datasetAndUsage.dataset_usage.chunk_count}{" "}
+                          </div>
+                          <button
+                             => {
+                              e.stopPropagation();
+                              e.preventDefault();
+                              reloadChunkCount(datasetAndUsage.dataset.id);
+                            }}
+                            class="ml-2 hover:text-fuchsia-500"
+                          >
+                            <TbReload />
+                          </button>
+                        </span>
                       </td>
                       <td
                         class="hidden whitespace-nowrap px-3 py-4 text-sm text-neutral-600 lg:block"

diff --git a/dashboard/src/pages/Dashboard/Dataset/DatasetStart.tsx b/dashboard/src/pages/Dashboard/Dataset/DatasetStart.tsx
@@ -1,4 +1,3 @@
-import { BiRegularInfoCircle, BiRegularLinkExternal } from "solid-icons/bi";
 import CreateChunkRequest from "../../../components/CreateChunkRequest.md";
 import HybridSearchReqeust from "../../../components/HybridSearchRequest.md";
 import { BuildingSomething } from "../../../components/BuildingSomething";
@@ -21,8 +20,10 @@ import {
 import { DatasetContext } from "../../../contexts/DatasetContext";
 import { FaRegularClipboard } from "solid-icons/fa";
 import { AiOutlineInfoCircle } from "solid-icons/ai";
-import { AddSampleDataModal } from "../../../components/DatasetExampleModal";
+import { TbReload } from "solid-icons/tb";
+import { BiRegularInfoCircle, BiRegularLinkExternal } from "solid-icons/bi";
 import { BsMagic } from "solid-icons/bs";
+import { AddSampleDataModal } from "../../../components/DatasetExampleModal";
 import { Tooltip } from "../../../components/Tooltip";
 
 const SAMPLE_DATASET_SIZE = 921;
@@ -63,7 +64,6 @@ export const DatasetStart = () => {
 
       if (response.ok) {
         const data = (await response.json()) as unknown as DatasetUsageCount;
-        console.log("Got good repsonse", data);
         return data;
       } else {
         createToast({
@@ -77,6 +77,60 @@ export const DatasetStart = () => {
     },
   );
 
+  const reloadChunkCount = () => {
+    const datasetId = currDatasetId();
+    if (!datasetId) {
+      console.error("Dataset ID is undefined.");
+      return;
+    }
+
+    fetch(`${api_host}/dataset/usage/${datasetId}`, {
+      method: "GET",
+      headers: {
+        "TR-Dataset": datasetId,
+        "Content-Type": "application/json",
+      },
+      credentials: "include",
+    })
+      .then((response) => {
+        if (!response.ok) {
+          throw new Error("Failed to fetch dataset usage");
+        }
+        return response.json();
+      })
+      .then((newData) => {
+        const currentUsage = usage();
+        const prevCount = currentUsage?.chunk_count || 0;
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+        const newCount: number = newData.chunk_count as number;
+        const countDifference = newCount - prevCount;
+
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+        mutateUsage(() => newData);
+        createToast({
+          title: "Updated",
+          type: "success",
+          message: `Successfully updated chunk count: ${countDifference} chunk${
+            Math.abs(countDifference) === 1 ? " has" : "s have"
+          } been ${
+            countDifference > 0
+              ? "added"
+              : countDifference < 0
+                ? "removed"
+                : "added or removed"
+          } since last update.`,
+          timeout: 3000,
+        });
+      })
+      .catch((error) => {
+        createToast({
+          title: "Error",
+          type: "error",
+          message: `Failed to reload chunk count: ${error}`,
+        });
+      });
+  };
+
   const [updatedTrackingId, setUpdatedTrackingId] = createSignal<
     string | undefined
   >(datasetContext.dataset?.()?.tracking_id);
@@ -241,6 +295,12 @@ export const DatasetStart = () => {
               <div class="flex items-center space-x-3">
                 <p class="block text-sm font-medium">Chunk Count:</p>
                 <p class="w-fit text-sm">{usage()?.chunk_count || 0}</p>
+                <button
+                  >
+                  class="hover:text-fuchsia-500"
+                >
+                  <TbReload />
+                </button>
               </div>
               <div class="flex items-center space-x-3">
                 <label class="block text-sm font-medium">tracking id:</label>