Skip to content

fix(openai): Token usage not working when using MoonshotAI official API #6618

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Soulter merged 1 commit into from
Mar 19, 2026
Merged

fix(openai): Token usage not working when using MoonshotAI official API #6618

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions astrbot/core/provider/sources/openai_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,7 +313,8 @@ async def _query_stream(
logger.warning("Saving chunk state error: " + str(e))
if not chunk.choices:
continue
delta = chunk.choices[0].delta
choice = chunk.choices[0]
delta = choice.delta
# logger.debug(f"chunk delta: {delta}")
# handle the content delta
reasoning = self._extract_reasoning_content(chunk)
Expand All @@ -331,6 +332,11 @@ async def _query_stream(
_y = True
if chunk.usage:
llm_response.usage = self._extract_usage(chunk.usage)
elif choice_usage := getattr(choice, "usage", None):
# Workaround for some providers that only return usage in choices[].usage, e.g. MoonshotAI
# See https://github.com/AstrBotDevs/AstrBot/issues/6614
llm_response.usage = self._extract_usage(choice_usage)
state.current_completion_snapshot.usage = choice_usage
Comment on lines +335 to +339
Copy link
Copy Markdown
Contributor

@sourcery-ai sourcery-ai Bot Mar 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (bug_risk): llm_response.usage and state.current_completion_snapshot.usage now hold different types, which can be surprising.

Here llm_response.usage is a normalized TokenUsage, while state.current_completion_snapshot.usage is a raw provider choice_usage. If other code assumes snapshot.usage is a TokenUsage, this mismatch can cause type errors. Consider either storing the raw data in a separate field (e.g. raw_usage) or normalizing both to TokenUsage for consistency.

if _y:
yield llm_response

Expand Down Expand Up @@ -359,13 +365,11 @@ def _extract_reasoning_content(
reasoning_text = str(reasoning_attr)
return reasoning_text

def _extract_usage(self, usage: CompletionUsage) -> TokenUsage:
ptd = usage.prompt_tokens_details
cached = ptd.cached_tokens if ptd and ptd.cached_tokens else 0
prompt_tokens = 0 if usage.prompt_tokens is None else usage.prompt_tokens
completion_tokens = (
0 if usage.completion_tokens is None else usage.completion_tokens
)
def _extract_usage(self, usage: CompletionUsage | dict) -> TokenUsage:
ptd = getattr(usage, "prompt_tokens_details", None)
cached = getattr(ptd, "cached_tokens", 0) if ptd else 0
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
return TokenUsage(
input_other=prompt_tokens - cached,
input_cached=cached,
Comment on lines -362 to 371
Copy link
Copy Markdown
Contributor

@sourcery-ai sourcery-ai Bot Mar 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (bug_risk): The updated _extract_usage implementation can mis-handle cached_tokens and does not really support dict inputs.

Two concrete issues:

  1. cached_tokens: with cached = getattr(ptd, "cached_tokens", 0) if ptd else 0, a present-but-None value now propagates into prompt_tokens - cached and can raise TypeError. The previous implementation treated None as 0. You can keep that behavior with something like cached = (getattr(ptd, "cached_tokens", 0) or 0).

  2. dict support: although the type now allows CompletionUsage | dict, the logic only uses getattr, so a plain dict like {"prompt_tokens": 10, "completion_tokens": 5} will yield all zeros. If dicts are truly supported, branch on isinstance(usage, dict) and use usage.get(...) (including nested prompt_tokens_details) instead of getattr.

Expand Down
Loading