From d5109bf3c4d9eae3217dece406c6e410d39b6c8d Mon Sep 17 00:00:00 2001
From: Ivy <ivy@acorn.io>
Date: Fri, 6 Dec 2024 12:35:16 -0500
Subject: [PATCH] fix: waiting for modelProvider models to load before showing
 default model selection (#796)

* fix: waiting for modelProvider models to load before showing default models selection

chore: name instead of id

Update ModelProviderModels.tsx

* chore: minor fixes
---
 .../ModelProviderConfigure.tsx                | 44 +++++++++++++++++--
 .../model-providers/ModelProviderModels.tsx   | 26 ++++++-----
 .../model/shared/DefaultModelAliasForm.tsx    |  8 +++-
 3 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/ui/admin/app/components/model-providers/ModelProviderConfigure.tsx b/ui/admin/app/components/model-providers/ModelProviderConfigure.tsx
index 47ccc876d..725aa1eae 100644
--- a/ui/admin/app/components/model-providers/ModelProviderConfigure.tsx
+++ b/ui/admin/app/components/model-providers/ModelProviderConfigure.tsx
@@ -1,4 +1,4 @@
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import useSWR from "swr";
 
 import { ModelProvider, ModelProviderConfig } from "~/lib/model/modelProviders";
@@ -29,6 +29,32 @@ export function ModelProviderConfigure({
     const [showDefaultModelAliasForm, setShowDefaultModelAliasForm] =
         useState(false);
 
+    const [loadingModelProviderId, setLoadingModelProviderId] = useState("");
+
+    const getLoadingModelProviderModels = useSWR(
+        ModelProviderApiService.getModelProviderById.key(
+            loadingModelProviderId
+        ),
+        ({ modelProviderId }) =>
+            ModelProviderApiService.getModelProviderById(modelProviderId),
+        {
+            revalidateOnFocus: false,
+            refreshInterval: 2000,
+        }
+    );
+
+    useEffect(() => {
+        if (!loadingModelProviderId) return;
+
+        const { isLoading, data } = getLoadingModelProviderModels;
+        if (isLoading) return;
+
+        if (data?.modelsBackPopulated) {
+            setShowDefaultModelAliasForm(true);
+            setLoadingModelProviderId("");
+        }
+    }, [getLoadingModelProviderModels, loadingModelProviderId]);
+
     const handleDone = () => {
         setDialogIsOpen(false);
         setShowDefaultModelAliasForm(false);
@@ -49,8 +75,16 @@ export function ModelProviderConfigure({
                 Configure Model Provider
             </DialogDescription>
 
-            <DialogContent className="p-0 gap-0">
-                {showDefaultModelAliasForm ? (
+            <DialogContent
+                className="p-0 gap-0"
+                hideCloseButton={loadingModelProviderId !== ""}
+            >
+                {loadingModelProviderId ? (
+                    <div className="flex items-center justify-center gap-1 p-2">
+                        <LoadingSpinner /> Loading {modelProvider.name}{" "}
+                        Models...
+                    </div>
+                ) : showDefaultModelAliasForm ? (
                     <div className="p-6">
                         <DialogHeader>
                             <DialogTitle className="flex items-center gap-2 pb-4">
@@ -62,7 +96,9 @@ export function ModelProviderConfigure({
                 ) : (
                     <ModelProviderConfigureContent
                         modelProvider={modelProvider}
-                        onSuccess={() => setShowDefaultModelAliasForm(true)}
+                        onSuccess={() =>
+                            setLoadingModelProviderId(modelProvider.id)
+                        }
                     />
                 )}
             </DialogContent>
diff --git a/ui/admin/app/components/model-providers/ModelProviderModels.tsx b/ui/admin/app/components/model-providers/ModelProviderModels.tsx
index 28f24ddb5..b29e2ffe7 100644
--- a/ui/admin/app/components/model-providers/ModelProviderModels.tsx
+++ b/ui/admin/app/components/model-providers/ModelProviderModels.tsx
@@ -75,7 +75,7 @@ export function ModelProvidersModels({ modelProvider }: ModelsConfigureProps) {
                     <DataTable
                         columns={getColumns()}
                         data={models}
-                        sort={[{ id: "id", desc: true }]}
+                        sort={[{ id: "usage", desc: true }]}
                     />
                 </ScrollArea>
             </DialogContent>
@@ -84,19 +84,21 @@ export function ModelProvidersModels({ modelProvider }: ModelsConfigureProps) {
 
     function getColumns(): ColumnDef<Model, string>[] {
         return [
-            columnHelper.accessor((model) => model.id, {
-                id: "id",
+            columnHelper.accessor((model) => model.name, {
+                id: "name",
                 header: "Model",
             }),
-            columnHelper.display({
-                id: "usage",
-                header: "Usage",
-                cell: ({ row }) => (
-                    <Badge variant="outline">
-                        {getModelUsageLabel(row.original.usage)}
-                    </Badge>
-                ),
-            }),
+            columnHelper.accessor(
+                (model) => getModelUsageLabel(model.usage) || "",
+                {
+                    id: "usage",
+                    header: "Usage",
+                    cell: ({ getValue }) =>
+                        getValue() ? (
+                            <Badge variant="outline">{getValue()}</Badge>
+                        ) : null,
+                }
+            ),
             columnHelper.display({
                 id: "actions",
                 cell: ({ row }) => (
diff --git a/ui/admin/app/components/model/shared/DefaultModelAliasForm.tsx b/ui/admin/app/components/model/shared/DefaultModelAliasForm.tsx
index 2592bd02d..3c21843e0 100644
--- a/ui/admin/app/components/model/shared/DefaultModelAliasForm.tsx
+++ b/ui/admin/app/components/model/shared/DefaultModelAliasForm.tsx
@@ -14,6 +14,7 @@ import {
 import { DefaultModelAliasApiService } from "~/lib/service/api/defaultModelAliasApiService";
 import { ModelApiService } from "~/lib/service/api/modelApiService";
 
+import { TypographyP } from "~/components/Typography";
 import { Button } from "~/components/ui/button";
 import {
     Dialog,
@@ -206,7 +207,12 @@ export function DefaultModelAliasForm({
         defaultModel: string
     ) {
         if (!modelOptions) {
-            if (!defaultModel) return null;
+            if (!defaultModel)
+                return (
+                    <TypographyP className="p-2 text-muted-foreground">
+                        No Models Available.
+                    </TypographyP>
+                );
             return <SelectItem value={defaultModel}>{defaultModel}</SelectItem>;
         }