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
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Unit tests verifying stream_chat() captures a trailing usage-only SSE chunk.
|
||||
|
||||
Exercises DeepSeekClient, QwenClient, and QwenVLClient directly (not through
|
||||
TrackedLLMClient) by mocking the underlying httpx.Client.stream() call — none
|
||||
of these tests make a real network call.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from app.services.llm.base_client import LLMConfig, LLMProvider
|
||||
from app.services.llm.deepseek_client import DeepSeekClient
|
||||
|
||||
|
||||
def _sse_lines(*chunks: str, usage: dict | None = None) -> list[str]:
|
||||
"""Build raw SSE 'data: ...' lines the way an OpenAI-compatible gateway sends them."""
|
||||
lines = [
|
||||
f'data: {json.dumps({"choices": [{"delta": {"content": c}}]})}'
|
||||
for c in chunks
|
||||
]
|
||||
if usage is not None:
|
||||
# Trailing usage-only chunk, as sent when stream_options.include_usage=true.
|
||||
lines.append(f'data: {json.dumps({"choices": [], "usage": usage})}')
|
||||
lines.append("data: [DONE]")
|
||||
return lines
|
||||
|
||||
|
||||
def _mock_streaming_client(lines: list[str]) -> MagicMock:
|
||||
"""Build a MagicMock standing in for httpx.Client, configured for .stream()."""
|
||||
fake_response = MagicMock()
|
||||
fake_response.raise_for_status.return_value = None
|
||||
fake_response.iter_lines.return_value = lines
|
||||
|
||||
stream_cm = MagicMock()
|
||||
stream_cm.__enter__.return_value = fake_response
|
||||
stream_cm.__exit__.return_value = False
|
||||
|
||||
client = MagicMock()
|
||||
client.stream.return_value = stream_cm
|
||||
return client
|
||||
|
||||
|
||||
def _drain(gen):
|
||||
"""Manually drive a generator, returning (yielded_chunks, stop_iteration_value)."""
|
||||
chunks = []
|
||||
value = None
|
||||
while True:
|
||||
try:
|
||||
chunks.append(next(gen))
|
||||
except StopIteration as stop:
|
||||
value = stop.value
|
||||
break
|
||||
return chunks, value
|
||||
|
||||
|
||||
def test_deepseek_stream_chat_returns_usage_from_trailing_chunk():
|
||||
"""DeepSeekClient.stream_chat() must return the trailing usage dict."""
|
||||
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
|
||||
client = DeepSeekClient(config)
|
||||
usage = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
|
||||
client._client = _mock_streaming_client(_sse_lines("Hello", " world", usage=usage))
|
||||
|
||||
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
|
||||
assert chunks == ["Hello", " world"]
|
||||
assert returned_usage == usage
|
||||
# The gateway must actually be asked to include usage in the stream.
|
||||
sent_payload = client._client.stream.call_args.kwargs["json"]
|
||||
assert sent_payload["stream_options"] == {"include_usage": True}
|
||||
|
||||
|
||||
def test_deepseek_stream_chat_without_usage_chunk_returns_none():
|
||||
"""If the gateway never sends a usage chunk, the generator returns None (unchanged behavior)."""
|
||||
config = LLMConfig(provider=LLMProvider.DEEPSEEK, model="deepseek-v4-flash", api_key="k", base_url="http://x/v1")
|
||||
client = DeepSeekClient(config)
|
||||
client._client = _mock_streaming_client(_sse_lines("Hi"))
|
||||
|
||||
chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
|
||||
assert chunks == ["Hi"]
|
||||
assert returned_usage is None
|
||||
Reference in New Issue
Block a user