From e2f24fe3c49f540340c040e0c21f27656fed388f Mon Sep 17 00:00:00 2001
From: Harish Mohan Raj <harish@airt.ai>
Date: Thu, 18 Apr 2024 17:48:12 +0530
Subject: [PATCH] WIP

---
 app/main.wasp                           | 14 ++++++--
 app/src/client/app/ModelsPage.tsx       | 46 +++++++++++++++++--------
 app/src/client/services/modelService.ts |  6 ++--
 app/src/server/actions.ts               | 44 ++++++++++++++++++++---
 app/src/server/common/constants.ts      |  1 +
 app/src/server/queries.ts               | 32 ++++++++++++++++-
 6 files changed, 118 insertions(+), 25 deletions(-)
 create mode 100644 app/src/server/common/constants.ts

diff --git a/app/main.wasp b/app/main.wasp
index 25f83b0..3a9e709 100644
--- a/app/main.wasp
+++ b/app/main.wasp
@@ -265,8 +265,13 @@ action updateUserById {
   entities: [User]
 }
 
-action getModels {
-  fn: import { getModels } from "@src/server/actions.js",
+action getAvailableModels {
+  fn: import { getAvailableModels } from "@src/server/actions.js",
+  entities: []
+}
+
+action updateUserModels {
+  fn: import { updateUserModels } from "@src/server/actions.js",
   entities: []
 }
 
@@ -287,6 +292,11 @@ query getPaginatedUsers {
   entities: [User]
 }
 
+query getModels {
+  fn: import { getModels } from "@src/server/queries.js",
+  entities: [User]
+}
+
 /*
  * 📡 These are custom Wasp API Endpoints. Use them for callbacks, webhooks, etc.
  * https://wasp-lang.dev/docs/advanced/apis
diff --git a/app/src/client/app/ModelsPage.tsx b/app/src/client/app/ModelsPage.tsx
index c34cb6a..eda798f 100644
--- a/app/src/client/app/ModelsPage.tsx
+++ b/app/src/client/app/ModelsPage.tsx
@@ -1,10 +1,12 @@
 import { useState } from 'react';
+import { getModels, useQuery, updateUserModels } from 'wasp/client/operations';
+
 import CustomLayout from './layout/CustomLayout';
 import CustomBreadcrumb from '../components/CustomBreadcrumb';
 import Button from '../components/Button';
 import DynamicFormBuilder from '../components/DynamicFormBuilder';
 import Loader from '../admin/common/Loader';
-import { getModels } from '../services/modelService';
+import { getAvailableModels } from '../services/modelService';
 import { ModelSchema, JsonSchema } from '../interfaces/models';
 import ModelFormContainer from '../components/ModelFormContainer';
 import NotificationBox from '../components/NotificationBox';
@@ -15,11 +17,14 @@ const ModelsPage = () => {
   const [selectedModel, setSelectedModel] = useState<string>('');
   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 fetchData = async () => {
+    setShowAddModel(true);
     setIsLoading(true);
     try {
-      const response = await getModels();
+      const response = await getAvailableModels();
       setModelsSchema(response);
       setInitialModelSchema(response.schemas[0].json_schema);
       setSelectedModel(response.schemas[0].name);
@@ -42,8 +47,10 @@ const ModelsPage = () => {
     }
   };
 
-  const onSuccessCallback = (data: any) => {
-    console.log('Form submitted with data:', data);
+  const onSuccessCallback = async (data: any) => {
+    await updateUserModels({ data });
+    refetchModels();
+    setShowAddModel(false);
   };
 
   return (
@@ -57,17 +64,26 @@ const ModelsPage = () => {
                 <Button onClick={() => fetchData()} label='Add Model' />
               </div>
               <div className='flex-col flex w-full'>
-                {modelsSchema && (
-                  <>
-                    <ModelFormContainer modelsSchema={modelsSchema} onModelChange={handleModelChange} />
-                    {initialModelSchema && (
-                      <DynamicFormBuilder
-                        jsonSchema={initialModelSchema}
-                        validationURL={`models/llms/${selectedModel}/validate`}
-                        onSuccessCallback={onSuccessCallback}
-                      />
-                    )}
-                  </>
+                {!showAddModel ? (
+                  modelsList ? (
+                    <div className='flex flex-col gap-3'>
+                      <h2 className='text-lg font-semibold'>Available Models</h2>
+                      <div className='flex flex-col gap-2'>{JSON.stringify(modelsList)}</div>
+                    </div>
+                  ) : null
+                ) : (
+                  modelsSchema && (
+                    <>
+                      <ModelFormContainer modelsSchema={modelsSchema} onModelChange={handleModelChange} />
+                      {initialModelSchema && (
+                        <DynamicFormBuilder
+                          jsonSchema={initialModelSchema}
+                          validationURL={`models/llms/${selectedModel}/validate`}
+                          onSuccessCallback={onSuccessCallback}
+                        />
+                      )}
+                    </>
+                  )
                 )}
                 {error && <NotificationBox type={'error'} onClick={() => setError(null)} message={error} />}
               </div>
diff --git a/app/src/client/services/modelService.ts b/app/src/client/services/modelService.ts
index d6e0841..0ff2e73 100644
--- a/app/src/client/services/modelService.ts
+++ b/app/src/client/services/modelService.ts
@@ -1,9 +1,9 @@
 // services/modelService.ts
-import { getModels as apiGetModels } from 'wasp/client/operations';
+import { getAvailableModels as apigetAvailableModels } from 'wasp/client/operations';
 
-export const getModels = async () => {
+export const getAvailableModels = async () => {
   try {
-    const response = await apiGetModels();
+    const response = await apigetAvailableModels();
     return response;
   } catch (error) {
     console.error('Failed to fetch models:', error);
diff --git a/app/src/server/actions.ts b/app/src/server/actions.ts
index 1111dcc..9e66ccc 100644
--- a/app/src/server/actions.ts
+++ b/app/src/server/actions.ts
@@ -4,15 +4,16 @@ import {
   type StripePayment,
   type UpdateCurrentUser,
   type UpdateUserById,
-  type GetModels,
+  type GetAvailableModels,
   type ValidateForm,
+  type UpdateUserModels,
 } from 'wasp/server/operations';
 import Stripe from 'stripe';
 import type { StripePaymentResult } from '../shared/types';
 import { fetchStripeCustomer, createStripeCheckoutSession } from './payments/stripeUtils.js';
 import { TierIds } from '../shared/constants.js';
 
-const FASTAGENCY_SERVER_URL = process.env.FASTAGENCY_SERVER_URL || 'http://127.0.0.1:8000';
+import { FASTAGENCY_SERVER_URL } from './common/constants';
 
 export const stripePayment: StripePayment<string, StripePaymentResult> = async (tier, context) => {
   if (!context.user || !context.user.email) {
@@ -91,7 +92,7 @@ export const updateCurrentUser: UpdateCurrentUser<Partial<User>, User> = async (
   });
 };
 
-export const getModels: GetModels<void, any> = async (user, context) => {
+export const getAvailableModels: GetAvailableModels<void, any> = async (user, context) => {
   if (!context.user) {
     throw new HttpError(401);
   }
@@ -115,6 +116,41 @@ export const getModels: GetModels<void, any> = async (user, context) => {
   }
 };
 
+type UpdateUserModelsValues = {
+  userId: number;
+  model: string;
+  base_url: string;
+  api_type: string;
+  api_version?: string;
+};
+
+type UpdateUserModelsPayload = {
+  data: UpdateUserModelsValues;
+};
+
+export const updateUserModels: UpdateUserModels<UpdateUserModelsPayload, void> = async (args, context) => {
+  if (!context.user) {
+    throw new HttpError(401);
+  }
+
+  try {
+    const response = await fetch(`${FASTAGENCY_SERVER_URL}/models/update`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ userId: context.user.id, ...args.data }),
+    });
+    const json: any = (await response.json()) as { detail?: string }; // Parse JSON once
+
+    if (!response.ok) {
+      const errorMsg = json.detail || `HTTP error with status code ${response.status}`;
+      console.error('Server Error:', errorMsg);
+      throw new Error(errorMsg);
+    }
+  } catch (error: any) {
+    throw new HttpError(500, error.message);
+  }
+};
+
 export const validateForm: ValidateForm<{ data: any; validationURL: string }, any> = async (
   { data, validationURL }: { data: any; validationURL: string },
   context: any
@@ -138,7 +174,7 @@ export const validateForm: ValidateForm<{ data: any; validationURL: string }, an
       );
     }
 
-    return json;
+    return data;
   } catch (error: any) {
     throw new HttpError(error.statusCode || 500, error.message || 'Server or network error occurred');
   }
diff --git a/app/src/server/common/constants.ts b/app/src/server/common/constants.ts
new file mode 100644
index 0000000..0534a13
--- /dev/null
+++ b/app/src/server/common/constants.ts
@@ -0,0 +1 @@
+export const FASTAGENCY_SERVER_URL = process.env.FASTAGENCY_SERVER_URL || 'http://127.0.0.1:8000';
diff --git a/app/src/server/queries.ts b/app/src/server/queries.ts
index e4dcc29..2f08029 100644
--- a/app/src/server/queries.ts
+++ b/app/src/server/queries.ts
@@ -1,6 +1,7 @@
 import { type DailyStats, type User, type PageViewSource } from 'wasp/entities';
 import { HttpError } from 'wasp/server';
-import { type GetDailyStats, type GetPaginatedUsers } from 'wasp/server/operations';
+import { type GetDailyStats, type GetPaginatedUsers, type GetModels } from 'wasp/server/operations';
+import { FASTAGENCY_SERVER_URL } from './common/constants';
 
 type DailyStatsWithSources = DailyStats & {
   sources: PageViewSource[];
@@ -104,3 +105,32 @@ export const getPaginatedUsers: GetPaginatedUsers<GetPaginatedUsersInput, GetPag
     totalPages,
   };
 };
+
+type GetModelsValues = {
+  model: string;
+  base_url: string;
+  api_type: string;
+  api_version?: string;
+};
+
+export const getModels: GetModels<void, GetModelsValues[]> = async (_args, context) => {
+  try {
+    const data = { userId: context.user.id };
+    const response = await fetch(`${FASTAGENCY_SERVER_URL}/models`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify(data),
+    });
+    const json: any = (await response.json()) as { detail?: string }; // Parse JSON once
+
+    if (!response.ok) {
+      const errorMsg = json.detail || `HTTP error with status code ${response.status}`;
+      console.error('Server Error:', errorMsg);
+      throw new Error(errorMsg);
+    }
+
+    return json;
+  } catch (error: any) {
+    throw new HttpError(500, error.message);
+  }
+};