diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
index 38b795935..9e98b07c9 100644
--- a/rag/llm/chat_model.py
+++ b/rag/llm/chat_model.py
@@ -192,6 +192,7 @@ class Base(ABC):
                 e = await self._exceptions_async(e, attempt)
                 if e:
                     yield e
+        yield total_tokens
 
     def _length_stop(self, ans):
         if is_chinese([ans]):