mirror of
https://gitee.com/infiniflow/ragflow.git
synced 2025-12-06 07:19:03 +08:00
fix(llm): handle None response in total_token_count_from_response (#10941)
### What problem does this PR solve? Fixes #10933 This PR fixes a `TypeError` in the Gemini model provider where the `total_token_count_from_response()` function could receive a `None` response object, causing the error: TypeError: argument of type 'NoneType' is not iterable **Root Cause:** The function attempted to use the `in` operator to check dictionary keys (lines 48, 54, 60) without first validating that `resp` was not `None`. When Gemini's `chat_streamly()` method returns `None`, this triggers the error. **Solution:** 1. Added a null check at the beginning of the function to return `0` if `resp is None` 2. Added `isinstance(resp, dict)` checks before all `in` operations to ensure type safety 3. This defensive programming approach prevents the TypeError while maintaining backward compatibility ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) ### Changes Made **File:** `rag/utils/__init__.py` - Line 36-38: Added `if resp is None: return 0` check - Line 52: Added `isinstance(resp, dict)` before `'usage' in resp` - Line 58: Added `isinstance(resp, dict)` before `'usage' in resp` - Line 64: Added `isinstance(resp, dict)` before `'meta' in resp` ### Testing - [x] Code compiles without errors - [x] Follows existing code style and conventions - [x] Change is minimal and focused on the specific issue ### Additional Notes This fix ensures robust handling of various response types from LLM providers, particularly Gemini, w --------- Signed-off-by: Zhang Zhefang <zhangzhefang@example.com>
This commit is contained in:
@@ -35,6 +35,12 @@ def num_tokens_from_string(string: str) -> int:
|
||||
return 0
|
||||
|
||||
def total_token_count_from_response(resp):
|
||||
"""
|
||||
Extract token count from LLM response in various formats.
|
||||
|
||||
Handles None responses and different response structures from various LLM providers.
|
||||
Returns 0 if token count cannot be determined.
|
||||
"""
|
||||
if resp is None:
|
||||
return 0
|
||||
|
||||
@@ -50,19 +56,19 @@ def total_token_count_from_response(resp):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if 'usage' in resp and 'total_tokens' in resp['usage']:
|
||||
if isinstance(resp, dict) and 'usage' in resp and 'total_tokens' in resp['usage']:
|
||||
try:
|
||||
return resp["usage"]["total_tokens"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if 'usage' in resp and 'input_tokens' in resp['usage'] and 'output_tokens' in resp['usage']:
|
||||
if isinstance(resp, dict) and 'usage' in resp and 'input_tokens' in resp['usage'] and 'output_tokens' in resp['usage']:
|
||||
try:
|
||||
return resp["usage"]["input_tokens"] + resp["usage"]["output_tokens"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if 'meta' in resp and 'tokens' in resp['meta'] and 'input_tokens' in resp['meta']['tokens'] and 'output_tokens' in resp['meta']['tokens']:
|
||||
if isinstance(resp, dict) and 'meta' in resp and 'tokens' in resp['meta'] and 'input_tokens' in resp['meta']['tokens'] and 'output_tokens' in resp['meta']['tokens']:
|
||||
try:
|
||||
return resp["meta"]["tokens"]["input_tokens"] + resp["meta"]["tokens"]["output_tokens"]
|
||||
except Exception:
|
||||
|
||||
Reference in New Issue
Block a user