From 399650001644ff8757781b2cddc93555c87f5c8f Mon Sep 17 00:00:00 2001
From: Harish Mohan Raj <harish@airt.ai>
Date: Mon, 22 Apr 2024 11:49:11 +0530
Subject: [PATCH] WIP

---
 app/src/client/app/ModelsPage.tsx             | 68 +++++++++++++++----
 .../client/components/DynamicFormBuilder.tsx  | 19 +++++-
 .../client/components/ModelFormContainer.tsx  |  9 ++-
 app/src/client/interfaces/models.ts           |  8 +++
 .../client/tests/DynamicFormBuilder.test.tsx  | 12 ++++
 app/src/server/queries.ts                     |  1 +
 6 files changed, 97 insertions(+), 20 deletions(-)

diff --git a/app/src/client/app/ModelsPage.tsx b/app/src/client/app/ModelsPage.tsx
index 285577c..0ffbdfa 100644
--- a/app/src/client/app/ModelsPage.tsx
+++ b/app/src/client/app/ModelsPage.tsx
@@ -1,4 +1,4 @@
-import { useState } from 'react';
+import { useState, useEffect } from 'react';
 import { getModels, useQuery, updateUserModels } from 'wasp/client/operations';
 
 import CustomLayout from './layout/CustomLayout';
@@ -10,18 +10,20 @@ import { getAvailableModels } from '../services/modelService';
 import { ModelSchema, JsonSchema } from '../interfaces/models';
 import ModelFormContainer from '../components/ModelFormContainer';
 import NotificationBox from '../components/NotificationBox';
