github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2026-06-09 10:17:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(xai): count reasoning tokens in output total (#34603)

Browse Source
This commit is contained in:
ccurme
2026-01-05 13:25:30 -05:00
committed by GitHub
parent 730a3676f8
commit 944b43dd25
2 changed files with 46 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
libs/partners/xai/langchain_xai/chat_models.py
View File

@@ -575,6 +575,21 @@ class ChatXAI(BaseChatOpenAI): # type: ignore[override]
response.citations
)
# Unlike OpenAI, xAI reports reasoning tokens < completion tokens. So we assume
# they are not counted in output tokens, and we add them here.
if (
(not self._use_responses_api({}))
and (usage_metadata := rtn.generations[0].message.usage_metadata) # type: ignore[attr-defined]
and (
reasoning_tokens := usage_metadata.get("output_token_details", {}).get(
"reasoning"
)
)
):
rtn.generations[0].message.usage_metadata["output_tokens"] += ( # type: ignore[attr-defined]
reasoning_tokens
)
return rtn
def _convert_chunk_to_generation_chunk(
@@ -609,6 +624,19 @@ class ChatXAI(BaseChatOpenAI): # type: ignore[override]
):
generation_chunk.message.additional_kwargs["citations"] = citations
# Unlike OpenAI, xAI reports reasoning tokens < completion tokens. So we assume
# they are not counted in output tokens, and we add them here.
if (
generation_chunk
and (not self._use_responses_api({}))
and (usage_metadata := generation_chunk.message.usage_metadata) # type: ignore[attr-defined]
and (
reasoning_tokens := usage_metadata.get("output_token_details", {}).get(
"reasoning"
)
)
):
generation_chunk.message.usage_metadata["output_tokens"] += reasoning_tokens # type: ignore[attr-defined]
return generation_chunk
def with_structured_output(