From 7496484eec8d26eb83ce13c0f26ebad838e07f41 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Tue, 9 Dec 2025 18:51:54 +0800
Subject: [PATCH] Update chat_model.py

---
 rag/llm/chat_model.py | 1 +
 1 file changed, 1 insertion(+)

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]):