feat: Bedrock API Keys & filter available models (#5343)

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>

Author: wenxi-onyx

backend/onyx/llm/llm_provider_options.py

@@ -226,21 +226,30 @@ def fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:
                     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.",
+                    description="If using IAM role or a long-term API key, leave this field 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.",
+                    description="If using IAM role or a long-term API key, leave this field blank.",
+                ),
+                CustomConfigKey(
+                    name="AWS_BEARER_TOKEN_BEDROCK",
+                    display_name="AWS Bedrock Long-term API Key",
+                    is_required=False,
+                    is_secret=True,
+                    description=(
+                        "If using IAM role or access key, leave this field blank."
+                    ),
                 ),
             ],
             model_configurations=fetch_model_configurations_for_provider(
                 BEDROCK_PROVIDER_NAME
             ),
             default_model=BEDROCK_DEFAULT_MODEL,
-            default_fast_model=BEDROCK_DEFAULT_MODEL,
+            default_fast_model=None,
         ),
         WellKnownLLMProviderDescriptor(
             name=VERTEXAI_PROVIDER_NAME,
@@ -304,6 +313,7 @@ def fetch_model_configurations_for_provider(
     visible_model_names = (
         fetch_visible_model_names_for_provider_as_set(provider_name) or set()
     )
+
     return [
         ModelConfigurationView(
             name=model_name,

backend/onyx/server/manage/llm/api.py

@@ -1,7 +1,12 @@
+import os
 from collections.abc import Callable
 from datetime import datetime
 from datetime import timezone
 
+import boto3
+from botocore.exceptions import BotoCoreError
+from botocore.exceptions import ClientError
+from botocore.exceptions import NoCredentialsError
 from fastapi import APIRouter
 from fastapi import Depends
 from fastapi import HTTPException
@@ -22,12 +27,14 @@
 from onyx.llm.factory import get_default_llms
 from onyx.llm.factory import get_llm
 from onyx.llm.factory import get_max_input_tokens_from_llm_provider
+from onyx.llm.llm_provider_options import BEDROCK_MODEL_NAMES
 from onyx.llm.llm_provider_options import fetch_available_well_known_llms
 from onyx.llm.llm_provider_options import WellKnownLLMProviderDescriptor
 from onyx.llm.utils import get_llm_contextual_cost
 from onyx.llm.utils import litellm_exception_to_error_msg
 from onyx.llm.utils import model_supports_image_input
 from onyx.llm.utils import test_llm
+from onyx.server.manage.llm.models import BedrockModelsRequest
 from onyx.server.manage.llm.models import LLMCost
 from onyx.server.manage.llm.models import LLMProviderDescriptor
 from onyx.server.manage.llm.models import LLMProviderUpsertRequest
@@ -386,3 +393,84 @@ def get_provider_contextual_cost(
             )
 
     return costs
+
+
+@admin_router.post("/bedrock/available-models")
+def get_bedrock_available_models(
+    request: BedrockModelsRequest,
+    _: User | None = Depends(current_admin_user),
+) -> list[str]:
+    """Fetch available Bedrock models for a specific region and credentials"""
+    try:
+        # Precedence: bearer → keys → IAM
+        if request.aws_bearer_token_bedrock:
+            os.environ["AWS_BEARER_TOKEN_BEDROCK"] = request.aws_bearer_token_bedrock
+            session = boto3.Session(region_name=request.aws_region_name)
+        elif request.aws_access_key_id and request.aws_secret_access_key:
+            session = boto3.Session(
+                aws_access_key_id=request.aws_access_key_id,
+                aws_secret_access_key=request.aws_secret_access_key,
+                region_name=request.aws_region_name,
+            )
+        else:
+            session = boto3.Session(region_name=request.aws_region_name)
+
+        try:
+            bedrock = session.client("bedrock")
+        except Exception as e:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Failed to create Bedrock client: {e}. Check AWS credentials and region.",
+            )
+
+        # Available Bedrock models: text-only, streaming supported
+        model_summaries = bedrock.list_foundation_models().get("modelSummaries", [])
+        available_models = {
+            model.get("modelId", "")
+            for model in model_summaries
+            if model.get("modelId")
+            and "embed" not in model.get("modelId", "").lower()
+            and model.get("responseStreamingSupported", False)
+        }
+
+        # Available inference profiles. Invoking these allows cross-region inference (preferred over base models).
+        profile_ids: set[str] = set()
+        cross_region_models: set[str] = set()
+        try:
+            inference_profiles = bedrock.list_inference_profiles(
+                typeEquals="SYSTEM_DEFINED"
+            ).get("inferenceProfileSummaries", [])
+            for profile in inference_profiles:
+                if profile_id := profile.get("inferenceProfileId"):
+                    profile_ids.add(profile_id)
+
+                    # The model id is everything after the first period in the profile id
+                    if "." in profile_id:
+                        model_id = profile_id.split(".", 1)[1]
+                        cross_region_models.add(model_id)
+        except Exception as e:
+            # Cross-region inference isn't guaranteed; ignore failures here.
+            logger.warning(f"Couldn't fetch inference profiles for Bedrock: {e}")
+
+        # Prefer profiles: de-dupe available models, then add profile IDs
+        candidates = (available_models - cross_region_models) | profile_ids
+
+        # Keep only models we support (compatibility with litellm)
+        filtered = sorted(
+            [model for model in candidates if model in BEDROCK_MODEL_NAMES],
+            reverse=True,
+        )
+
+        # Unset the environment variable, even though it is set again in DefaultMultiLLM init
+        os.environ.pop("AWS_BEARER_TOKEN_BEDROCK", None)
+
+        return filtered
+
+    except (ClientError, NoCredentialsError, BotoCoreError) as e:
+        raise HTTPException(
+            status_code=400, detail=f"Failed to connect to AWS Bedrock: {e}"
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"Unexpected error fetching Bedrock models: {e}"
+        )

backend/onyx/server/manage/llm/models.py

@@ -188,3 +188,11 @@ class LLMCost(BaseModel):
     provider: str
     model_name: str
     cost: float
+
+
+class BedrockModelsRequest(BaseModel):
+    aws_region_name: str
+    aws_access_key_id: str | None = None
+    aws_secret_access_key: str | None = None
+    aws_bearer_token_bedrock: str | None = None
+    provider_name: str | None = None  # Optional: to save models to existing provider

backend/requirements/default.txt

@@ -1,10 +1,10 @@
-aioboto3==14.0.0
+aioboto3==15.1.0
 aiohttp==3.11.16
 alembic==1.10.4
 asyncpg==0.30.0
 atlassian-python-api==3.41.16
 beautifulsoup4==4.12.3
-boto3==1.36.23
+boto3==1.39.11
 celery==5.5.1
 chardet==5.2.0
 chonkie==1.0.10
@@ -93,7 +93,7 @@ zulip==0.8.2
 hubspot-api-client==8.1.0
 asana==5.0.8
 dropbox==11.36.2
-boto3-stubs[s3]==1.34.133
+boto3-stubs[s3]==1.39.11
 shapely==2.0.6
 stripe==10.12.0
 urllib3==2.2.3

backend/requirements/dev.txt

@@ -1,5 +1,5 @@
 black==25.1.0
-boto3-stubs[s3]==1.34.133
+boto3-stubs[s3]==1.39.11
 celery-types==0.19.0
 cohere==5.6.1
 faker==37.1.0

backend/requirements/model_server.txt

@@ -17,5 +17,5 @@ uvicorn==0.35.0
 voyageai==0.2.3
 litellm==1.76.2
 sentry-sdk[fastapi,celery,starlette]==2.14.0
-aioboto3==14.0.0
+aioboto3==15.1.0
 prometheus_fastapi_instrumentator==7.1.0

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

@@ -48,6 +48,8 @@ export function LLMProviderUpdateForm({
 
   const [isTesting, setIsTesting] = useState(false);
   const [testError, setTestError] = useState<string>("");
+  const [isFetchingModels, setIsFetchingModels] = useState(false);
+  const [fetchModelsError, setFetchModelsError] = useState<string>("");
 
   const [showAdvancedOptions, setShowAdvancedOptions] = useState(false);
 
@@ -97,6 +99,9 @@ export function LLMProviderUpdateForm({
         (llmProviderDescriptor.model_configurations
           .filter((modelConfiguration) => modelConfiguration.is_visible)
           .map((modelConfiguration) => modelConfiguration.name) as string[]),
+
+    // Helper field to force re-renders when model list updates
+    _modelListUpdated: 0,
   };
 
   // Setup validation schema if required
@@ -179,6 +184,99 @@ export function LLMProviderUpdateForm({
     );
   };
 
+  const fetchBedrockModels = async (values: any, setFieldValue: any) => {
+    if (llmProviderDescriptor.name !== "bedrock") {
+      return;
+    }
+
+    setIsFetchingModels(true);
+    setFetchModelsError("");
+
+    try {
+      const response = await fetch("/api/admin/llm/bedrock/available-models", {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          aws_region_name: values.custom_config?.AWS_REGION_NAME,
+          aws_access_key_id: values.custom_config?.AWS_ACCESS_KEY_ID,
+          aws_secret_access_key: values.custom_config?.AWS_SECRET_ACCESS_KEY,
+          aws_bearer_token_bedrock:
+            values.custom_config?.AWS_BEARER_TOKEN_BEDROCK,
+          provider_name: existingLlmProvider?.name, // Save models to existing provider if editing
+        }),
+      });
+
+      if (!response.ok) {
+        const errorData = await response.json();
+        throw new Error(errorData.detail || "Failed to fetch models");
+      }
+
+      const availableModels: string[] = await response.json();
+
+      // Update the model configurations with the fetched models
+      const updatedModelConfigs = availableModels.map((modelName) => {
+        // Find existing configuration to preserve is_visible setting
+        const existingConfig = llmProviderDescriptor.model_configurations.find(
+          (config) => config.name === modelName
+        );
+
+        return {
+          name: modelName,
+          is_visible: existingConfig?.is_visible ?? false, // Preserve existing visibility or default to false
+          max_input_tokens: null,
+          supports_image_input: false, // Will be determined by the backend
+        };
+      });
+
+      // Update the descriptor and form values
+      llmProviderDescriptor.model_configurations = updatedModelConfigs;
+
+      // Update selected model names to only include previously visible models that are available
+      const previouslySelectedModels = values.selected_model_names || [];
+      const stillAvailableSelectedModels = previouslySelectedModels.filter(
+        (modelName: string) => availableModels.includes(modelName)
+      );
+      setFieldValue("selected_model_names", stillAvailableSelectedModels);
+
+      // Set a default model if none is set
+      if (
+        (!values.default_model_name ||
+          !availableModels.includes(values.default_model_name)) &&
+        availableModels.length > 0
+      ) {
+        setFieldValue("default_model_name", availableModels[0]);
+      }
+
+      // Clear fast model if it's not in the new list
+      if (
+        values.fast_default_model_name &&
+        !availableModels.includes(values.fast_default_model_name)
+      ) {
+        setFieldValue("fast_default_model_name", null);
+      }
+
+      // Force a re-render by updating a timestamp or counter
+      setFieldValue("_modelListUpdated", Date.now());
+
+      setPopup?.({
+        message: `Successfully fetched ${availableModels.length} models for the selected region (including cross-region inference models).`,
+        type: "success",
+      });
+    } catch (error) {
+      const errorMessage =
+        error instanceof Error ? error.message : "Unknown error";
+      setFetchModelsError(errorMessage);
+      setPopup?.({
+        message: `Failed to fetch models: ${errorMessage}`,
+        type: "error",
+      });
+    } finally {
+      setIsFetchingModels(false);
+    }
+  };
+
   return (
     <Formik
       initialValues={initialValues}
@@ -191,6 +289,7 @@ export function LLMProviderUpdateForm({
           selected_model_names: visibleModels,
           model_configurations: modelConfigurations,
           target_uri,
+          _modelListUpdated,
           ...rest
         } = values;
 
@@ -390,6 +489,7 @@ export function LLMProviderUpdateForm({
                       </ReactMarkdown>
                     }
                     placeholder={customConfigKey.default_value || undefined}
+                    type={customConfigKey.is_secret ? "password" : "text"}
                   />
                 </div>
               );
@@ -407,6 +507,47 @@ export function LLMProviderUpdateForm({
             }
           })}
 
+          {/* Bedrock-specific fetch models button */}
+          {llmProviderDescriptor.name === "bedrock" && (
+            <div className="flex flex-col gap-2">
+              <Button
+                type="button"
+                onClick={() =>
+                  fetchBedrockModels(
+                    formikProps.values,
+                    formikProps.setFieldValue
+                  )
+                }
+                disabled={
+                  isFetchingModels ||
+                  !formikProps.values.custom_config?.AWS_REGION_NAME
+                }
+                className="w-fit"
+              >
+                {isFetchingModels ? (
+                  <>
+                    <LoadingAnimation size="text-sm" />
+                    <span className="ml-2">Fetching Models...</span>
+                  </>
+                ) : (
+                  "Fetch Available Models for Region"
+                )}
+              </Button>
+
+              {fetchModelsError && (
+                <Text className="text-red-600 text-sm">{fetchModelsError}</Text>
+              )}
+
+              <Text className="text-sm text-gray-600">
+                Enter your AWS region, then click this button to fetch available
+                Bedrock models.
+                <br />
+                If you&apos;re updating your existing provider, you&apos;ll need
+                to click this button to fetch the latest models.
+              </Text>
+            </div>
+          )}
+
           {!firstTimeConfiguration && (
             <>
               <Separator />