get models from backend, save current modal in localstorage

Author: Bogdan Tsechoev

ui/packages/platform/src/api/bot/getLLMModels.ts

@@ -0,0 +1,23 @@
+import {request} from "../../helpers/request";
+import { LLMModel } from "../../types/api/entities/bot";
+
+export const getLLMModels = async (): Promise<{ response: LLMModel[] | null; error: Response | null }> => {
+  const apiServer = process.env.REACT_APP_API_URL_PREFIX || '';
+
+  try {
+    const response = await request(`${apiServer}/llm_models`, {
+      method: 'GET',
+    });
+
+    if (!response.ok) {
+      return { response: null, error: response };
+    }
+
+    const responseData: LLMModel[] = await response.json();
+
+    return { response: responseData, error: null };
+
+  } catch (error) {
+    return { response: null, error: error as Response };
+  }
+}

ui/packages/platform/src/pages/Bot/SettingsDialog/SettingsDialog.tsx

@@ -27,6 +27,7 @@ import { Spinner } from '@postgres.ai/shared/components/Spinner'
 import { colors } from "@postgres.ai/shared/styles/colors";
 import FormLabel from '@mui/material/FormLabel'
 import { Model } from '../hooks'
+import { LLMModel } from "../../../types/api/entities/bot";
 
 export type Visibility = 'public' | 'private';
 
