From 0043feae2a7500298a02e84e280536b0606c4493 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E0=AE=AE=E0=AE=A9=E0=AF=8B=E0=AE=9C=E0=AF=8D=E0=AE=95?=
 =?UTF-8?q?=E0=AF=81=E0=AE=AE=E0=AE=BE=E0=AE=B0=E0=AF=8D=20=E0=AE=AA?=
 =?UTF-8?q?=E0=AE=B4=E0=AE=A9=E0=AE=BF=E0=AE=9A=E0=AF=8D=E0=AE=9A=E0=AE=BE?=
 =?UTF-8?q?=E0=AE=AE=E0=AE=BF?= <smartmanoj42857@gmail.com>
Date: Sun, 26 Jan 2025 19:39:11 +0530
Subject: [PATCH] Fix ChatBedrock integration with Amazon Bedrock model

Fixes #308

Update the `ChatBedrock` class to handle the `region_name` parameter correctly.

* Modify `libs/aws/langchain_aws/chat_models/bedrock.py` to ensure the `region_name` parameter is passed to the `ChatBedrockConverse` class and add logic to take `region_name` from `AWS_REGION` environment variable if not provided.
* Add test cases in `libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py` to confirm that `ChatBedrockConverse` works with the `amazon.nova-pro-v1:0` model and takes `region_name` from `AWS_REGION` environment variable if not provided.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/langchain-ai/langchain-aws/issues/308?shareId=XXXX-XXXX-XXXX-XXXX).
---
 libs/aws/langchain_aws/chat_models/bedrock.py |  2 +-
 .../chat_models/test_bedrock_converse.py      | 21 ++++++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/libs/aws/langchain_aws/chat_models/bedrock.py b/libs/aws/langchain_aws/chat_models/bedrock.py
index 73e53202..e5090c56 100644
--- a/libs/aws/langchain_aws/chat_models/bedrock.py
+++ b/libs/aws/langchain_aws/chat_models/bedrock.py
@@ -852,7 +852,7 @@ def _as_converse(self) -> ChatBedrockConverse:
             kwargs["temperature"] = self.temperature
         return ChatBedrockConverse(
             model=self.model_id,
-            region_name=self.region_name,
+            region_name=self.region_name or os.getenv("AWS_REGION"),
             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,
diff --git a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py
index 4841b775..72ddb872 100644
--- a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py
+++ b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py
@@ -55,7 +55,7 @@ def chat_model_class(self) -> Type[BaseChatModel]:
 
     @property
     def chat_model_params(self) -> dict:
-        return {"model": "us.amazon.nova-pro-v1:0"}
+        return {"model": "us.amazon.nova-pro-v1:0", "region_name": "us-east-1"}
 
     @property
     def standard_chat_model_params(self) -> dict:
@@ -242,3 +242,22 @@ def test_guardrails() -> None:
     )
     assert response.response_metadata["stopReason"] == "guardrail_intervened"
     assert response.response_metadata["trace"] is not None
+
+
+def test_chat_bedrock_converse_with_region_name() -> None:
+    model = ChatBedrockConverse(
+        model="us.amazon.nova-pro-v1:0", region_name="us-east-1", temperature=0
+    )
+    response = model.invoke("Create a list of 3 pop songs")
+    assert isinstance(response, BaseModel)
+    assert response.content is not None
+
+
+def test_chat_bedrock_converse_with_aws_region_env() -> None:
+    import os
+
+    os.environ["AWS_REGION"] = "us-east-1"
+    model = ChatBedrockConverse(model="us.amazon.nova-pro-v1:0", temperature=0)
+    response = model.invoke("Create a list of 3 pop songs")
+    assert isinstance(response, BaseModel)
+    assert response.content is not None