Adds Google VertexAI as model provider

README.md

@@ -112,6 +112,7 @@ You need to edit the ports accordingly.
 1a: Copy the `sample.env` file to `.env`
 1b: Copy the `deploy/gcp/sample.env` file to `deploy/gcp/.env`
 2a: Fillout desired LLM provider access keys etc. in `.env`
+
 - Note: you will have to comeback and edit this file again once you have the address of the K8s backend deploy
   2b: Fillout the GCP info in `deploy/gcp/.env`
   3: Edit `GCP_REPO` to the correct docker image repo path if you are using something other than Container registry

docker-compose.yaml

@@ -17,7 +17,7 @@ services:
       args:
         - SEARXNG_API_URL=null
     environment:
-      SEARXNG_API_URL: "http://searxng:8080"
+      SEARXNG_API_URL: 'http://searxng:8080'
       SUPER_SECRET_KEY: ${SUPER_SECRET_KEY}
       OPENAI: ${OPENAI}
       GROQ: ${GROQ}

src/app.ts

@@ -21,7 +21,7 @@ app.use(cors(corsOptions));
 
 if (getAccessKey()) {
   app.all('*', requireAccessKey);
-};
+}
 
 app.use(express.json());
 

src/auth.ts

@@ -1,6 +1,5 @@
-import {
-  getAccessKey,
-} from './config';
+import { auth } from 'google-auth-library';
+import { getAccessKey } from './config';
 
 export const requireAccessKey = (req, res, next) => {
   const authHeader = req.headers.authorization;
@@ -17,5 +16,14 @@ export const requireAccessKey = (req, res, next) => {
 
 export const checkAccessKey = (authHeader) => {
   const token = authHeader.split(' ')[1];
-      return Boolean(authHeader && (token === getAccessKey()));
+  return Boolean(authHeader && token === getAccessKey());
+};
+
+export const hasGCPCredentials = async () => {
+  try {
+    const credentials = await auth.getCredentials();
+    return Object.keys(credentials).length > 0;
+  } catch (e) {
+    return false;
+  }
 };

src/config.ts

@@ -34,33 +34,38 @@ const loadEnv = () => {
     GENERAL: {
       PORT: Number(process.env.PORT),
       SIMILARITY_MEASURE: process.env.SIMILARITY_MEASURE,
-      SUPER_SECRET_KEY: process.env.SUPER_SECRET_KEY
+      SUPER_SECRET_KEY: process.env.SUPER_SECRET_KEY,
     },
     API_KEYS: {
       OPENAI: process.env.OPENAI,
-      GROQ: process.env.GROQ
+      GROQ: process.env.GROQ,
     },
     API_ENDPOINTS: {
       SEARXNG: process.env.SEARXNG_API_URL,
-      OLLAMA: process.env.OLLAMA_API_URL
-    }
+      OLLAMA: process.env.OLLAMA_API_URL,
+    },
   } as Config;
 };
 
 export const getPort = () => loadConfig().GENERAL.PORT;
 
-export const getAccessKey = () => loadEnv().GENERAL.SUPER_SECRET_KEY || loadConfig().GENERAL.SUPER_SECRET_KEY;
+export const getAccessKey = () =>
+  loadEnv().GENERAL.SUPER_SECRET_KEY || loadConfig().GENERAL.SUPER_SECRET_KEY;
 
 export const getSimilarityMeasure = () =>
   loadConfig().GENERAL.SIMILARITY_MEASURE;
 
-export const getOpenaiApiKey = () => loadEnv().API_KEYS.OPENAI || loadConfig().API_KEYS.OPENAI;
+export const getOpenaiApiKey = () =>
+  loadEnv().API_KEYS.OPENAI || loadConfig().API_KEYS.OPENAI;
 
-export const getGroqApiKey = () => loadEnv().API_KEYS.GROQ || loadConfig().API_KEYS.GROQ;
+export const getGroqApiKey = () =>
+  loadEnv().API_KEYS.GROQ || loadConfig().API_KEYS.GROQ;
 
-export const getSearxngApiEndpoint = () => loadEnv().API_ENDPOINTS.SEARXNG || loadConfig().API_ENDPOINTS.SEARXNG;
+export const getSearxngApiEndpoint = () =>
+  loadEnv().API_ENDPOINTS.SEARXNG || loadConfig().API_ENDPOINTS.SEARXNG;
 
