geckon01/anything-llm
1
0
Fork 0
mirror of https://github.com/Mintplex-Labs/anything-llm.git synced 2026-06-15 23:20:32 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat: support for iam roles for bedrock client (#2632)

Browse Source 
* feat: implement iam role auth for bedrock

* fix: make client refreshes properly when switching between iam_user and iam_role

* checkout agent flow

* fix aiprovider for bedrock in agent use

---------

Co-authored-by: timothycarambat <rambat1010@gmail.com>
This commit is contained in:
Danny Steenman
2025-05-06 22:48:15 +02:00
committed by GitHub
parent b64a77f29f
commit 5500fa2bc5
6 changed files with 147 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
server/utils/AiProviders/bedrock/index.js
View File

@@ -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";

4
server/utils/AiProviders/bedrock/utils.js
View File

@@ -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

13
server/utils/agents/aibitat/providers/ai-provider.js
View File

@@ -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":

48
server/utils/agents/aibitat/providers/bedrock.js
View File

@@ -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() {

5
server/utils/helpers/updateENV.js
View File

@@ -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: {