Add AWS client generator utility

libs/aws/langchain_aws/chat_models/bedrock_converse.py

@@ -1,7 +1,6 @@
 import base64
 import json
 import logging
-import os
 import re
 import warnings
 from operator import itemgetter
@@ -22,7 +21,6 @@
     cast,
 )
 
-import boto3
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models import BaseChatModel, LanguageModelInput
@@ -56,6 +54,7 @@
 from typing_extensions import Self
 
 from langchain_aws.function_calling import ToolsOutputParser
+from langchain_aws.utils import create_aws_client
 
 logger = logging.getLogger(__name__)
 _BM = TypeVar("_BM", bound=BaseModel)
@@ -568,53 +567,16 @@ def validate_environment(self) -> Self:
 
         # Skip creating new client if passed in constructor
         if self.client is None:
-            creds = {
-                "aws_access_key_id": self.aws_access_key_id,
-                "aws_secret_access_key": self.aws_secret_access_key,
-                "aws_session_token": self.aws_session_token,
-            }
-            if creds["aws_access_key_id"] and creds["aws_secret_access_key"]:
-                session_params = {
-                    k: v.get_secret_value() for k, v in creds.items() if v
-                }
-            elif any(creds.values()):
-                raise ValueError(
-                    f"If any of aws_access_key_id, aws_secret_access_key, or "
-                    f"aws_session_token are specified then both aws_access_key_id and "
-                    f"aws_secret_access_key must be specified. Only received "
-                    f"{(k for k, v in creds.items() if v)}."
-                )
-            elif self.credentials_profile_name is not None:
-                session_params = {"profile_name": self.credentials_profile_name}
-            else:
-                # use default credentials
-                session_params = {}
-
-            try:
-                session = boto3.Session(**session_params)
-
-                self.region_name = (
-                    self.region_name
-                    or os.getenv("AWS_REGION")
-                    or os.getenv("AWS_DEFAULT_REGION")
-                    or session.region_name
-                )
-
-                client_params = {
-                    "endpoint_url": self.endpoint_url,
-                    "config": self.config,
-                    "region_name": self.region_name,
-                }
-                client_params = {k: v for k, v in client_params.items() if v}
-                self.client = session.client("bedrock-runtime", **client_params)
-            except ValueError as e:
-                raise ValueError(f"Error raised by bedrock service:\n\n{e}") from e
-            except Exception as e:
-                raise ValueError(
-                    "Could not load credentials to authenticate with AWS client. "
-                    "Please check that credentials in the specified "
-                    f"profile name are valid. Bedrock error:\n\n{e}"
-                ) from e
+            self.client = create_aws_client(
+                region_name=self.region_name,
+                credentials_profile_name=self.credentials_profile_name,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+                aws_session_token=self.aws_session_token,
+                endpoint_url=self.endpoint_url,
+                config=self.config,
+                service_name="bedrock-runtime",
+            )
 
         return self
 

libs/aws/langchain_aws/embeddings/bedrock.py

@@ -7,9 +7,12 @@
 import numpy as np
 from langchain_core.embeddings import Embeddings
 from langchain_core.runnables.config import run_in_executor
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from langchain_core.utils import secret_from_env
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
 from typing_extensions import Self
 
+from langchain_aws.utils import create_aws_client
+
 logger = logging.getLogger(__name__)
 
 
@@ -47,7 +50,7 @@ class BedrockEmbeddings(BaseModel, Embeddings):
     client: Any = Field(default=None, exclude=True)  #: :meta private:
     """Bedrock client."""
     region_name: Optional[str] = None
-    """The aws region e.g., `us-west-2`. Fallsback to AWS_DEFAULT_REGION env variable
+    """The aws region e.g., `us-west-2`. Falls back to AWS_REGION/AWS_DEFAULT_REGION env variable
     or region specified in ~/.aws/config in case it is not provided here.
     """
 
@@ -59,6 +62,44 @@ class BedrockEmbeddings(BaseModel, Embeddings):
     See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
     """
 
+    aws_access_key_id: Optional[SecretStr] = Field(
+        default_factory=secret_from_env("AWS_ACCESS_KEY_ID", default=None)
+    )
+    """AWS access key id. 
+
+    If provided, aws_secret_access_key must also be provided.
+    If not specified, the default credential profile or, if on an EC2 instance,
+    credentials from IMDS will be used.
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+
+    If not provided, will be read from 'AWS_ACCESS_KEY_ID' environment variable.
+    """
+
+    aws_secret_access_key: Optional[SecretStr] = Field(
+        default_factory=secret_from_env("AWS_SECRET_ACCESS_KEY", default=None)
+    )
+    """AWS secret_access_key. 
+
+    If provided, aws_access_key_id must also be provided.
+    If not specified, the default credential profile or, if on an EC2 instance,
+    credentials from IMDS will be used.
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+
+    If not provided, will be read from 'AWS_SECRET_ACCESS_KEY' environment variable.
+    """
+
+    aws_session_token: Optional[SecretStr] = Field(
+        default_factory=secret_from_env("AWS_SESSION_TOKEN", default=None)
+    )
+    """AWS session token. 
+
+    If provided, aws_access_key_id and aws_secret_access_key must also be provided.
+    Not required unless using temporary credentials.
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+
+    If not provided, will be read from 'AWS_SESSION_TOKEN' environment variable.
+    """
+
     model_id: str = "amazon.titan-embed-text-v1"
     """Id of the model to call, e.g., amazon.titan-embed-text-v1, this is
     equivalent to the modelId property in the list-foundation-models api"""