-export const getOllamaApiEndpoint = () => loadEnv().API_ENDPOINTS.OLLAMA || loadConfig().API_ENDPOINTS.OLLAMA;
+export const getOllamaApiEndpoint = () =>
+  loadEnv().API_ENDPOINTS.OLLAMA || loadConfig().API_ENDPOINTS.OLLAMA;
 
 export const updateConfig = (config: RecursivePartial<Config>) => {
   const currentConfig = loadConfig();

src/lib/providers.ts

@@ -1,6 +1,7 @@
 import { ChatOpenAI, OpenAIEmbeddings } from '@langchain/openai';
 import { ChatOllama } from '@langchain/community/chat_models/ollama';
 import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama';
+import { hasGCPCredentials } from '../auth';
 import {
   getGroqApiKey,
   getOllamaApiEndpoint,
@@ -111,6 +112,23 @@ export const getAvailableChatModelProviders = async () => {
     }
   }
 
+  if (await hasGCPCredentials()) {
+    try {
+      models['vertexai'] = {
+        'gemini-1.5-pro (preview-0409)': new VertexAI({
+          temperature: 0.7,
+          modelName: 'gemini-1.5-pro-preview-0409',
+        }),
+        'gemini-1.0-pro (Latest)': new VertexAI({
+          temperature: 0.7,
+          modelName: 'gemini-1.0-pro',
+        }),
+      };
+    } catch (err) {
+      logger.error(`Error loading VertexAI models: ${err}`);
+    }
+  }
+
   models['custom_openai'] = {};
 
   return models;

src/websocket/connectionManager.ts

@@ -31,7 +31,7 @@ export const handleConnection = async (
           }),
         );
         ws.close();
-      };
+      }
     }
 
     const [chatModelProviders, embeddingModelProviders] = await Promise.all([

ui/components/ChatWindow.tsx

@@ -36,14 +36,11 @@ const useSocket = (url: string) => {
           !embeddingModel ||
           !embeddingModelProvider
         ) {
-          const providers = await clientFetch(
-            '/models',
-            {
+          const providers = await clientFetch('/models', {
             headers: {
               'Content-Type': 'application/json',
             },
-            },
-          ).then(async (res) => await res.json());
+          }).then(async (res) => await res.json());
 
           const chatModelProviders = providers.chatModelProviders;
           const embeddingModelProviders = providers.embeddingModelProviders;
@@ -103,8 +100,8 @@ const useSocket = (url: string) => {
         const secretToken = getAccessKey();
 
         if (secretToken) {
-          protocols = ["Authorization", `${secretToken}`];
-        };
+          protocols = ['Authorization', `${secretToken}`];
+        }
 
         const ws = new WebSocket(wsURL.toString(), protocols);
 

ui/components/SearchImages.tsx

@@ -34,9 +34,7 @@ const SearchImages = ({
             const chatModelProvider = localStorage.getItem('chatModelProvider');
             const chatModel = localStorage.getItem('chatModel');
 
-            const res = await clientFetch(
-              '/images',
-              {
+            const res = await clientFetch('/images', {
               method: 'POST',
               headers: {
                 'Content-Type': 'application/json',
@@ -47,8 +45,7 @@ const SearchImages = ({
                 chat_model_provider: chatModelProvider,
                 chat_model: chatModel,
               }),
-              },
-            );
+            });
 
             const data = await res.json();
 

ui/components/SearchVideos.tsx

@@ -47,9 +47,7 @@ const Searchvideos = ({
             const chatModelProvider = localStorage.getItem('chatModelProvider');
             const chatModel = localStorage.getItem('chatModel');
 
-            const res = await clientFetch(
-              '/videos',
-              {
+            const res = await clientFetch('/videos', {
               method: 'POST',
               headers: {
                 'Content-Type': 'application/json',
@@ -60,8 +58,7 @@ const Searchvideos = ({
                 chat_model_provider: chatModelProvider,
                 chat_model: chatModel,
               }),
-              },
-            );
+            });
 
             const data = await res.json();
 

ui/lib/config.ts

@@ -11,11 +11,12 @@ const loadEnv = () => {
     GENERAL: {
       NEXT_PUBLIC_SUPER_SECRET_KEY: process.env.NEXT_PUBLIC_SUPER_SECRET_KEY!,
       NEXT_PUBLIC_API_URL: process.env.NEXT_PUBLIC_API_URL!,
-      NEXT_PUBLIC_WS_URL: process.env.NEXT_PUBLIC_WS_URL!
+      NEXT_PUBLIC_WS_URL: process.env.NEXT_PUBLIC_WS_URL!,
     },
   } as Config;
 };
 
-export const getAccessKey = () => loadEnv().GENERAL.NEXT_PUBLIC_SUPER_SECRET_KEY;
+export const getAccessKey = () =>
+  loadEnv().GENERAL.NEXT_PUBLIC_SUPER_SECRET_KEY;
 
 export const getBackendURL = () => loadEnv().GENERAL.NEXT_PUBLIC_API_URL;

ui/lib/utils.ts

@@ -1,6 +1,6 @@
 import clsx, { ClassValue } from 'clsx';
 import { twMerge } from 'tailwind-merge';
-import { getAccessKey, getBackendURL } from './config'
+import { getAccessKey, getBackendURL } from './config';
 
 export const cn = (...classes: ClassValue[]) => twMerge(clsx(...classes));
 
@@ -29,11 +29,11 @@ export const clientFetch = async (path: string, payload: any): Promise<any> => {
   if (secretToken) {
     if (headers == null) {
       headers = {};
-    };
+    }
 
     headers['Authorization'] = `Bearer ${secretToken}`;
     payload.headers = headers;
-  };
+  }
 
   return await fetch(url, payload);
 };