fix: handle None values in _extract_usage to prevent TypeError (#4244)
* fix: handle None values in _extract_usage to prevent TypeError Some LLM providers (especially API proxies) may return None for prompt_tokens and completion_tokens in the usage response. This causes a TypeError when attempting arithmetic operations. Added null checks with fallback to 0 for both prompt_tokens and completion_tokens before performing calculations. * refactor: use explicit None check and reuse cached variable - Use `is None` instead of `or 0` to avoid masking unexpected falsy values - Reuse `cached` variable for `input_cached` to avoid redundant calculation * ruff format --------- Co-authored-by: Soulter <905617992@qq.com>
This commit is contained in:
@@ -251,10 +251,14 @@ class ProviderOpenAIOfficial(Provider):
|
||||
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
|
||||
)
|
||||
return TokenUsage(
|
||||
input_other=usage.prompt_tokens - cached,
|
||||
input_cached=ptd.cached_tokens if ptd and ptd.cached_tokens else 0,
|
||||
output=usage.completion_tokens,
|
||||
input_other=prompt_tokens - cached,
|
||||
input_cached=cached,
|
||||
output=completion_tokens,
|
||||
)
|
||||
|
||||
async def _parse_openai_completion(
|
||||
|
||||
Reference in New Issue
Block a user