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 33b66b36f38..30911995a69 100644 --- a/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py +++ b/libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py @@ -32,6 +32,7 @@ from langchain_core.messages import ( SystemMessage, ToolMessage, ) +from langchain_core.messages.ai import UsageMetadata from langchain_core.output_parsers.base import OutputParserLike from langchain_core.output_parsers.openai_tools import ( JsonOutputKeyToolsParser, @@ -103,6 +104,17 @@ def _convert_dict_to_message(_dict: Mapping[str, Any]) -> AIMessage: } ] + if usage := additional_kwargs.get("usage", None): + return 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, + ), + ) + return AIMessage( content=content, additional_kwargs=msg_additional_kwargs, @@ -579,6 +591,7 @@ class QianfanChatEndpoint(BaseChatModel): content=msg.content, role="assistant", additional_kwargs=additional_kwargs, + usage_metadata=msg.usage_metadata, ), generation_info=msg.additional_kwargs, ) @@ -606,6 +619,7 @@ class QianfanChatEndpoint(BaseChatModel): content=msg.content, role="assistant", additional_kwargs=additional_kwargs, + usage_metadata=msg.usage_metadata, ), generation_info=msg.additional_kwargs, )