From 81a6d54fffbcd13191e120fe605f30ec7fb1cf5e Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 23 Jul 2026 11:24:31 +0800 Subject: [PATCH] feat: capture streaming token usage in DeepSeekClient.stream_chat Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- backend/app/services/llm/deepseek_client.py | 20 ++++- .../test_stream_chat_usage_capture.py | 83 +++++++++++++++++++ 2 files changed, 99 insertions(+), 4 deletions(-) create mode 100644 backend/tests/observability/test_stream_chat_usage_capture.py diff --git a/backend/app/services/llm/deepseek_client.py b/backend/app/services/llm/deepseek_client.py index d0af5cb..adb0aaa 100644 --- a/backend/app/services/llm/deepseek_client.py +++ b/backend/app/services/llm/deepseek_client.py @@ -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 diff --git a/backend/tests/observability/test_stream_chat_usage_capture.py b/backend/tests/observability/test_stream_chat_usage_capture.py new file mode 100644 index 0000000..346ece5 --- /dev/null +++ b/backend/tests/observability/test_stream_chat_usage_capture.py @@ -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