From 89ec86c53029aa7d412b4635a36cc5bb5eaca277 Mon Sep 17 00:00:00 2001
From: Guanchao Wang <wangcham233@gmail.com>
Date: Tue, 27 Jan 2026 20:28:19 +0800
Subject: [PATCH] fix: issue 1936 (#1937)

---
 .../provider/modelmgr/requesters/deepseekchatcmpl.py  | 11 +++++++++--
 .../modelmgr/requesters/modelscopechatcmpl.py         |  7 +++++--
 .../provider/modelmgr/requesters/moonshotchatcmpl.py  | 11 +++++++++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/langbot/pkg/provider/modelmgr/requesters/deepseekchatcmpl.py b/src/langbot/pkg/provider/modelmgr/requesters/deepseekchatcmpl.py
index a95371da..1b846161 100644
--- a/src/langbot/pkg/provider/modelmgr/requesters/deepseekchatcmpl.py
+++ b/src/langbot/pkg/provider/modelmgr/requesters/deepseekchatcmpl.py
@@ -25,7 +25,7 @@ class DeepseekChatCompletions(chatcmpl.OpenAIChatCompletions):
         use_funcs: list[resource_tool.LLMTool] = None,
         extra_args: dict[str, typing.Any] = {},
         remove_think: bool = False,
-    ) -> provider_message.Message:
+    ) -> tuple[provider_message.Message, dict]:
         self.client.api_key = use_model.provider.token_mgr.get_token()
 
         args = {}
@@ -57,4 +57,11 @@ class DeepseekChatCompletions(chatcmpl.OpenAIChatCompletions):
         # 处理请求结果
         message = await self._make_msg(resp, remove_think)
 
-        return message
+        # Extract token usage from response
+        usage_info = {}
+        if hasattr(resp, 'usage') and resp.usage:
+            usage_info['input_tokens'] = resp.usage.prompt_tokens or 0
+            usage_info['output_tokens'] = resp.usage.completion_tokens or 0
+            usage_info['total_tokens'] = resp.usage.total_tokens or 0
+
+        return message, usage_info
diff --git a/src/langbot/pkg/provider/modelmgr/requesters/modelscopechatcmpl.py b/src/langbot/pkg/provider/modelmgr/requesters/modelscopechatcmpl.py
index 0d92bb51..2ed4e3b5 100644
--- a/src/langbot/pkg/provider/modelmgr/requesters/modelscopechatcmpl.py
+++ b/src/langbot/pkg/provider/modelmgr/requesters/modelscopechatcmpl.py
@@ -130,7 +130,7 @@ class ModelScopeChatCompletions(requester.ProviderAPIRequester):
         use_funcs: list[resource_tool.LLMTool] = None,
         extra_args: dict[str, typing.Any] = {},
         remove_think: bool = False,
-    ) -> provider_message.Message:
+    ) -> tuple[provider_message.Message, dict]:
         self.client.api_key = use_model.provider.token_mgr.get_token()
 
         args = {}
@@ -162,7 +162,10 @@ class ModelScopeChatCompletions(requester.ProviderAPIRequester):
         # 处理请求结果
         message = await self._make_msg(resp)
 
-        return message
+        # ModelScope uses streaming, usage info not available
+        usage_info = {}
+
+        return message, usage_info
 
     async def _req_stream(
         self,
diff --git a/src/langbot/pkg/provider/modelmgr/requesters/moonshotchatcmpl.py b/src/langbot/pkg/provider/modelmgr/requesters/moonshotchatcmpl.py
index 969392b4..b6852963 100644
--- a/src/langbot/pkg/provider/modelmgr/requesters/moonshotchatcmpl.py
+++ b/src/langbot/pkg/provider/modelmgr/requesters/moonshotchatcmpl.py
@@ -26,7 +26,7 @@ class MoonshotChatCompletions(chatcmpl.OpenAIChatCompletions):
         use_funcs: list[resource_tool.LLMTool] = None,
         extra_args: dict[str, typing.Any] = {},
         remove_think: bool = False,
-    ) -> provider_message.Message:
+    ) -> tuple[provider_message.Message, dict]:
         self.client.api_key = use_model.provider.token_mgr.get_token()
 
         args = {}
@@ -57,4 +57,11 @@ class MoonshotChatCompletions(chatcmpl.OpenAIChatCompletions):
         # 处理请求结果
         message = await self._make_msg(resp, remove_think)
 
-        return message
+        # Extract token usage from response
+        usage_info = {}
+        if hasattr(resp, 'usage') and resp.usage:
+            usage_info['input_tokens'] = resp.usage.prompt_tokens or 0
+            usage_info['output_tokens'] = resp.usage.completion_tokens or 0
+            usage_info['total_tokens'] = resp.usage.total_tokens or 0
+
+        return message, usage_info