From b235e69e868178fd74942d39fc5438d1d7eb9d9a Mon Sep 17 00:00:00 2001 From: dannykok Date: Sun, 16 Jun 2024 23:30:38 +0800 Subject: [PATCH] fix usage int value for merge_dicts --- libs/aws/langchain_aws/llms/bedrock.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/libs/aws/langchain_aws/llms/bedrock.py b/libs/aws/langchain_aws/llms/bedrock.py index 24416eac..7d1ce404 100644 --- a/libs/aws/langchain_aws/llms/bedrock.py +++ b/libs/aws/langchain_aws/llms/bedrock.py @@ -117,6 +117,12 @@ def _stream_response_to_generation_chunk( else: # chunk obj format varies with provider generation_info = {k: v for k, v in stream_response.items() if k != output_key} + if provider == "meta": + if "prompt_token_count" in generation_info: + generation_info["prompt_token_count"] = [generation_info["prompt_token_count"]] + if "generation_token_count" in generation_info: + generation_info["generation_token_count"] = [generation_info["generation_token_count"]] + return GenerationChunk( text=( stream_response[output_key]