@@ -46,6 +47,7 @@ type PublicChatDialogProps = {
   onSaveChanges: SaveChangesFunction
   isLoading: boolean
   threadId: string | null
+  llmModels: LLMModel[] | null
 }
 
 const useDialogTitleStyles = makeStyles(
@@ -173,7 +175,8 @@ export const SettingsDialog = (props: PublicChatDialogProps) => {
     onClose,
     isOpen,
     isLoading,
-    threadId
+    threadId,
+    llmModels
   } = props;
 
   const [visibility, setVisibility] = useState<Visibility>(defaultVisibility);
@@ -273,38 +276,32 @@ export const SettingsDialog = (props: PublicChatDialogProps) => {
               label="Anyone with a special link and members of the organization can view"
             />
           </RadioGroup>
-          {/*{shareUrl.remark && (
-              <Typography className={classes.remark}>
-                <span className={classes.remarkIcon}>{icons.warningIcon}</span>
-                {shareUrl.remark}
-              </Typography>
-            )}*/}
           {visibility && (
               <div className={classes.urlContainer}>{urlField}</div>
             )}
         </>}
-        <FormLabel component="legend">Model</FormLabel>
-        <RadioGroup
-          aria-label="model"
-          name="model"
-          value={model}
-          onChange={(event) => {
-            setModel(event.target.value as Model)
-          }}
-          className={classes.radioLabel}
-        >
-          <FormControlLabel
-            value="gemini"
-            control={<Radio />}
-            label="gemini-1.5-pro"
-          />
-
-          <FormControlLabel
-            value="gpt"
-            control={<Radio />}
-            label="gpt-4-turbo"
-          />
-        </RadioGroup>
+        {llmModels && <>
+          <FormLabel component="legend">Model</FormLabel>
+          <RadioGroup
+            aria-label="model"
+            name="model"
+            value={model}
+            onChange={(event) => {
+              setModel(event.target.value as Model)
+            }}
+            className={classes.radioLabel}
+          >
+            {llmModels.map((model) =>
+                <FormControlLabel
+                  key={`${model.vendor}/${model.name}`}
+                  value={`${model.vendor}/${model.name}`}
+                  control={<Radio />}
+                  label={model.name}
+                />
+              )
+            }
+          </RadioGroup>
+        </>}
       </DialogContent>
 
       <DialogActions>

ui/packages/platform/src/pages/Bot/hooks.ts

@@ -5,23 +5,26 @@
  *--------------------------------------------------------------------------
  */
 
-import { Dispatch, SetStateAction, useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import useWebSocket, {ReadyState} from "react-use-websocket";
 import { useLocation } from "react-router-dom";
-import {BotMessage} from "../../types/api/entities/bot";
+import { BotMessage, LLMModel } from "../../types/api/entities/bot";
 import {getChatsWithWholeThreads} from "../../api/bot/getChatsWithWholeThreads";
 import {getChats} from "api/bot/getChats";
 import {useAlertSnackbar} from "@postgres.ai/shared/components/AlertSnackbar/useAlertSnackbar";
 import {localStorage} from "../../helpers/localStorage";
 import { makeChatPublic } from "../../api/bot/makeChatPublic";
-import { aiModelBuilder } from "./utils";
+import { getLLMModels } from "../../api/bot/getLLMModels";
 
 
 const WS_URL = process.env.REACT_APP_WS_URL || '';
 
-export type Model = 'oai' | 'gcp'
+export type Model = `${LLMModel['vendor']}/${LLMModel['name']}`
 
-export const DEFAULT_MODEL: Model = 'gcp'
+const DEFAULT_MODEL_NAME = 'gemini-1.5-pro'
+const DEFAULT_VENDOR = 'gcp';
+
+export const DEFAULT_MODEL: Model = `${DEFAULT_VENDOR}/${DEFAULT_MODEL_NAME}`
 
 type ErrorType = {
   code?: number;
@@ -47,8 +50,9 @@ type UseAiBotReturnType = {
   changeChatVisibility: (threadId: string, isPublic: boolean) => void;
   isChangeVisibilityLoading: boolean;
   unsubscribe: (threadId: string) => void;
-  model: Model,
-  setModel: Dispatch<SetStateAction<Model>>
+  model: UseLLMModelsList['model'],
+  setModel: UseLLMModelsList['setModel']
+  llmModels: UseLLMModelsList['llmModels']
 }
 
 type UseAiBotArgs = {
@@ -60,14 +64,15 @@ type UseAiBotArgs = {
 export const useAiBot = (args: UseAiBotArgs): UseAiBotReturnType => {
   const { threadId, onChatLoadingError } = args;
   const { showMessage, closeSnackbar } = useAlertSnackbar();
+  const { llmModels, model, setModel } = useLLMModelsList();
   let location = useLocation<{skipReloading?: boolean}>();
 
   const [messages, setMessages] = useState<BotMessage[] | null>(null);
   const [isLoading, setLoading] = useState<boolean>(false);
   const [error, setError] = useState<ErrorType | null>(null);
   const [wsLoading, setWsLoading] = useState<boolean>(false);
   const [isChangeVisibilityLoading, setIsChangeVisibilityLoading] = useState<boolean>(false);
-  const [model, setModel] = useState<Model>(DEFAULT_MODEL);
+
 
   const token = localStorage.getAuthToken()
 
@@ -144,7 +149,6 @@ export const useAiBot = (args: UseAiBotArgs): UseAiBotReturnType => {
         subscribe(threadId)
         if (response && response.length > 0) {
           setMessages(response);
-          setModel(aiModelBuilder(response?.[response.length - 1]?.ai_model || ''))
         } else {
           if (onChatLoadingError) onChatLoadingError();
           setError({
@@ -239,7 +243,6 @@ export const useAiBot = (args: UseAiBotArgs): UseAiBotReturnType => {
 
   const clearChat = () => {
     setMessages(null);
-    setModel(DEFAULT_MODEL);
   }
 
   const changeChatVisibility = async (threadId: string, isPublic: boolean) => {
@@ -307,7 +310,8 @@ export const useAiBot = (args: UseAiBotArgs): UseAiBotReturnType => {
     messages,
     unsubscribe,
     model,
-    setModel
+    setModel,
+    llmModels
   }
 }
 
@@ -360,4 +364,63 @@ export const useBotChatsList = (orgId?: number): UseBotChatsListHook => {
     getChatsList,
     loading: isLoading
   }
+}
+
+type UseLLMModelsList = {
+  llmModels: LLMModel[] | null
+  error: Response | null
+  model: Model
+  setModel: (model: Model) => void
+}
+
+const useLLMModelsList = (): UseLLMModelsList => {
+  const [llmModels, setLLMModels] = useState<UseLLMModelsList['llmModels']>(null);
+  const [error, setError] = useState<Response | null>(null);
+  const [userModel, setUserModel] = useState<Model | null>(null)
+
+  const getModels = useCallback(async () => {
+    let models = null;
+    try {
+      const { response } = await getLLMModels();
+      setLLMModels(response)
+      const currentModel = window.localStorage.getItem('bot.llm_model')
+
+      if (currentModel && currentModel !== userModel) {
+        setUserModel(currentModel as Model)
+      } else if (!currentModel) {
+        setModel(DEFAULT_MODEL)
+      }
+    } catch (e) {
+      setError(e as unknown as Response)
+    }
+    return models
+  }, []);
+
+  useEffect(() => {
+    let isCancelled = false;
+
+    getModels()
+      .catch((e) => {
+        if (!isCancelled) {
+          setError(e);
+        }
+      });
+    return () => {
+      isCancelled = true;
+    };
+  }, [getModels]);
+
+  const setModel = (model: Model) => {
+    if (model !== userModel) {
+      setUserModel(model);
+      window.localStorage.setItem('bot.llm_model', model)
+    }
+  }
+
+  return {
+    llmModels,
+    error,
+    setModel,
+    model: userModel || DEFAULT_MODEL,
+  }
 }

ui/packages/platform/src/pages/Bot/index.tsx

@@ -17,11 +17,11 @@ import { Messages } from './Messages/Messages';
 import { Command } from './Command/Command';
 import { ChatsList } from "./ChatsList/ChatsList";
 import { BotWrapperProps } from "./BotWrapper";
-import { useBotChatsList, useAiBot, Model } from "./hooks";
+import { useBotChatsList, useAiBot } from "./hooks";
 import { usePrev } from "../../hooks/usePrev";
 import {HeaderButtons} from "./HeaderButtons/HeaderButtons";
 import settings from "../../utils/settings";
-import { SaveChangesFunction, SettingsDialog, Visibility } from "./SettingsDialog/SettingsDialog";
+import { SaveChangesFunction, SettingsDialog } from "./SettingsDialog/SettingsDialog";
 import { theme } from "@postgres.ai/shared/styles/theme";
 import { colors } from "@postgres.ai/shared/styles/colors";
 import { SettingsWithLabel } from "./SettingsWithLabel/SettingsWithLabel";
@@ -105,7 +105,8 @@ export const BotPage = (props: BotPageProps) => {
     changeChatVisibility,
     unsubscribe,
     model,
-    setModel
+    setModel,
+    llmModels
   } = useAiBot({
     threadId: match.params.threadId,
   });
@@ -226,6 +227,7 @@ export const BotPage = (props: BotPageProps) => {
         onClose={toggleSettingsDialog}
         onSaveChanges={handleSaveSettings}
         threadId={match.params.threadId || null}
+        llmModels={llmModels}
       />
       <ChatsList
         isOpen={isChatsListVisible}

ui/packages/platform/src/pages/Bot/utils.ts

@@ -6,15 +6,9 @@
  */
 
 import { API_URL_PREFIX } from "../../config/env";
-import { DEFAULT_MODEL, Model } from "./hooks";
 
 export const permalinkLinkBuilder = (id: string): string => {
   const apiUrl = process.env.REACT_APP_API_URL_PREFIX || API_URL_PREFIX;
   const isV2API = /https?:\/\/.*v2\.postgres\.ai\b/.test(apiUrl);
   return `https://${isV2API ? 'v2.' : ''}postgres.ai/chats/${id}`;
-};
-
-
-export const aiModelBuilder = (model: string): Model => {
-  return /oai|gpt/.test(model) ? "oai" : /gemini|gcp/.test(model) ? "gcp" : DEFAULT_MODEL;
-}
+};

ui/packages/platform/src/types/api/entities/bot.ts

@@ -14,4 +14,11 @@ export type BotMessage = {
   org_id: string
   thread_id: string
   ai_model: string
-}
+}
+
+export type LLMModel = {
+  name: string;
+  vendor: string;
+  isThirdParty: boolean;
+  freeUseAvailable: boolean;
+};