Adds Google VertexAI as model provider

src/app.ts

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

src/auth.ts

@@ -1,21 +1,29 @@
-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;
+  const authHeader = req.headers.authorization;
 
-    if (authHeader) {
-        if (!checkAccessKey(authHeader)) {
-            return res.sendStatus(403);
-        }
-        next();
-    } else {
-        res.sendStatus(401);
+  if (authHeader) {
+    if (!checkAccessKey(authHeader)) {
+      return res.sendStatus(403);
     }
+    next();
+  } else {
+    res.sendStatus(401);
+  }
 };
 
 export const checkAccessKey = (authHeader) => {
-      const token = authHeader.split(' ')[1];
-      return Boolean(authHeader && (token === getAccessKey()));
+  const token = authHeader.split(' ')[1];
+  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([