feature: show queue length in dashboard

vid277 · cdxker · commit 3837faeb52a2 · 2025-05-29T12:48:46.000-07:00
diff --git a/clients/ts-sdk/openapi.json b/clients/ts-sdk/openapi.json
@@ -4239,6 +4239,57 @@
         ]
       }
     },
+    "/api/dataset/get_dataset_queue_lengths": {
+      "get": {
+        "tags": [
+          "Dataset"
+        ],
+        "summary": "Get file and chunk creation queue lengths",
+        "description": "Get the queue lengths for file and chunk creation.",
+        "operationId": "get_dataset_queue_lengths",
+        "parameters": [
+          {
+            "name": "TR-Dataset",
+            "in": "header",
+            "description": "The dataset id or tracking_id to use for the request. We assume you intend to use an id if the value is a valid uuid.",
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid"
+            }
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Queue lengths for file and chunk creation",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/DatasetQueueLengthsResponse"
+                }
+              }
+            }
+          },
+          "400": {
+            "description": "Service error relating to getting the queue lengths",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ErrorResponseBody"
+                }
+              }
+            }
+          }
+        },
+        "security": [
+          {
+            "ApiKey": [
+              "readonly"
+            ]
+          }
+        ]
+      }
+    },
     "/api/dataset/groups/{dataset_id}/{page}": {
       "get": {
         "tags": [
@@ -12391,6 +12442,25 @@
           }
         }
       },
