github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-01 00:49:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

baidu qianfan AiMessage with usage_metadata (#24288)

Browse Source 
add usage_metadata to qianfan AIMessage. Thanks
This commit is contained in:
Chen Xiabin 2024-07-16 21:30:50 +08:00 committed by GitHub
parent d83164f837
commit 69b1603173
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
libs/community/langchain_community/chat_models/baidu_qianfan_endpoint.py
View File

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