From ecf0322165017501f931f8193779dc25ce2c44db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E5=93=B2=E8=8A=B3?=
 <34058239+zhangzhefang-github@users.noreply.github.com>
Date: Thu, 20 Nov 2025 10:04:03 +0800
Subject: [PATCH] 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>
---
 common/token_utils.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/common/token_utils.py b/common/token_utils.py
index 29f10f7eb..4d5220fc9 100644
--- a/common/token_utils.py
+++ b/common/token_utils.py
@@ -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: