From 63c60a31f06d3a1ae32765f03b203aba1f8850d1 Mon Sep 17 00:00:00 2001 From: Chen Xiabin <128658978+bimslab@users.noreply.github.com> Date: Thu, 18 Jul 2024 21:28:16 +0800 Subject: [PATCH] [fix] baidu qianfan AiMessage with usage_metadata (#24389) make AIMessage usage_metadata has error --- .../chat_models/baidu_qianfan_endpoint.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py b/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py index 30911995a69..8b09a473ce7 100644 --- a/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py +++ b/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py @@ -109,9 +109,9 @@ def _convert_dict_to_message(_dict: Mapping[str, Any]) -> AIMessage: content=content, additional_kwargs=msg_additional_kwargs, usage_metadata=UsageMetadata( - input_tokens=usage.prompt_tokens, - output_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, + input_tokens=usage.get("prompt_tokens", 0), + output_tokens=usage.get("completion_tokens", 0), + total_tokens=usage.get("total_tokens", 0), ), )