From d8c40253c30f5c695acddb5b061144aaf0b60593 Mon Sep 17 00:00:00 2001 From: Dwai Banerjee Date: Thu, 20 Jul 2023 19:55:59 +0530 Subject: [PATCH] Adding endpoint_url to embeddings/bedrock.py and updated docs (#7927) BedrockEmbeddings does not have endpoint_url so that switching to custom endpoint is not possible. I have access to Bedrock custom endpoint and cannot use BedrockEmbeddings --------- Co-authored-by: Bagatur --- .../text_embedding/integrations/bedrock.ipynb | 4 +++- .../modules/model_io/models/llms/integrations/bedrock.ipynb | 4 +++- langchain/embeddings/bedrock.py | 6 ++++++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/docs/extras/modules/data_connection/text_embedding/integrations/bedrock.ipynb b/docs/extras/modules/data_connection/text_embedding/integrations/bedrock.ipynb index ae161a5285..a69c99de1d 100644 --- a/docs/extras/modules/data_connection/text_embedding/integrations/bedrock.ipynb +++ b/docs/extras/modules/data_connection/text_embedding/integrations/bedrock.ipynb @@ -27,7 +27,9 @@ "source": [ "from langchain.embeddings import BedrockEmbeddings\n", "\n", - "embeddings = BedrockEmbeddings(credentials_profile_name=\"bedrock-admin\")" + "embeddings = BedrockEmbeddings(\n", + " credentials_profile_name=\"bedrock-admin\", endpoint_url=\"custom_endpoint_url\"\n", + ")" ] }, { diff --git a/docs/extras/modules/model_io/models/llms/integrations/bedrock.ipynb b/docs/extras/modules/model_io/models/llms/integrations/bedrock.ipynb index 6a297d26a4..56847a00fd 100644 --- a/docs/extras/modules/model_io/models/llms/integrations/bedrock.ipynb +++ b/docs/extras/modules/model_io/models/llms/integrations/bedrock.ipynb @@ -34,7 +34,9 @@ "from langchain.llms.bedrock import Bedrock\n", "\n", "llm = Bedrock(\n", - " credentials_profile_name=\"bedrock-admin\", model_id=\"amazon.titan-tg1-large\"\n", + " credentials_profile_name=\"bedrock-admin\",\n", + " model_id=\"amazon.titan-tg1-large\",\n", + " endpoint_url=\"custom_endpoint_url\",\n", ")" ] }, diff --git a/langchain/embeddings/bedrock.py b/langchain/embeddings/bedrock.py index ca1fae15dd..93b3d69bbe 100644 --- a/langchain/embeddings/bedrock.py +++ b/langchain/embeddings/bedrock.py @@ -60,6 +60,9 @@ class BedrockEmbeddings(BaseModel, Embeddings): model_kwargs: Optional[Dict] = None """Key word arguments to pass to the model.""" + endpoint_url: Optional[str] = None + """Needed if you don't want to default to us-east-1 endpoint""" + class Config: """Configuration for this pydantic object.""" @@ -85,6 +88,9 @@ class BedrockEmbeddings(BaseModel, Embeddings): if values["region_name"]: client_params["region_name"] = values["region_name"] + if values["endpoint_url"]: + client_params["endpoint_url"] = values["endpoint_url"] + values["client"] = session.client("bedrock", **client_params) except ImportError: