aws[patch]: make ChatBedrockConverse serializable (#212)

langchain-ai · Sep 20, 2024 · 7c20d8b · 7c20d8b
1 parent 019041e
commit 7c20d8b
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 8 deletions.
diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py
@@ -397,14 +397,6 @@ class Joke(BaseModel):
         populate_by_name=True,
     )
 
-    @property
-    def lc_secrets(self) -> Dict[str, str]:
-        return {
-            "aws_access_key_id": "AWS_ACCESS_KEY_ID",
-            "aws_secret_access_key": "AWS_SECRET_ACCESS_KEY",
-            "aws_session_token": "AWS_SESSION_TOKEN",
-        }
-
     @model_validator(mode="before")
     @classmethod
     def set_disable_streaming(cls, values: Dict) -> Any:
@@ -662,6 +654,22 @@ def _llm_type(self) -> str:
         """Return type of chat model."""
         return "amazon_bedrock_converse_chat"
 
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        return True
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        return ["langchain_aws", "chat_models"]
+
+    @property
+    def lc_secrets(self) -> Dict[str, str]:
+        return {
+            "aws_access_key_id": "AWS_ACCESS_KEY_ID",
+            "aws_secret_access_key": "AWS_SECRET_ACCESS_KEY",
+            "aws_session_token": "AWS_SESSION_TOKEN",
+        }
+
 
 def _messages_to_bedrock(
     messages: List[BaseMessage],

diff --git a/libs/aws/tests/unit_tests/chat_models/__snapshots__/test_bedrock_converse.ambr b/libs/aws/tests/unit_tests/chat_models/__snapshots__/test_bedrock_converse.ambr
@@ -0,0 +1,27 @@
+# serializer version: 1
+# name: TestBedrockStandard.test_serdes[serialized]
+  dict({
+    'id': list([
+      'langchain_aws',
+      'chat_models',
+      'ChatBedrockConverse',
+    ]),
+    'kwargs': dict({
+      'max_tokens': 100,
+      'model_id': 'anthropic.claude-3-sonnet-20240229-v1:0',
+      'provider': 'anthropic',
+      'region_name': 'us-west-1',
+      'stop_sequences': list([
+      ]),
+      'supports_tool_choice_values': list([
+        'auto',
+        'any',
+        'tool',
+      ]),
+      'temperature': 0.0,
+    }),
+    'lc': 1,
+    'name': 'ChatBedrockConverse',
+    'type': 'constructor',
+  })
+# ---