+      "DatasetQueueLengthsResponse": {
+        "type": "object",
+        "required": [
+          "file_queue_length",
+          "chunk_queue_length"
+        ],
+        "properties": {
+          "chunk_queue_length": {
+            "type": "integer",
+            "format": "int64",
+            "description": "Number of chunks in the queue for the dataset"
+          },
+          "file_queue_length": {
+            "type": "integer",
+            "format": "int64",
+            "description": "Number of files in the queue for the dataset"
+          }
+        }
+      },
       "DatasetUsageCount": {
         "type": "object",
         "required": [
diff --git a/clients/ts-sdk/src/functions/datasets/index.ts b/clients/ts-sdk/src/functions/datasets/index.ts
@@ -10,6 +10,7 @@ import {
   CreateDatasetReqPayload,
   Dataset,
   DatasetAndUsage,
+  DatasetQueueLengthsResponse,
   DatasetUsageCount,
   EventReturn,
   FileData,
@@ -300,3 +301,19 @@ export async function getPagefindUrl(
     signal
   ) as Promise<GetPagefindIndexResponse>;
 }
+
+export async function getDatasetQueueLengths(
+  /** @hidden */
+  this: TrieveSDK,
+  datasetId: string,
+  signal?: AbortSignal
+): Promise<DatasetQueueLengthsResponse> {
+  return this.trieve.fetch(
+    "/api/dataset/get_dataset_queue_lengths",
+    "get",
+    {
+      datasetId,
+    },
+    signal
+  ) as Promise<DatasetQueueLengthsResponse>;
+}
diff --git a/clients/ts-sdk/src/types.gen.ts b/clients/ts-sdk/src/types.gen.ts
@@ -1470,6 +1470,17 @@ export type DatasetFilePathParams = {
     page: number;
 };
 
+export type DatasetQueueLengthsResponse = {
+    /**
+     * Number of chunks in the queue for the dataset
+     */
+    chunk_queue_length: number;
+    /**
+     * Number of files in the queue for the dataset
+     */
+    file_queue_length: number;
+};
+
 export type DatasetUsageCount = {
     chunk_count: number;
     dataset_id: string;
@@ -6012,6 +6023,15 @@ export type GetAllTagsData = {
 
 export type GetAllTagsResponse2 = (GetAllTagsResponse);
 
+export type GetDatasetQueueLengthsData = {
+    /**
+     * The dataset id or tracking_id to use for the request. We assume you intend to use an id if the value is a valid uuid.
+     */
+    trDataset: string;
+};
+
+export type GetDatasetQueueLengthsResponse = (DatasetQueueLengthsResponse);
+
 export type GetGroupsForDatasetData = {
     /**
      * The cursor offset for. Requires `use_cursor` = True. Defaults to `00000000-00000000-00000000-00000000`. Group ids are compared to the cursor using a greater than or equal to.
@@ -7811,6 +7831,21 @@ export type $OpenApiTs = {
             };
         };
     };
+    '/api/dataset/get_dataset_queue_lengths': {
+        get: {
+            req: GetDatasetQueueLengthsData;
+            res: {
+                /**
+                 * Queue lengths for file and chunk creation
+                 */
+                200: DatasetQueueLengthsResponse;
+                /**
+                 * Service error relating to getting the queue lengths
+                 */
+                400: ErrorResponseBody;
+            };
+        };
+    };
     '/api/dataset/groups/{dataset_id}/{page}': {
         get: {
             req: GetGroupsForDatasetData;
diff --git a/frontends/dashboard/src/components/Sidebar.tsx b/frontends/dashboard/src/components/Sidebar.tsx
@@ -11,6 +11,7 @@ import {
   AiOutlineLeft,
   AiOutlineMessage,
   AiOutlineSearch,
+  AiOutlineReload,
 } from "solid-icons/ai";
 import { Spacer } from "./Spacer";
 import { Portal } from "solid-js/web";
@@ -24,14 +25,19 @@ import { TbSparkles, TbTimelineEventText, TbTransform } from "solid-icons/tb";
 import { createSignal } from "solid-js";
 import NewDatasetModal from "../components/NewDatasetModal";
 import { ImNewspaper } from "solid-icons/im";
+import { useTrieve } from "../hooks/useTrieve";
+import { createQuery } from "@tanstack/solid-query";
 
-const searchUiURL = import.meta.env.VITE_SEARCH_UI_URL as string;
-const chatUiURL = import.meta.env.VITE_CHAT_UI_URL as string;
+const searchUiURL =
+  (import.meta.env.VITE_SEARCH_UI_URL as string | undefined) ?? "";
+const chatUiURL =
+  (import.meta.env.VITE_CHAT_UI_URL as string | undefined) ?? "";
 
 export const DashboardSidebar = () => {
   const { datasetId } = useContext(DatasetContext);
   const userContext = useContext(UserContext);
   const pathname = useLocation();
+  const trieve = useTrieve();
 
   const [newDatasetModalOpen, setNewDatasetModalOpen] =
     createSignal<boolean>(false);
@@ -63,6 +69,91 @@ export const DashboardSidebar = () => {
     );
   });
 
+  const uploadStatusQuery = createQuery(() => ({
+    queryKey: ["upload-status", datasetId()],
+    queryFn: async () => {
+      return await trieve.fetch(
+        "/api/dataset/get_dataset_queue_lengths",
+        "get",
+        {
+          datasetId: datasetId(),
+        },
+      );
+    },
+    refetchInterval: 5000,
+    refetchOnMount: true,
+    refetchOnWindowFocus: true,
+    enabled: !!datasetId(),
+  }));
+
+  const UploadStatusBar = () => {
+    return (
+      <div class="space-y-2 p-2">
+        <div class="flex flex-col justify-between">
+          <div class="text-sm font-medium text-neutral-700">
+            Processing Queue
+          </div>
+          <div class="flex items-center justify-between text-xs">
+            <span class="text-neutral-500">Updates every 5 seconds</span>
+            <button
+              onClick={() => {
+                void uploadStatusQuery.refetch();
+              }}
+              class="flex items-center gap-1 text-neutral-500 hover:text-fuchsia-500"
+              title="Reload status"
+            >
+              <AiOutlineReload
+                size={14}
+                classList={{
+                  "animate-spin": uploadStatusQuery.isFetching,
+                }}
+              />
+            </button>
+          </div>
+        </div>
+
+        <div class="space-y-1">
+          <div class="flex justify-between text-xs text-neutral-600">
+            <span>Files</span>
+            <span>{uploadStatusQuery.data?.file_queue_length ?? 0} files</span>
+          </div>
+          <div class="h-2 w-full rounded-full bg-neutral-200">
+            <div
+              class="h-2 rounded-full bg-blue-500 transition-all duration-300"
+              style={{
+                width: `${Math.min(
+                  ((uploadStatusQuery.data?.file_queue_length ?? 0) / 20) * 100,
+                  100,
+                )}%`,
+              }}
+            />
+          </div>
+        </div>
+
+        <div class="space-y-1">
+          <div class="flex justify-between text-xs text-neutral-600">
+            <span>Chunk Batches</span>
+            <span>
+              {uploadStatusQuery.data?.chunk_queue_length ?? 0} batches
+            </span>
+          </div>
+          <div class="h-2 w-full rounded-full bg-neutral-200">
+            <div
+              class="h-2 rounded-full bg-orange-500 transition-all duration-300"
+              style={{
+                width: `${Math.min(
+                  ((uploadStatusQuery.data?.chunk_queue_length ?? 0) / 1000) *
+                    100,
+                  100,
+                )}%`,
+              }}
+            />
+          </div>
+        </div>
+      </div>
+    );
+  };
+
   const Link = (props: {
     href: string;
     label: JSX.Element;
@@ -239,6 +330,10 @@ export const DashboardSidebar = () => {
               />
             </Show>
           </div>
+          <div class="gap flex flex-col pt-4">
+            <SectionLabel>Upload Status</SectionLabel>
+            <UploadStatusBar />
+          </div>
         </div>
       </div>
     </>
diff --git a/server/src/bin/file-worker.rs b/server/src/bin/file-worker.rs
@@ -11,7 +11,7 @@ use std::{
     sync::{atomic::AtomicBool, Arc},
 };
 use trieve_server::{
-    data::models::{self, ChunkGroup, FileWorkerMessage, Pool},
+    data::models::{self, ChunkGroup, FileWorkerMessage},
     establish_connection, get_env,
     handlers::chunk_handler::ChunkReqPayload,
     operators::{
@@ -796,4 +796,4 @@ async fn upload_file(
     .await?;
 
     Ok(None)
-}
+}
diff --git a/server/src/handlers/dataset_handler.rs b/server/src/handlers/dataset_handler.rs
@@ -8,6 +8,7 @@ use crate::{
     get_env,
     middleware::auth_middleware::{verify_admin, verify_owner},
     operators::{
+        chunk_operator::get_chunk_queue_length,
         dataset_operator::{
             clear_dataset_by_dataset_id_query, create_dataset_query, create_datasets_query,
             get_dataset_by_id_query, get_dataset_by_tracking_id_query, get_dataset_usage_query,
@@ -17,7 +18,6 @@ use crate::{
         dittofeed_operator::{
             send_ditto_event, DittoDatasetCreated, DittoTrackProperties, DittoTrackRequest,
         },
-        chunk_operator::get_chunk_queue_length,
         file_operator::get_file_queue_length,
         organization_operator::{get_org_dataset_count, get_org_from_id_query},
     },
@@ -1028,14 +1028,15 @@ pub async fn get_dataset_queue_lengths(
     broccoli_queue: web::Data<BroccoliQueue>,
     dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan,
 ) -> Result<HttpResponse, ServiceError> {
-    log::info!("Getting file queue length for dataset: {}", dataset_org_plan_sub.dataset.id);
-    let file_queue_length = get_file_queue_length(dataset_org_plan_sub.dataset.id, &broccoli_queue).await
-        .map_err(|e| ServiceError::InternalServerError(e.to_string()))?;
-    let chunk_queue_length = get_chunk_queue_length(dataset_org_plan_sub.dataset.id, &broccoli_queue).await
+    let file_queue_length = get_file_queue_length(dataset_org_plan_sub.dataset.id, &broccoli_queue)
+        .await
         .map_err(|e| ServiceError::InternalServerError(e.to_string()))?;
-    log::info!("file_queue_length: {}", file_queue_length);
+    let chunk_queue_length =
+        get_chunk_queue_length(dataset_org_plan_sub.dataset.id, &broccoli_queue)
+            .await
+            .map_err(|e| ServiceError::InternalServerError(e.to_string()))?;
     Ok(HttpResponse::Ok().json(DatasetQueueLengthsResponse {
         file_queue_length,
         chunk_queue_length,
     }))
-}
+}
diff --git a/server/src/lib.rs b/server/src/lib.rs
@@ -261,6 +261,7 @@ impl Modify for SecurityAddon {
         handlers::dataset_handler::get_pagefind_index_for_dataset,
         handlers::dataset_handler::clear_dataset,
         handlers::dataset_handler::clone_dataset,
+        handlers::dataset_handler::get_dataset_queue_lengths,
         handlers::payment_handler::direct_to_payment_link,
         handlers::payment_handler::cancel_subscription,
         handlers::payment_handler::update_subscription_plan,
diff --git a/server/src/operators/chunk_operator.rs b/server/src/operators/chunk_operator.rs
@@ -10,12 +10,12 @@ use crate::operators::parse_operator::convert_html_to_text;
 use crate::operators::qdrant_operator::{
     delete_points_from_qdrant, get_qdrant_collection_from_dataset_config, scroll_dataset_points,
 };
-use broccoli_queue::queue::BroccoliQueue;
 use crate::{
     data::models::{ChunkMetadata, Pool},
     errors::ServiceError,
 };
 use actix_web::web;
+use broccoli_queue::queue::BroccoliQueue;
 use chrono::NaiveDateTime;
 use clickhouse::Row;
 use dateparser::DateTimeUtc;
@@ -2972,9 +2972,14 @@ pub async fn get_chunk_queue_length(
         .map_err(|e| ServiceError::InternalServerError(e.to_string()))?;
 
     let premium_ingestion_queue_status = broccoli_queue
-        .queue_status("premium_ingestion".to_string(), Some(dataset_id.to_string()))
+        .queue_status(
+            "premium_ingestion".to_string(),
+            Some(dataset_id.to_string()),
+        )
         .await
         .map_err(|e| ServiceError::InternalServerError(e.to_string()))?;
 
-    Ok(openai_ingestion_queue_status.size as i64 + regular_ingestion_queue_status.size as i64 + premium_ingestion_queue_status.size as i64)
+    Ok(openai_ingestion_queue_status.size as i64
+        + regular_ingestion_queue_status.size as i64
+        + premium_ingestion_queue_status.size as i64)
 }
diff --git a/server/src/operators/file_operator.rs b/server/src/operators/file_operator.rs
diff --git a/yarn.lock b/yarn.lock
