From 53cda7f307be14699d62b6ea5287443bd200a9a8 Mon Sep 17 00:00:00 2001
From: Piyush Jain <piyushjain@duck.com>
Date: Thu, 7 Nov 2024 16:50:28 -0800
Subject: [PATCH] Fixes type for latencyMs from numeric to list to prepare for
 response metadata (#275)

Fixes #223
---
 .../chat_models/bedrock_converse.py           | 11 +++++++++-
 .../chat_models/test_bedrock_converse.py      | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py
index 7982438e..c6858296 100644
--- a/libs/aws/langchain_aws/chat_models/bedrock_converse.py
+++ b/libs/aws/langchain_aws/chat_models/bedrock_converse.py
@@ -726,6 +726,15 @@ def _messages_to_bedrock(
     return bedrock_messages, bedrock_system
 
 
+def _extract_response_metadata(response: Dict[str, Any]) -> Dict[str, Any]:
+    response_metadata = response
+    # response_metadata only supports string, list or dict
+    if "metrics" in response and "latencyMs" in response["metrics"]:
+        response_metadata["metrics"]["latencyMs"] = [response["metrics"]["latencyMs"]]
+
+    return response_metadata
+
+
 def _parse_response(response: Dict[str, Any]) -> AIMessage:
     anthropic_content = _bedrock_to_anthropic(
         response.pop("output")["message"]["content"]
@@ -735,7 +744,7 @@ def _parse_response(response: Dict[str, Any]) -> AIMessage:
     return AIMessage(
         content=_str_if_single_text_block(anthropic_content),  # type: ignore[arg-type]
         usage_metadata=usage,
-        response_metadata=response,
+        response_metadata=_extract_response_metadata(response),
         tool_calls=tool_calls,
     )
 
diff --git a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py
index 37211bcd..e76719af 100644
--- a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py
+++ b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py
@@ -21,6 +21,7 @@
     _bedrock_to_anthropic,
     _camel_to_snake,
     _camel_to_snake_keys,
+    _extract_response_metadata,
     _messages_to_bedrock,
     _snake_to_camel,
     _snake_to_camel_keys,
@@ -415,3 +416,24 @@ def test_set_disable_streaming(
 ) -> None:
     llm = ChatBedrockConverse(model=model_id, region_name="us-west-2")
     assert llm.disable_streaming == disable_streaming
+
+
+def test__extract_response_metadata() -> None:
+    response = {
+        "ResponseMetadata": {
+            "RequestId": "xxxxxx",
+            "HTTPStatusCode": 200,
+            "HTTPHeaders": {
+                "date": "Wed, 06 Nov 2024 23:28:31 GMT",
+                "content-type": "application/json",
+                "content-length": "212",
+                "connection": "keep-alive",
+                "x-amzn-requestid": "xxxxx",
+            },
+            "RetryAttempts": 0,
+        },
+        "stopReason": "end_turn",
+        "metrics": {"latencyMs": 191},
+    }
+    response_metadata = _extract_response_metadata(response)
+    assert response_metadata["metrics"]["latencyMs"] == [191]