feat: Vertex AI support (onyx-dot-app#4458)

* Add gemini well-known-llm-provider

* Edit styling of anonymous function

* Remove space

* Edit how advanced options are displayed

* Add VertexAI to acceptable llm providers

* Add new `FileUploadFormField` component

* Edit FileUpload component

* Clean up logic for displaying native llm providers; add support for more complex `CustomConfigKey` types

* Fix minor nits in web app

* Add ability to pass vertex credentials to `litellm`

* Remove unused prop

* Change name of enum value

* Add back ability to change form based on first time configurations

* Create new Error with string instead of throwing raw string

* Add more Gemini models

* Edit mappings for Gemini models

* Edit comment

* Rearrange llm models

* Run black formatter

* Remove complex configurations during first time registration

* Fix nit

* Update llm provider name

* Edit temporary formik field to also have the filename

* Run reformatter

* Reorder commits

* Add advanced configurations for enabled LLM Providers

Author: raunakab

backend/onyx/llm/chat_llm.py

@@ -36,6 +36,7 @@
 from onyx.llm.interfaces import LLM
 from onyx.llm.interfaces import LLMConfig
 from onyx.llm.interfaces import ToolChoiceOptions
+from onyx.llm.llm_provider_options import CREDENTIALS_FILE_CUSTOM_CONFIG_KEY
 from onyx.llm.utils import model_is_reasoning_model
 from onyx.server.utils import mask_string
 from onyx.utils.logger import setup_logger
@@ -454,6 +455,13 @@ def _completion(
                     if structured_response_format
                     else {}
                 ),
+                **(
+                    {
+                        "vertex_credentials": self.config.credentials_file,
+                    }
+                    if self.config.model_provider == "vertex_ai"
+                    else {}
+                ),
                 **self._model_kwargs,
             )
         except Exception as e:
@@ -469,6 +477,12 @@ def _completion(
 
     @property
     def config(self) -> LLMConfig:
+        credentials_file: str | None = (
+            self._custom_config.get(CREDENTIALS_FILE_CUSTOM_CONFIG_KEY, None)
+            if self._custom_config
+            else None
+        )
+
         return LLMConfig(
             model_provider=self._model_provider,
             model_name=self._model_version,
@@ -477,6 +491,7 @@ def config(self) -> LLMConfig:
             api_base=self._api_base,
             api_version=self._api_version,
             deployment_name=self._deployment_name,
+            credentials_file=credentials_file,
         )
 
     def _invoke_implementation(

backend/onyx/llm/interfaces.py

@@ -26,6 +26,7 @@ class LLMConfig(BaseModel):
     api_base: str | None = None
     api_version: str | None = None
     deployment_name: str | None = None
+    credentials_file: str | None = None
     # This disables the "model_" protected namespace for pydantic
     model_config = {"protected_namespaces": ()}
 

backend/onyx/llm/llm_provider_options.py

@@ -1,12 +1,26 @@
+from enum import Enum
+
 import litellm  # type: ignore
 from pydantic import BaseModel
 
 
+class CustomConfigKeyType(Enum):
+    # used for configuration values that require manual input
+    # i.e., textual API keys (e.g., "abcd1234")
+    TEXT_INPUT = "text_input"
+
+    # used for configuration values that require a file to be selected/drag-and-dropped
+    # i.e., file based credentials (e.g., "/path/to/credentials/file.json")
+    FILE_INPUT = "file_input"
+
+
 class CustomConfigKey(BaseModel):
     name: str
+    display_name: str
     description: str | None = None
     is_required: bool = True
     is_secret: bool = False
+    key_type: CustomConfigKeyType = CustomConfigKeyType.TEXT_INPUT
 
 
 class WellKnownLLMProviderDescriptor(BaseModel):
@@ -77,17 +91,50 @@ class WellKnownLLMProviderDescriptor(BaseModel):
 AZURE_PROVIDER_NAME = "azure"
 
 
+VERTEXAI_PROVIDER_NAME = "vertex_ai"
+VERTEXAI_DEFAULT_MODEL = "gemini-2.0-flash"
+VERTEXAI_DEFAULT_FAST_MODEL = "gemini-2.0-flash-lite"
+VERTEXAI_MODEL_NAMES = [
+    # 2.5 pro models
+    "gemini-2.5-pro-exp-03-25",
+    # 2.0 flash-lite models
+    VERTEXAI_DEFAULT_FAST_MODEL,
+    "gemini-2.0-flash-lite-001",
+    # "gemini-2.0-flash-lite-preview-02-05",
+    # 2.0 flash models
+    VERTEXAI_DEFAULT_MODEL,
+    "gemini-2.0-flash-001",
+    "gemini-2.0-flash-exp",
+    # "gemini-2.0-flash-exp-image-generation",
+    # "gemini-2.0-flash-thinking-exp-01-21",
+    # 1.5 pro models
+    "gemini-1.5-pro-latest",
+    "gemini-1.5-pro",
+    "gemini-1.5-pro-001",
+    "gemini-1.5-pro-002",
+    # 1.5 flash models
+    "gemini-1.5-flash-latest",
+    "gemini-1.5-flash",
+    "gemini-1.5-flash-001",
+    "gemini-1.5-flash-002",
+]
+
+
 _PROVIDER_TO_MODELS_MAP = {
     OPENAI_PROVIDER_NAME: OPEN_AI_MODEL_NAMES,
     BEDROCK_PROVIDER_NAME: BEDROCK_MODEL_NAMES,
     ANTHROPIC_PROVIDER_NAME: ANTHROPIC_MODEL_NAMES,
+    VERTEXAI_PROVIDER_NAME: VERTEXAI_MODEL_NAMES,
 }
 
 
+CREDENTIALS_FILE_CUSTOM_CONFIG_KEY = "CREDENTIALS_FILE"
+
+
 def fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:
     return [
         WellKnownLLMProviderDescriptor(
-            name="openai",
+            name=OPENAI_PROVIDER_NAME,
             display_name="OpenAI",
             api_key_required=True,
             api_base_required=False,
@@ -126,14 +173,19 @@ def fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:
             api_base_required=False,
             api_version_required=False,
             custom_config_keys=[
-                CustomConfigKey(name="AWS_REGION_NAME"),
+                CustomConfigKey(
+                    name="AWS_REGION_NAME",
+                    display_name="AWS Region Name",
+                ),
                 CustomConfigKey(
                     name="AWS_ACCESS_KEY_ID",
+                    display_name="AWS Access Key ID",
                     is_required=False,
                     description="If using AWS IAM roles, AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY can be left blank.",
                 ),
                 CustomConfigKey(
                     name="AWS_SECRET_ACCESS_KEY",
+                    display_name="AWS Secret Access Key",
                     is_required=False,
                     is_secret=True,
                     description="If using AWS IAM roles, AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY can be left blank.",
@@ -143,6 +195,26 @@ def fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:
             default_model="anthropic.claude-3-5-sonnet-20241022-v2:0",
             default_fast_model="anthropic.claude-3-5-sonnet-20241022-v2:0",
         ),
+        WellKnownLLMProviderDescriptor(
+            name=VERTEXAI_PROVIDER_NAME,
+            display_name="GCP Vertex AI",
+            api_key_required=False,
+            api_base_required=False,
+            api_version_required=False,
+            llm_names=fetch_models_for_provider(VERTEXAI_PROVIDER_NAME),
+            custom_config_keys=[
+                CustomConfigKey(
+                    name=CREDENTIALS_FILE_CUSTOM_CONFIG_KEY,
+                    display_name="Credentials File",
+                    description="This should be a JSON file containing some private credentials.",
+                    is_required=True,
+                    is_secret=False,
+                    key_type=CustomConfigKeyType.FILE_INPUT,
+                ),
+            ],
+            default_model=VERTEXAI_DEFAULT_MODEL,
+            default_fast_model=VERTEXAI_DEFAULT_MODEL,
+        ),
     ]
 
 

web/src/app/admin/configuration/llm/ConfiguredLLMProviderDisplay.tsx

@@ -28,6 +28,9 @@ function LLMProviderUpdateModal({
     : llmProviderDescriptor?.display_name ||
       llmProviderDescriptor?.name ||
       "Custom LLM Provider";
+
+  const hasAdvancedOptions = llmProviderDescriptor?.name != "azure";
+
   return (
     <Modal
       title={`${llmProviderDescriptor ? "Configure" : "Setup"} ${providerName}`}
@@ -41,6 +44,7 @@ function LLMProviderUpdateModal({
             existingLlmProvider={existingLlmProvider}
             shouldMarkAsDefault={shouldMarkAsDefault}
             setPopup={setPopup}
+            hasAdvancedOptions={hasAdvancedOptions}
           />
         ) : (
           <CustomLLMProviderUpdateForm

web/src/app/admin/configuration/llm/LLMConfiguration.tsx

@@ -34,6 +34,9 @@ function LLMProviderUpdateModal({
     llmProviderDescriptor?.name ||
     existingLlmProvider?.name ||
     "Custom LLM Provider";
+
+  const hasAdvancedOptions = llmProviderDescriptor?.name != "azure";
+
   return (
     <Modal title={`Setup ${providerName}`} onOutsideClick={() => onClose()}>
       <div className="max-h-[70vh] overflow-y-auto px-4">
@@ -44,6 +47,7 @@ function LLMProviderUpdateModal({
             existingLlmProvider={existingLlmProvider}
             shouldMarkAsDefault={shouldMarkAsDefault}
             setPopup={setPopup}
+            hasAdvancedOptions={hasAdvancedOptions}
           />
         ) : (
           <CustomLLMProviderUpdateForm
@@ -75,9 +79,8 @@ function DefaultLLMProviderDisplay({
       {popup}
       <div className="border border-border p-3 dark:bg-neutral-800 dark:border-neutral-700 rounded w-96 flex shadow-md">
         <div className="my-auto">
-          <div className="font-bold">{providerName} </div>
+          <div className="font-bold">{providerName}</div>
         </div>
-
         <div className="ml-auto">
           <Button variant="navigate" onClick={() => setFormIsVisible(true)}>
             Set up
@@ -169,15 +172,13 @@ export function LLMConfiguration() {
       </Text>
 
       <div className="gap-y-4 flex flex-col">
-        {llmProviderDescriptors.map((llmProviderDescriptor) => {
-          return (
-            <DefaultLLMProviderDisplay
-              key={llmProviderDescriptor.name}
-              llmProviderDescriptor={llmProviderDescriptor}
-              shouldMarkAsDefault={existingLlmProviders.length === 0}
-            />
-          );
-        })}
+        {llmProviderDescriptors.map((llmProviderDescriptor) => (
+          <DefaultLLMProviderDisplay
+            key={llmProviderDescriptor.name}
+            llmProviderDescriptor={llmProviderDescriptor}
+            shouldMarkAsDefault={existingLlmProviders.length === 0}
+          />
+        ))}
       </div>
 
       <div className="mt-4">