[fix] baidu qianfan AiMessage with usage_metadata (#24389)

make AIMessage usage_metadata has error
This commit is contained in:
Chen Xiabin 2024-07-18 21:28:16 +08:00 committed by GitHub
parent 242de9aa5e
commit 63c60a31f0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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),
),
)