@@ -88,42 +129,17 @@ def provider(self) -> str:
     @model_validator(mode="after")
     def validate_environment(self) -> Self:
         """Validate that AWS credentials to and python package exists in environment."""
-
-        if self.client is not None:
-            return self
-
-        try:
-            import boto3
-
-            if self.credentials_profile_name is not None:
-                session = boto3.Session(profile_name=self.credentials_profile_name)
-            else:
-                # use default credentials
-                session = boto3.Session()
-
-            client_params = {}
-            if self.region_name:
-                client_params["region_name"] = self.region_name
-
-            if self.endpoint_url:
-                client_params["endpoint_url"] = self.endpoint_url
-
-            if self.config:
-                client_params["config"] = self.config
-
-            self.client = session.client("bedrock-runtime", **client_params)
-
-        except ImportError:
-            raise ModuleNotFoundError(
-                "Could not import boto3 python package. "
-                "Please install it with `pip install boto3`."
+        if self.client is None:
+            self.client = create_aws_client(
+                region_name=self.region_name,
+                credentials_profile_name=self.credentials_profile_name,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+                aws_session_token=self.aws_session_token,
+                endpoint_url=self.endpoint_url,
+                config=self.config,
+                service_name="bedrock-runtime",
             )
-        except Exception as e:
-            raise ValueError(
-                "Could not load credentials to authenticate with AWS client. "
-                "Please check that credentials in the specified "
-                f"profile name are valid. Bedrock error: {e}"
-            ) from e
 
         return self
 

libs/aws/langchain_aws/llms/bedrock.py

@@ -18,7 +18,6 @@
     Union,
 )
 
-import boto3
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
@@ -34,6 +33,7 @@
 from langchain_aws.function_calling import _tools_in_params
 from langchain_aws.utils import (
     anthropic_tokens_supported,
+    create_aws_client,
     enforce_stop_tokens,
     get_num_tokens_anthropic,
     get_token_ids_anthropic,
@@ -577,7 +577,7 @@ class BedrockBase(BaseLanguageModel, ABC):
     client: Any = Field(default=None, exclude=True)  #: :meta private:
 
     region_name: Optional[str] = Field(default=None, alias="region")
-    """The aws region e.g., `us-west-2`. Fallsback to AWS_REGION or AWS_DEFAULT_REGION 
+    """The aws region e.g., `us-west-2`. Falls back to AWS_REGION or AWS_DEFAULT_REGION 
     env variable or region specified in ~/.aws/config in case it is not provided here.
     """
 
@@ -742,54 +742,17 @@ def validate_environment(self) -> Self:
                 self.model_kwargs.pop("max_tokens")
 
         # Skip creating new client if passed in constructor
-        if self.client is not None:
-            return self
-
-        creds = {
-            "aws_access_key_id": self.aws_access_key_id,
-            "aws_secret_access_key": self.aws_secret_access_key,
-            "aws_session_token": self.aws_session_token,
-        }
-        if creds["aws_access_key_id"] and creds["aws_secret_access_key"]:
-            session_params = {k: v.get_secret_value() for k, v in creds.items() if v}
-        elif any(creds.values()):
-            raise ValueError(
-                f"If any of aws_access_key_id, aws_secret_access_key, or "
-                f"aws_session_token are specified then both aws_access_key_id and "
-                f"aws_secret_access_key must be specified. Only received "
-                f"{(k for k, v in creds.items() if v)}."
+        if self.client is None:
+            self.client = create_aws_client(
+                region_name=self.region_name,
+                credentials_profile_name=self.credentials_profile_name,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+                aws_session_token=self.aws_session_token,
+                endpoint_url=self.endpoint_url,
+                config=self.config,
+                service_name="bedrock-runtime",
             )
-        elif self.credentials_profile_name is not None:
-            session_params = {"profile_name": self.credentials_profile_name}
-        else:
-            # use default credentials
-            session_params = {}
-
-        try:
-            session = boto3.Session(**session_params)
-
-            self.region_name = (
-                self.region_name
-                or os.getenv("AWS_REGION")
-                or os.getenv("AWS_DEFAULT_REGION")
-                or session.region_name
-            )
-
-            client_params = {
-                "endpoint_url": self.endpoint_url,
-                "config": self.config,
-                "region_name": self.region_name,
-            }
-            client_params = {k: v for k, v in client_params.items() if v}
-            self.client = session.client("bedrock-runtime", **client_params)
-        except ValueError as e:
-            raise ValueError(f"Error raised by bedrock service:\n\n{e}") from e
-        except Exception as e:
-            raise ValueError(
-                "Could not load credentials to authenticate with AWS client. "
-                "Please check that credentials in the specified "
-                f"profile name are valid. Bedrock error:\n\n{e}"
-            ) from e
 
         return self