diff --git a/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx b/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
index 00f44a35ff616b0466c02016bb1af4bea02a2325..569ec4395ba4d2d43b48768c57cd7b345f67fbe5 100644
--- a/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
+++ b/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
@@ -1,7 +1,12 @@
import { ArrowSquareOut, Info } from "@phosphor-icons/react";
import { AWS_REGIONS } from "./regions";
+import { useState } from "react";
export default function AwsBedrockLLMOptions({ settings }) {
+ const [useSessionToken, setUseSessionToken] = useState(
+ settings?.AwsBedrockLLMConnectionMethod === "sessionToken"
+ );
+
return (
<div className="w-full flex flex-col">
{!settings?.credentialsOnly && (
@@ -24,6 +29,43 @@ export default function AwsBedrockLLMOptions({ settings }) {
</div>
)}
+ <div className="flex flex-col gap-y-2">
+ <input
+ type="hidden"
+ name="AwsBedrockLLMConnectionMethod"
+ value={useSessionToken ? "sessionToken" : "iam"}
+ />
+ <div className="flex flex-col w-full">
+ <label className="text-white text-sm font-semibold block mb-3">
+ Use session token
+ </label>
+ <p className="text-white/50 text-sm">
+ Select the method to authenticate with AWS Bedrock.
+ </p>
+ </div>
+ <div className="flex items-center justify-start gap-x-4 bg-zinc-900 p-2.5 rounded-lg w-fit">
+ <span
+ className={`text-sm ${!useSessionToken ? "text-white" : "text-white/50"}`}
+ >
+ 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-zinc-700 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-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-white" : "text-white/50"}`}
+ >
+ Session Token
+ </span>
+ </div>
+ </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">
@@ -59,6 +101,25 @@ export default function AwsBedrockLLMOptions({ settings }) {
spellCheck={false}
/>
</div>
+ {useSessionToken && (
+ <div className="flex flex-col w-60">
+ <label className="text-white text-sm font-semibold block mb-3">
+ AWS Bedrock Session Token
+ </label>
+ <input
+ type="password"
+ name="AwsBedrockLLMSessionToken"
+ className="border-none bg-zinc-900 text-white placeholder:text-white/20 text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+ placeholder="AWS Bedrock Session Token"
+ defaultValue={
+ settings?.AwsBedrockLLMSessionToken ? "*".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 region
diff --git a/server/models/systemSettings.js b/server/models/systemSettings.js
index 55569be078d0d796cd4c95b93481323f18538611..41dfc92931687001c38d42201114b38f09d0baea 100644
--- a/server/models/systemSettings.js
+++ b/server/models/systemSettings.js
@@ -505,8 +505,11 @@ const SystemSettings = {
GenericOpenAiKey: !!process.env.GENERIC_OPEN_AI_API_KEY,
GenericOpenAiMaxTokens: process.env.GENERIC_OPEN_AI_MAX_TOKENS,
+ AwsBedrockLLMConnectionMethod:
+ process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam",
AwsBedrockLLMAccessKeyId: !!process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
AwsBedrockLLMAccessKey: !!process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+ AwsBedrockLLMSessionToken: !!process.env.AWS_BEDROCK_LLM_SESSION_TOKEN,
AwsBedrockLLMRegion: process.env.AWS_BEDROCK_LLM_REGION,
AwsBedrockLLMModel: process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE,
AwsBedrockLLMTokenLimit: process.env.AWS_BEDROCK_LLM_MODEL_TOKEN_LIMIT,
diff --git a/server/utils/AiProviders/bedrock/index.js b/server/utils/AiProviders/bedrock/index.js
index c271f7297e7d25aa7a9db0571eb54bc56411aaf4..8d9be2d392862818bff40bbb8f9237bb6721d3dd 100644
--- a/server/utils/AiProviders/bedrock/index.js
+++ b/server/utils/AiProviders/bedrock/index.js
@@ -31,6 +31,14 @@ class AWSBedrockLLM {
if (!process.env.AWS_BEDROCK_LLM_REGION)
throw new Error("No AWS Bedrock LLM region was set.");
+ if (
+ process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD === "sessionToken" &&
+ !process.env.AWS_BEDROCK_LLM_SESSION_TOKEN
+ )
+ throw new Error(
+ "No AWS Bedrock LLM session token was set while using session token as the authentication method."
+ );
+
this.model =
modelPreference || process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE;
this.limits = {
@@ -41,6 +49,20 @@ class AWSBedrockLLM {
this.embedder = embedder ?? new NativeEmbedder();
this.defaultTemp = 0.7;
+ this.#log(
+ `Loaded with model: ${this.model}. Will communicate with AWS Bedrock using ${this.authMethod} authentication.`
+ );
+ }
+
+ /**
+ * 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"}
+ */
+ get authMethod() {
+ const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";
+ if (!["iam", "sessionToken"].includes(method)) return "iam";
+ return method;
}
#bedrockClient({ temperature = 0.7 }) {
@@ -51,6 +73,9 @@ class AWSBedrockLLM {
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 }
+ : {}),
},
temperature,
});
diff --git a/server/utils/agents/aibitat/providers/bedrock.js b/server/utils/agents/aibitat/providers/bedrock.js
index 819dfc4053ff601e0908afe843bb26336d8c87fb..6cdb92897a6601f189b8ad7a73bfd8d7e77408e2 100644
--- a/server/utils/agents/aibitat/providers/bedrock.js
+++ b/server/utils/agents/aibitat/providers/bedrock.js
@@ -22,6 +22,11 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
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 }
+ : {}),
},
model,
});
@@ -31,6 +36,17 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
this.verbose = true;
}
+ /**
+ * 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"}
+ */
+ get authMethod() {
+ const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";
+ if (!["iam", "sessionToken"].includes(method)) return "iam";
+ return method;
+ }
+
get client() {
return this._client;
}
diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js
index a884f1324ba3d1949de426d09bb99e984a95a2ae..6081159a5242f0b72bc2befaf1532a60f6d88882 100644
--- a/server/utils/helpers/updateENV.js
+++ b/server/utils/helpers/updateENV.js
@@ -213,6 +213,13 @@ const KEY_MAPPING = {
},
// AWS Bedrock LLM InferenceSettings
+ AwsBedrockLLMConnectionMethod: {
+ envKey: "AWS_BEDROCK_LLM_CONNECTION_METHOD",
+ checks: [
+ (input) =>
+ ["iam", "sessionToken"].includes(input) ? null : "Invalid value",
+ ],
+ },
AwsBedrockLLMAccessKeyId: {
envKey: "AWS_BEDROCK_LLM_ACCESS_KEY_ID",
checks: [isNotEmpty],
@@ -221,6 +228,10 @@ const KEY_MAPPING = {
envKey: "AWS_BEDROCK_LLM_ACCESS_KEY",
checks: [isNotEmpty],
},
+ AwsBedrockLLMSessionToken: {
+ envKey: "AWS_BEDROCK_LLM_SESSION_TOKEN",
+ checks: [],
+ },
AwsBedrockLLMRegion: {
envKey: "AWS_BEDROCK_LLM_REGION",
checks: [isNotEmpty],