+import { UpdateExistingModelType } from '../interfaces/models';
 
 const ModelsPage = () => {
   const [modelsSchema, setModelsSchema] = useState<ModelSchema | null>(null);
   const [initialModelSchema, setInitialModelSchema] = useState<JsonSchema | null>(null);
   const [selectedModel, setSelectedModel] = useState<string>('');
+  const [updateExistingModel, setUpdateExistingModel] = useState<UpdateExistingModelType | null>();
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [showAddModel, setShowAddModel] = useState(false);
-  const { data: modelsList, refetch: refetchModels } = useQuery(getModels);
+  const { data: modelsList, refetch: refetchModels, isLoading: getModelsIsLoading } = useQuery(getModels);
 
   const fetchData = async () => {
-    setShowAddModel(true);
+    // setShowAddModel(true);
     setIsLoading(true);
     try {
       const response = await getAvailableModels();
@@ -35,9 +37,9 @@ const ModelsPage = () => {
     setIsLoading(false);
   };
 
-  // useEffect(() => {
-  //   fetchData();
-  // }, []);
+  useEffect(() => {
+    fetchData();
+  }, []);
 
   const handleModelChange = (newModel: string) => {
     const foundSchema = modelsSchema?.schemas.find((schema) => schema.name === newModel);
@@ -57,6 +59,20 @@ const ModelsPage = () => {
     setShowAddModel(false);
   };
 
+  const updateSelectedModel = (index: number) => {
+    if (modelsList) {
+      const selectedModel = modelsList[index];
+      const selectedModelSchemaName = selectedModel.api_type === 'openai' ? 'openai' : 'azureoai';
+      const foundSchema = modelsSchema?.schemas.find((schema) => schema.name.toLowerCase() === selectedModelSchemaName);
+      if (foundSchema) {
+        setInitialModelSchema(foundSchema.json_schema);
+        setSelectedModel(foundSchema.name);
+        setUpdateExistingModel(selectedModel);
+        setShowAddModel(true);
+      }
+    }
+  };
+
   return (
     <CustomLayout>
       <CustomBreadcrumb pageName='Models' />
@@ -64,8 +80,16 @@ const ModelsPage = () => {
         <div className='flex flex-col gap-4'>
           <div className='rounded-sm border border-stroke bg-white shadow-default dark:border-strokedark dark:bg-boxdark min-h-[300px] sm:min-h-[600px]'>
             <div className='flex-col flex items-start p-6 gap-3 w-full'>
-              <div className='flex justify-end w-full px-6.5 py-3'>
-                <Button onClick={() => fetchData()} label='Add Model' />
+              <div className={`${showAddModel ? 'hidden' : ''} flex justify-end w-full px-6.5 py-3`}>
+                <Button
+                  onClick={async () => {
+                    setIsLoading(true);
+                    await fetchData();
+                    setShowAddModel(true);
+                    setIsLoading(false);
+                  }}
+                  label='Add Model'
+                />
               </div>
               <div className='flex-col flex w-full'>
                 {!showAddModel ? (
@@ -78,6 +102,8 @@ const ModelsPage = () => {
                             <div
                               key={i}
                               className='group relative cursor-pointer overflow-hidden bg-airt-primary text-airt-font-base px-6 pt-10 pb-8  transition-all duration-300 hover:-translate-y-1 sm:max-w-sm sm:rounded-lg sm:pl-8 sm:pr-24'
+                              // add a click event to the div
+                              onClick={() => updateSelectedModel(i)}
                             >
                               <span className='absolute top-10 z-0 h-9 w-9 rounded-full bg-airt-hero-gradient-start transition-all duration-300 group-hover:scale-[30]'></span>
                               <div className='relative z-10 mx-auto max-w-md'>
@@ -126,17 +152,28 @@ const ModelsPage = () => {
                   ) : (
                     <div className='flex flex-col gap-3'>
                       <h2 className='text-lg font-semibold text-airt-primary'>Available Models</h2>
-                      <p className='text-airt-primary/50 pt-2'>Please add a new model...</p>
+                      <p className='text-airt-primary/50 pt-2'>No models available. Please add one.</p>
                     </div>
                   )
                 ) : (
                   modelsSchema && (
                     <>
-                      <ModelFormContainer modelsSchema={modelsSchema} onModelChange={handleModelChange} />
+                      <ModelFormContainer
+                        selectedModel={
+                          updateExistingModel
+                            ? updateExistingModel.api_type === 'openai'
+                              ? 'openai'
+                              : 'azureoai'
+                            : null
+                        }
+                        modelsSchema={modelsSchema}
+                        onModelChange={handleModelChange}
+                      />
                       {initialModelSchema && (
                         <DynamicFormBuilder
                           jsonSchema={initialModelSchema}
                           validationURL={`models/llms/${selectedModel}/validate`}
+                          updateExistingModel={updateExistingModel ?? null}
                           onSuccessCallback={onSuccessCallback}
                           onCancelCallback={onCancelCallback}
                         />
@@ -149,11 +186,12 @@ const ModelsPage = () => {
             </div>
           </div>
         </div>
-        {isLoading && (
-          <div className='absolute inset-0 flex items-center justify-center bg-white bg-opacity-50'>
-            <Loader />
-          </div>
-        )}
+        {isLoading ||
+          (getModelsIsLoading && (
+            <div className='absolute inset-0 flex items-center justify-center bg-white bg-opacity-50'>
+              <Loader />
+            </div>
+          ))}
       </div>
     </CustomLayout>
   );
diff --git a/app/src/client/components/DynamicFormBuilder.tsx b/app/src/client/components/DynamicFormBuilder.tsx
index d9076d1..fc566b8 100644
--- a/app/src/client/components/DynamicFormBuilder.tsx
+++ b/app/src/client/components/DynamicFormBuilder.tsx
@@ -7,9 +7,12 @@ import { validateForm } from '../services/commonService';
 import { parseValidationErrors } from '../app/utils/formHelpers';
 import Loader from '../admin/common/Loader';
 
+import { UpdateExistingModelType } from '../interfaces/models';
+
 interface DynamicFormBuilderProps {
   jsonSchema: JsonSchema;
   validationURL: string;
+  updateExistingModel: UpdateExistingModelType | null;
   onSuccessCallback: (data: any) => void;
   onCancelCallback: (data: any) => void;
 }
@@ -17,6 +20,7 @@ interface DynamicFormBuilderProps {
 const DynamicFormBuilder: React.FC<DynamicFormBuilderProps> = ({
   jsonSchema,
   validationURL,
+  updateExistingModel,
   onSuccessCallback,
   onCancelCallback,
 }) => {
@@ -36,7 +40,6 @@ const DynamicFormBuilder: React.FC<DynamicFormBuilderProps> = ({
     }
     setIsLoading(false);
   };
-
   return (
     <>
       <form onSubmit={handleSubmit} className='grid grid-cols-1 md:grid-cols-2 gap-9 p-6.5'>
@@ -47,14 +50,24 @@ const DynamicFormBuilder: React.FC<DynamicFormBuilderProps> = ({
               {property.enum ? (
                 <SelectInput
                   id={key}
-                  value={formData[key]}
+                  value={
+                    updateExistingModel && updateExistingModel.hasOwnProperty(key)
+                      ? // @ts-ignore
+                        updateExistingModel[key]
+                      : formData[key]
+                  }
                   options={property.enum}
                   onChange={(value) => handleChange(key, value)}
                 />
               ) : (
                 <TextInput
                   id={key}
-                  value={formData[key]}
+                  value={
+                    updateExistingModel && updateExistingModel.hasOwnProperty(key)
+                      ? // @ts-ignore
+                        updateExistingModel[key]
+                      : formData[key]
+                  }
                   placeholder={property.description || ''}
                   onChange={(value) => handleChange(key, value)}
                 />
diff --git a/app/src/client/components/ModelFormContainer.tsx b/app/src/client/components/ModelFormContainer.tsx
index 10131a2..e73ef86 100644
--- a/app/src/client/components/ModelFormContainer.tsx
+++ b/app/src/client/components/ModelFormContainer.tsx
@@ -3,19 +3,24 @@ import React from 'react';
 import { ModelSchema } from '../interfaces/models';
 
 interface ModelFormContainerProps {
+  selectedModel: string | null;
   modelsSchema: ModelSchema;
   onModelChange: (selectedModel: string) => void;
 }
 
-const ModelFormContainer: React.FC<ModelFormContainerProps> = ({ modelsSchema, onModelChange }) => {
+const ModelFormContainer: React.FC<ModelFormContainerProps> = ({ selectedModel, modelsSchema, onModelChange }) => {
   return (
     <div className='flex flex-col gap-9'>
       <div className='flex flex-col gap-5.5 px-6.5'>
-        <label className='mb-3 block text-black dark:text-white'>Select Model</label>
+        <h2 className='text-lg font-semibold text-airt-primary mt-6 '>Add a new model</h2>
+        <label className='-mb-3 block text-black dark:text-white'>Select Model</label>
         <div className='relative z-20 bg-white dark:bg-form-input'>
           <select
             className='w-full rounded border bg-transparent py-3 px-12 outline-none focus:border-primary'
             onChange={(e) => onModelChange(e.target.value)}
+            defaultValue={
+              selectedModel ? (selectedModel === 'openai' ? 'OpenAI' : 'AzureOAI') : modelsSchema.schemas[0].name
+            }
           >
             {modelsSchema.schemas.map((schema) => (
               <option key={schema.name} value={schema.name}>
diff --git a/app/src/client/interfaces/models.ts b/app/src/client/interfaces/models.ts
index 54112bb..79dc102 100644
--- a/app/src/client/interfaces/models.ts
+++ b/app/src/client/interfaces/models.ts
@@ -26,3 +26,11 @@ export interface Schema {
 export interface ModelSchema {
   schemas: Schema[];
 }
+
+export type UpdateExistingModelType = {
+  model: string;
+  base_url: string;
+  api_type: string;
+  api_version?: string;
+  uuid: string;
+};
diff --git a/app/src/client/tests/DynamicFormBuilder.test.tsx b/app/src/client/tests/DynamicFormBuilder.test.tsx
index 85d3c16..7f2290b 100644
--- a/app/src/client/tests/DynamicFormBuilder.test.tsx
+++ b/app/src/client/tests/DynamicFormBuilder.test.tsx
@@ -6,6 +6,7 @@ import { renderInContext } from 'wasp/client/test';
 import DynamicFormBuilder from '../components/DynamicFormBuilder';
 import { JsonSchema } from '../interfaces/models';
 import { validateForm } from '../services/commonService';
+import { UpdateExistingModelType } from '../interfaces/models';
 
 const setFormErrors = vi.fn();
 const handleChange = vi.fn();
@@ -64,12 +65,20 @@ const jsonSchema: JsonSchema = {
   type: '',
 };
 
+const updateExistingModel: UpdateExistingModelType = {
+  model: 'gpt-3.5-turbo',
+  base_url: 'https://api.openai.com/v1',
+  api_type: 'openai',
+  uuid: '1234',
+};
+
 describe('DynamicFormBuilder', () => {
   test('renders form fields correctly', () => {
     renderInContext(
       <DynamicFormBuilder
         jsonSchema={jsonSchema}
         validationURL='https://some-domain/some-route'
+        updateExistingModel={updateExistingModel}
         onSuccessCallback={vi.fn()}
         onCancelCallback={vi.fn()}
       />
@@ -85,6 +94,7 @@ describe('DynamicFormBuilder', () => {
       <DynamicFormBuilder
         jsonSchema={jsonSchema}
         validationURL='https://some-domain/some-route'
+        updateExistingModel={updateExistingModel}
         onSuccessCallback={onSuccessCallback}
         onCancelCallback={vi.fn()}
       />
@@ -135,6 +145,7 @@ describe('DynamicFormBuilder', () => {
       <DynamicFormBuilder
         jsonSchema={jsonSchema}
         validationURL='https://some-domain/some-route'
+        updateExistingModel={updateExistingModel}
         onSuccessCallback={onSuccessCallback}
         onCancelCallback={vi.fn()}
       />
@@ -154,6 +165,7 @@ describe('DynamicFormBuilder', () => {
       <DynamicFormBuilder
         jsonSchema={jsonSchema}
         validationURL='https://some-domain/some-route'
+        updateExistingModel={updateExistingModel}
         onSuccessCallback={vi.fn()}
         onCancelCallback={vi.fn()}
       />
diff --git a/app/src/server/queries.ts b/app/src/server/queries.ts
index 2f08029..d24299a 100644
--- a/app/src/server/queries.ts
+++ b/app/src/server/queries.ts
@@ -111,6 +111,7 @@ type GetModelsValues = {
   base_url: string;
   api_type: string;
   api_version?: string;
+  uuid: string;
 };
 
 export const getModels: GetModels<void, GetModelsValues[]> = async (_args, context) => {