feat: capture streaming token usage in DeepSeekClient.stream_chat
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user