community[patch]: support AWS_... env vars (langchain-ai#6866)

FilipZmijewski · Sep 25, 2024 · dd6399b · dd6399b
1 parent f0c5d02
commit dd6399b
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 25 deletions.
diff --git a/libs/langchain-community/src/chat_models/bedrock/index.ts b/libs/langchain-community/src/chat_models/bedrock/index.ts
@@ -21,9 +21,9 @@ export interface BedrockChatFields
  *
  * ```bash
  * npm install @langchain/openai
- * export BEDROCK_AWS_REGION="your-aws-region"
- * export BEDROCK_AWS_SECRET_ACCESS_KEY="your-aws-secret-access-key"
- * export BEDROCK_AWS_ACCESS_KEY_ID="your-aws-access-key-id"
+ * export AWS_REGION="your-aws-region"
+ * export AWS_SECRET_ACCESS_KEY="your-aws-secret-access-key"
+ * export AWS_ACCESS_KEY_ID="your-aws-access-key-id"
  * ```
  *
  * ## [Constructor args](/classes/langchain_community_chat_models_bedrock.BedrockChat.html#constructor)
@@ -60,15 +60,21 @@ export interface BedrockChatFields
  * const llm = new BedrockChat({
  *   region: process.env.BEDROCK_AWS_REGION,
  *   maxRetries: 0,
- *   credentials: {
- *     secretAccessKey: process.env.BEDROCK_AWS_SECRET_ACCESS_KEY!,
- *     accessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID!,
- *   },
  *   model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
  *   temperature: 0,
  *   maxTokens: undefined,
  *   // other params...
  * });
+ *
+ * // You can also pass credentials in explicitly:
+ * const llmWithCredentials = new BedrockChat({
+ *   region: process.env.BEDROCK_AWS_REGION,
+ *   model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
+ *   credentials: {
+ *     secretAccessKey: process.env.BEDROCK_AWS_SECRET_ACCESS_KEY!,
+ *     accessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID!,
+ *   },
+ * });
  * ```
  * </details>
  *

diff --git a/libs/langchain-community/src/chat_models/bedrock/web.ts b/libs/langchain-community/src/chat_models/bedrock/web.ts
@@ -183,9 +183,9 @@ export interface BedrockChatFields
  *
  * ```bash
  * npm install @langchain/openai
- * export BEDROCK_AWS_REGION="your-aws-region"
- * export BEDROCK_AWS_SECRET_ACCESS_KEY="your-aws-secret-access-key"
- * export BEDROCK_AWS_ACCESS_KEY_ID="your-aws-access-key-id"
+ * export AWS_REGION="your-aws-region"
+ * export AWS_SECRET_ACCESS_KEY="your-aws-secret-access-key"
+ * export AWS_ACCESS_KEY_ID="your-aws-access-key-id"
  * ```
  *
  * ## [Constructor args](/classes/langchain_community_chat_models_bedrock.BedrockChat.html#constructor)
@@ -217,20 +217,26 @@ export interface BedrockChatFields
  * <summary><strong>Instantiate</strong></summary>
  *
  * ```typescript
- * import { BedrockChat } from '@langchain/community/chat_models/bedrock';
+ * import { BedrockChat } from '@langchain/community/chat_models/bedrock/web';
  *
  * const llm = new BedrockChat({
- *   region: process.env.BEDROCK_AWS_REGION,
+ *   region: process.env.AWS_REGION,
  *   maxRetries: 0,
- *   credentials: {
- *     secretAccessKey: process.env.BEDROCK_AWS_SECRET_ACCESS_KEY!,
- *     accessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID!,
- *   },
  *   model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
  *   temperature: 0,
  *   maxTokens: undefined,
  *   // other params...
  * });
+ *
+ * // You can also pass credentials in explicitly:
+ * const llmWithCredentials = new BedrockChat({
+ *   region: process.env.BEDROCK_AWS_REGION,
+ *   model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
+ *   credentials: {
+ *     secretAccessKey: process.env.BEDROCK_AWS_SECRET_ACCESS_KEY!,
+ *     accessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID!,
+ *   },
+ * });
  * ```
  * </details>
  *
@@ -542,8 +548,12 @@ export class BedrockChat
 
   get lc_secrets(): { [key: string]: string } | undefined {
     return {
-      "credentials.accessKeyId": "BEDROCK_AWS_ACCESS_KEY_ID",
-      "credentials.secretAccessKey": "BEDROCK_AWS_SECRET_ACCESS_KEY",
+      "credentials.accessKeyId": "AWS_ACCESS_KEY_ID",
+      "credentials.secretAccessKey": "AWS_SECRET_ACCESS_KEY",
+      "credentials.sessionToken": "AWS_SECRET_ACCESS_KEY",
+      awsAccessKeyId: "AWS_ACCESS_KEY_ID",
+      awsSecretAccessKey: "AWS_SECRET_ACCESS_KEY",
+      awsSessionToken: "AWS_SESSION_TOKEN",
     };
   }
 
@@ -566,7 +576,32 @@ export class BedrockChat
   }
 
   constructor(fields?: BedrockChatFields) {
-    super(fields ?? {});
+    const awsAccessKeyId =
+      fields?.awsAccessKeyId ?? getEnvironmentVariable("AWS_ACCESS_KEY_ID");
+    const awsSecretAccessKey =
+      fields?.awsSecretAccessKey ??
+      getEnvironmentVariable("AWS_SECRET_ACCESS_KEY");
+    const awsSessionToken =
+      fields?.awsSessionToken ?? getEnvironmentVariable("AWS_SESSION_TOKEN");
+
+    let credentials = fields?.credentials;
+    if (credentials === undefined) {
+      if (awsAccessKeyId === undefined || awsSecretAccessKey === undefined) {
+        throw new Error(
+          "Please set your AWS credentials in the 'credentials' field or set env vars AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, and optionally AWS_SESSION_TOKEN."
+        );
+      }
+      credentials = {
+        accessKeyId: awsAccessKeyId,
+        secretAccessKey: awsSecretAccessKey,
+        sessionToken: awsSessionToken,
+      };
+    }
+
+    // eslint-disable-next-line no-param-reassign
+    fields = { ...fields, awsAccessKeyId, awsSecretAccessKey, awsSessionToken };
+
+    super(fields);
 
     this.model = fields?.model ?? this.model;
     this.modelProvider = getModelProvider(this.model);
@@ -585,12 +620,6 @@ export class BedrockChat
     }
     this.region = region;
 
-    const credentials = fields?.credentials;
-    if (!credentials) {
-      throw new Error(
-        "Please set the AWS credentials in the 'credentials' field."
-      );
-    }
     this.credentials = credentials;
 
     this.temperature = fields?.temperature ?? this.temperature;

diff --git a/libs/langchain-community/src/chat_models/tests/chatbedrock.test.ts b/libs/langchain-community/src/chat_models/tests/chatbedrock.test.ts
@@ -26,3 +26,34 @@ test("Test Bedrock identifying params", async () => {
     model,
   });
 });
+
+test("Test Bedrock serialization", async () => {
+  delete process.env.AWS_ACCESS_KEY_ID;
+  delete process.env.AWS_SECRET_ACCESS_KEY;
+  const bedrock = new BedrockChat({
+    region: "us-west-2",
+    model: "anthropic.claude-3-sonnet-20240229-v1:0",
+    credentials: {
+      accessKeyId: "unused",
+      secretAccessKey: "unused",
+      sessionToken: "unused",
+    },
+  });
+
+  expect(JSON.stringify(bedrock)).toEqual(
+    `{"lc":1,"type":"constructor","id":["langchain","chat_models","bedrock","BedrockChat"],"kwargs":{"region_name":"us-west-2","model_id":"anthropic.claude-3-sonnet-20240229-v1:0","credentials":{"accessKeyId":{"lc":1,"type":"secret","id":["AWS_ACCESS_KEY_ID"]},"secretAccessKey":{"lc":1,"type":"secret","id":["AWS_SECRET_ACCESS_KEY"]},"sessionToken":{"lc":1,"type":"secret","id":["AWS_SECRET_ACCESS_KEY"]}}}}`
+  );
+});
+
+test("Test Bedrock serialization from environment variables", async () => {
+  process.env.AWS_ACCESS_KEY_ID = "foo";
+  process.env.AWS_SECRET_ACCESS_KEY = "bar";
+  const bedrock = new BedrockChat({
+    region: "us-west-2",
+    model: "anthropic.claude-3-sonnet-20240229-v1:0",
+  });
+
+  expect(JSON.stringify(bedrock)).toEqual(
+    `{"lc":1,"type":"constructor","id":["langchain","chat_models","bedrock","BedrockChat"],"kwargs":{"region_name":"us-west-2","model_id":"anthropic.claude-3-sonnet-20240229-v1:0","aws_access_key_id":{"lc":1,"type":"secret","id":["AWS_ACCESS_KEY_ID"]},"aws_secret_access_key":{"lc":1,"type":"secret","id":["AWS_SECRET_ACCESS_KEY"]},"credentials":{"accessKeyId":{"lc":1,"type":"secret","id":["AWS_ACCESS_KEY_ID"]},"secretAccessKey":{"lc":1,"type":"secret","id":["AWS_SECRET_ACCESS_KEY"]}}}}`
+  );
+});
diff --git a/libs/langchain-community/src/utils/bedrock/index.ts b/libs/langchain-community/src/utils/bedrock/index.ts
@@ -168,6 +168,12 @@ export interface BaseBedrockInput {
     tagSuffix: string;
     streamProcessingMode: "SYNCHRONOUS" | "ASYNCHRONOUS";
   };
+
+  awsAccessKeyId?: string;
+
+  awsSecretAccessKey?: string;
+
+  awsSessionToken?: string;
 }
 
 type Dict = { [key: string]: unknown };