feat: capture streaming token usage in DeepSeekClient.stream_chat

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-23 11:24:31 +08:00
co-authored by Copilot
parent beddc2d976
commit 81a6d54fff
2 changed files with 99 additions and 4 deletions
+16 -4
View File
@@ -5,7 +5,7 @@ from the model response so that callers can dispatch tool invocations.
"""
import time
from typing import List, Dict, Optional
from typing import List, Dict, Optional, Generator
from loguru import logger
import httpx
@@ -130,8 +130,14 @@ class DeepSeekClient(BaseLLMClient):
max_tokens: Optional[int] = None,
temperature: Optional[float] = None,
**kwargs
):
"""Stream chat for the Deep Seek Client instance."""
) -> Generator[str, None, Optional[Dict[str, int]]]:
"""Stream chat for the Deep Seek Client instance.
Returns the trailing token-usage dict as the generator's return value
(read via StopIteration.value when manually driven with next()) when
the gateway sends one via stream_options.include_usage, else None.
"""
usage: Optional[Dict[str, int]] = None
try:
payload = {
"model": self.config.model,
@@ -139,7 +145,8 @@ class DeepSeekClient(BaseLLMClient):
"max_tokens": max_tokens or self.config.max_tokens,
"temperature": temperature or self.config.temperature,
"top_p": kwargs.get("top_p", self.config.top_p),
"stream": True
"stream": True,
"stream_options": {"include_usage": True}
}
with self._client.stream("POST", "/chat/completions", json=payload) as response:
@@ -168,6 +175,9 @@ class DeepSeekClient(BaseLLMClient):
content = delta.get("content", "")
if content:
yield content
elif data.get("usage"):
# Trailing usage-only chunk — no content to yield, just capture it.
usage = data["usage"]
except json.JSONDecodeError:
continue
@@ -178,6 +188,8 @@ class DeepSeekClient(BaseLLMClient):
logger.error(f"DeepSeek Stream调用失败: {e}")
yield ""
return usage
def get_available_models(self) -> List[str]:
"""Return available models for the Deep Seek Client instance."""
return self.SUPPORTED_MODELS