Mintplex-Labs · timothycarambat · May 6, 2025 · Nov 14, 2024 · Nov 14, 2024 · May 6, 2025
diff --git a/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx b/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
@@ -3,10 +3,11 @@ import { AWS_REGIONS } from "./regions";
 import { useState } from "react";
 
 export default function AwsBedrockLLMOptions({ settings }) {
-  const [useSessionToken, setUseSessionToken] = useState(
-    settings?.AwsBedrockLLMConnectionMethod === "sessionToken"
+  const [connectionMethod, setConnectionMethod] = useState(
+    settings?.AwsBedrockLLMConnectionMethod ?? "iam"
   );
 
+  console.log("connectionMethod", connectionMethod);
   return (
     <div className="w-full flex flex-col">
       {!settings?.credentialsOnly && (
@@ -29,11 +30,11 @@ export default function AwsBedrockLLMOptions({ settings }) {
         </div>
       )}
 
-      <div className="flex flex-col gap-y-2">
+      <div className="flex flex-col gap-y-2 mb-2">
         <input
           type="hidden"
           name="AwsBedrockLLMConnectionMethod"
-          value={useSessionToken ? "sessionToken" : "iam"}
+          value={connectionMethod}
         />
         <div className="flex flex-col w-full">
           <label className="text-theme-text-primary text-sm font-semibold block mb-3">
@@ -43,73 +44,61 @@ export default function AwsBedrockLLMOptions({ settings }) {
             Select the method to authenticate with AWS Bedrock.
           </p>
         </div>
-        <div className="flex items-center justify-start gap-x-4 bg-theme-settings-input-bg p-2.5 rounded-lg w-fit">
-          <span
-            className={`text-sm ${
-              !useSessionToken
-                ? "text-theme-text-primary"
-                : "text-theme-text-secondary"
-            }`}
-          >
-            IAM
-          </span>
-          <label className="relative inline-flex items-center cursor-pointer">
-            <input
-              type="checkbox"
-              className="sr-only peer"
-              checked={useSessionToken}
-              onChange={(e) => setUseSessionToken(e.target.checked)}
-            />
-            <div className="w-11 h-6 bg-[#4b5563] light:bg-[#e5e7eb] peer-focus:outline-none rounded-full peer peer-checked:after:translate-x-full after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-[#d1d5db] light:after:bg-white after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-primary-button"></div>
-          </label>
-          <span
-            className={`text-sm ${
-              useSessionToken
-                ? "text-theme-text-primary"
-                : "text-theme-text-secondary"
-            }`}
-          >
-            Session Token
-          </span>
-        </div>
+        <select
+          name="AwsBedrockLLMConnectionMethod"
+          value={connectionMethod}
+          required={true}
+          onChange={(e) => setConnectionMethod(e.target.value)}
+          className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-fit p-2.5"
+        >
+          <option value="iam">IAM (Explicit Credentials)</option>
+          <option value="sessionToken">
+            Session Token (Temporary Credentials)
+          </option>
+          <option value="iam_role">IAM Role (Implied Credentials)</option>
+        </select>
       </div>
 
       <div className="w-full flex items-center gap-[36px] my-1.5">
-        <div className="flex flex-col w-60">
-          <label className="text-white text-sm font-semibold block mb-3">
-            AWS Bedrock IAM Access ID
-          </label>
-          <input
-            type="password"
-            name="AwsBedrockLLMAccessKeyId"
-            className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
-            placeholder="AWS Bedrock IAM User Access ID"
-            defaultValue={
-              settings?.AwsBedrockLLMAccessKeyId ? "*".repeat(20) : ""
-            }
-            required={true}
-            autoComplete="off"
-            spellCheck={false}
-          />
-        </div>
-        <div className="flex flex-col w-60">
-          <label className="text-white text-sm font-semibold block mb-3">
-            AWS Bedrock IAM Access Key
-          </label>
-          <input
-            type="password"
-            name="AwsBedrockLLMAccessKey"
-            className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
-            placeholder="AWS Bedrock IAM User Access Key"
-            defaultValue={
-              settings?.AwsBedrockLLMAccessKey ? "*".repeat(20) : ""
-            }
-            required={true}
-            autoComplete="off"
-            spellCheck={false}
-          />
-        </div>
-        {useSessionToken && (
+        {["iam", "sessionToken"].includes(connectionMethod) && (
+          <>
+            <div className="flex flex-col w-60">
+              <label className="text-white text-sm font-semibold block mb-3">
+                AWS Bedrock IAM Access ID
+              </label>
+              <input
+                type="password"
+                name="AwsBedrockLLMAccessKeyId"
+                className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+                placeholder="AWS Bedrock IAM User Access ID"
+                defaultValue={
+                  settings?.AwsBedrockLLMAccessKeyId ? "*".repeat(20) : ""
+                }
+                required={true}
+                autoComplete="off"
+                spellCheck={false}
+              />
+            </div>
+            <div className="flex flex-col w-60">
+              <label className="text-white text-sm font-semibold block mb-3">
+                AWS Bedrock IAM Access Key
+              </label>
+              <input
+                type="password"
+                name="AwsBedrockLLMAccessKey"
+                className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+                placeholder="AWS Bedrock IAM User Access Key"
+                defaultValue={
+                  settings?.AwsBedrockLLMAccessKey ? "*".repeat(20) : ""
+                }
+                required={true}
+                autoComplete="off"
+                spellCheck={false}
+              />
+            </div>
+          </>
+        )}
+        {connectionMethod === "sessionToken" && (
           <div className="flex flex-col w-60">
             <label className="text-theme-text-primary text-sm font-semibold block mb-3">
               AWS Bedrock Session Token

diff --git a/server/utils/AiProviders/bedrock/index.js b/server/utils/AiProviders/bedrock/index.js
@@ -42,8 +42,19 @@ class AWSBedrockLLM {
    */
   constructor(embedder = null, modelPreference = null) {
     const requiredEnvVars = [
-      "AWS_BEDROCK_LLM_ACCESS_KEY_ID",
-      "AWS_BEDROCK_LLM_ACCESS_KEY",
+      ...(this.authMethod !== "iam_role"
+        ? [
+            // required for iam and sessionToken
+            "AWS_BEDROCK_LLM_ACCESS_KEY_ID",
+            "AWS_BEDROCK_LLM_ACCESS_KEY",
+          ]
+        : []),
+      ...(this.authMethod === "sessionToken"
+        ? [
+            // required for sessionToken
+            "AWS_BEDROCK_LLM_SESSION_TOKEN",
+          ]
+        : []),
       "AWS_BEDROCK_LLM_REGION",
       "AWS_BEDROCK_LLM_MODEL_PREFERENCE",
     ];
@@ -54,15 +65,6 @@ class AWSBedrockLLM {
         throw new Error(`Required environment variable ${envVar} is not set.`);
     }
 
-    if (
-      process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD === "sessionToken" &&
-      !process.env.AWS_BEDROCK_LLM_SESSION_TOKEN
-    ) {
-      throw new Error(
-        "AWS_BEDROCK_LLM_SESSION_TOKEN is not set for sessionToken authentication method."
-      );
-    }
-
     this.model =
       modelPreference || process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE;
 
@@ -75,13 +77,7 @@ class AWSBedrockLLM {
 
     this.bedrockClient = new BedrockRuntimeClient({
       region: process.env.AWS_BEDROCK_LLM_REGION,
-      credentials: {
-        accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
-        secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
-        ...(this.authMethod === "sessionToken"
-          ? { sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN }
-          : {}),
-      },
+      credentials: this.credentials,
     });
 
     this.embedder = embedder ?? new NativeEmbedder();
@@ -91,10 +87,36 @@ class AWSBedrockLLM {
     );
   }
 
+  /**
+   * Gets the credentials for the AWS Bedrock LLM based on the authentication method provided.
+   * @returns {object} The credentials object.
+   */
+  get credentials() {
+    switch (this.authMethod) {
+      case "iam": // explicit credentials
+        return {
+          accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
+          secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+        };
+      case "sessionToken": // Session token is used for temporary credentials
+        return {
+          accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
+          secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+          sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN,
+        };
+      // IAM role is used for long-term credentials implied by system process
+      // is filled by the AWS SDK automatically if we pass in no credentials
+      case "iam_role":
+        return {};
+      default:
+        return {};
+    }
+  }
+
   /**
    * Gets the configured AWS authentication method ('iam' or 'sessionToken').
    * Defaults to 'iam' if the environment variable is invalid.
-   * @returns {"iam" | "sessionToken"} The authentication method.
+   * @returns {"iam" | "iam_role" | "sessionToken"} The authentication method.
    */
   get authMethod() {
     const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";

diff --git a/server/utils/AiProviders/bedrock/utils.js b/server/utils/AiProviders/bedrock/utils.js
@@ -7,8 +7,8 @@ const DEFAULT_MAX_OUTPUT_TOKENS = 4096;
 /** @type {number} */
 const DEFAULT_CONTEXT_WINDOW_TOKENS = 8191;
 
-/** @type {'iam' | 'sessionToken'} */
-const SUPPORTED_CONNECTION_METHODS = ["iam", "sessionToken"];
+/** @type {'iam' | 'iam_role' | 'sessionToken'} */
+const SUPPORTED_CONNECTION_METHODS = ["iam", "iam_role", "sessionToken"];
 
 /**
  * Parses a MIME type string (e.g., "image/jpeg") to extract and validate the image format

diff --git a/server/utils/agents/aibitat/providers/ai-provider.js b/server/utils/agents/aibitat/providers/ai-provider.js
@@ -117,13 +117,18 @@ class Provider {
           ...config,
         });
       case "bedrock":
+        // Grab just the credentials from the bedrock provider
+        // using a closure to avoid circular dependency + to avoid instantiating the provider
+        const credentials = (() => {
+          const AWSBedrockProvider = require("./bedrock");
+          const bedrockProvider = new AWSBedrockProvider();
+          return bedrockProvider.credentials;
+        })();
+
         return new ChatBedrockConverse({
           model: process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE,
           region: process.env.AWS_BEDROCK_LLM_REGION,
-          credentials: {
-            accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
-            secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
-          },
+          credentials: credentials,
           ...config,
         });
       case "fireworksai":

diff --git a/server/utils/agents/aibitat/providers/bedrock.js b/server/utils/agents/aibitat/providers/bedrock.js
@@ -1,3 +1,6 @@
+const {
+  SUPPORTED_CONNECTION_METHODS,
+} = require("../../../AiProviders/bedrock/utils.js");
 const Provider = require("./ai-provider.js");
 const InheritMultiple = require("./helpers/classes.js");
 const UnTooled = require("./helpers/untooled.js");
@@ -19,15 +22,7 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
     const model = process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE ?? null;
     const client = new ChatBedrockConverse({
       region: process.env.AWS_BEDROCK_LLM_REGION,
-      credentials: {
-        accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
-        secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
-        // If we're using a session token, we need to pass it in as a credential
-        // otherwise we must omit it so it does not conflict if using IAM auth
-        ...(this.authMethod === "sessionToken"
-          ? { sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN }
-          : {}),
-      },
+      credentials: this.credentials,
       model,
     });
 
@@ -37,14 +32,39 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
   }
 
   /**
-   * Get the authentication method for the AWS Bedrock LLM.
-   * There are only two valid values for this setting - anything else will default to "iam".
-   * @returns {"iam"|"sessionToken"}
+   * Gets the credentials for the AWS Bedrock LLM based on the authentication method provided.
+   * @returns {object} The credentials object.
+   */
+  get credentials() {
+    switch (this.authMethod) {
+      case "iam": // explicit credentials
+        return {
+          accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
+          secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+        };
+      case "sessionToken": // Session token is used for temporary credentials
+        return {
+          accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
+          secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+          sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN,
+        };
+      // IAM role is used for long-term credentials implied by system process
+      // is filled by the AWS SDK automatically if we pass in no credentials
+      case "iam_role":
+        return {};
+      default:
+        return {};
+    }
+  }
+
+  /**
+   * Gets the configured AWS authentication method ('iam' or 'sessionToken').
+   * Defaults to 'iam' if the environment variable is invalid.
+   * @returns {"iam" | "iam_role" | "sessionToken"} The authentication method.
    */
   get authMethod() {
     const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";
-    if (!["iam", "sessionToken"].includes(method)) return "iam";
-    return method;
+    return SUPPORTED_CONNECTION_METHODS.includes(method) ? method : "iam";
   }
 
   get client() {

diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js
@@ -1,4 +1,7 @@
 const { Telemetry } = require("../../models/telemetry");
+const {
+  SUPPORTED_CONNECTION_METHODS,
+} = require("../AiProviders/bedrock/utils");
 const { resetAllVectorStores } = require("../vectorStore/resetAllVectorStores");
 
 const KEY_MAPPING = {
@@ -227,7 +230,7 @@ const KEY_MAPPING = {
     envKey: "AWS_BEDROCK_LLM_CONNECTION_METHOD",
     checks: [
       (input) =>
-        ["iam", "sessionToken"].includes(input) ? null : "Invalid value",
+        SUPPORTED_CONNECTION_METHODS.includes(input) ? null : "invalid Value",
     ],
   },
   AwsBedrockLLMAccessKeyId: {