From cc1fdc0d5c46293489051292d1d3aeb01a4e2ddf Mon Sep 17 00:00:00 2001
From: Bogdan Tsechoev <bogdan@postgres.ai>
Date: Fri, 20 Dec 2024 02:35:08 +0000
Subject: [PATCH] Assistant UI: Switch to bot_llm_models API endpoint

---
 .../platform/src/api/bot/getAiModels.ts       | 19 +++++++++++++------
 ui/packages/platform/src/pages/Bot/hooks.tsx  |  6 +++---
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/ui/packages/platform/src/api/bot/getAiModels.ts b/ui/packages/platform/src/api/bot/getAiModels.ts
index 2e4e2056..b8e1afb0 100644
--- a/ui/packages/platform/src/api/bot/getAiModels.ts
+++ b/ui/packages/platform/src/api/bot/getAiModels.ts
@@ -1,21 +1,28 @@
 import {request} from "../../helpers/request";
 import { AiModel } from "../../types/api/entities/bot";
 
-export const getAiModels = async (): Promise<{ response: AiModel[] | null; error: Response | null }> => {
+export const getAiModels = async (orgId?: number): Promise<{ response: AiModel[] | null; error: Response | null }> => {
   const apiServer = process.env.REACT_APP_API_URL_PREFIX || '';
-
+  const body = {
+    org_id: orgId
+  }
   try {
-    const response = await request(`${apiServer}/llm_models`, {
-      method: 'GET',
+    const response = await request(`${apiServer}/rpc/bot_llm_models`, {
+      method: 'POST',
+      body: JSON.stringify(body),
+      headers: {
+        'Accept': 'application/vnd.pgrst.object+json',
+        'Prefer': 'return=representation',
+      }
     });
 
     if (!response.ok) {
       return { response: null, error: response };
     }
 
-    const responseData: AiModel[] = await response.json();
+    const responseData: { bot_llm_models: AiModel[] | null } = await response.json();
 
-    return { response: responseData, error: null };
+    return { response: responseData?.bot_llm_models, error: null };
 
   } catch (error) {
     return { response: null, error: error as Response };
diff --git a/ui/packages/platform/src/pages/Bot/hooks.tsx b/ui/packages/platform/src/pages/Bot/hooks.tsx
index 0eecc1a4..840b3785 100644
--- a/ui/packages/platform/src/pages/Bot/hooks.tsx
+++ b/ui/packages/platform/src/pages/Bot/hooks.tsx
@@ -89,7 +89,7 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
     aiModel,
     setAiModel,
     loading: aiModelsLoading
-  } = useAiModelsList();
+  } = useAiModelsList(orgId);
   let location = useLocation<{skipReloading?: boolean}>();
 
   const {
@@ -567,7 +567,7 @@ type UseAiModelsList = {
   setAiModel: (model: AiModel) => void
 }
 
-export const useAiModelsList = (): UseAiModelsList => {
+export const useAiModelsList = (orgId?: number): UseAiModelsList => {
   const [llmModels, setLLMModels] = useState<UseAiModelsList['aiModels']>(null);
   const [error, setError] = useState<Response | null>(null);
   const [userModel, setUserModel] = useState<AiModel | null>(null);
@@ -577,7 +577,7 @@ export const useAiModelsList = (): UseAiModelsList => {
     let models = null;
     setLoading(true);
     try {
-      const { response } = await getAiModels();
+      const { response } = await getAiModels(orgId);
       setLLMModels(response);
       const currentModel = window.localStorage.getItem('bot.ai_model');
       const parsedModel: AiModel = currentModel ? JSON.parse(currentModel) : null;