From ca320a8c30282708c13e5279039b1bf0d100b050 Mon Sep 17 00:00:00 2001 From: Stephen Hu Date: Tue, 26 Aug 2025 10:47:20 +0800 Subject: [PATCH] Refactor: for total_token_count method use if to check first. (#9707) ### What problem does this PR solve? for total_token_count method use if to check first, to improve the performance when we need to handle exception cases ### Type of change - [x] Refactoring --- rag/llm/rerank_model.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 22da0110f..a69efa7eb 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -44,14 +44,17 @@ class Base(ABC): raise NotImplementedError("Please implement encode method!") def total_token_count(self, resp): - try: - return resp.usage.total_tokens - except Exception: - pass - try: - return resp["usage"]["total_tokens"] - except Exception: - pass + if hasattr(resp, "usage") and hasattr(resp.usage, "total_tokens"): + try: + return resp.usage.total_tokens + except Exception: + pass + + if 'usage' in resp and 'total_tokens' in resp['usage']: + try: + return resp["usage"]["total_tokens"] + except Exception: + pass return 0