From f2bd0deeb3316565ca5ff90641378f9adb1f91ad Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 23 Jul 2026 13:19:23 +0800 Subject: [PATCH] feat: capture streaming token usage in QwenClient and QwenVLClient Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- backend/app/services/llm/qwen_client.py | 36 +++++++++++++++---- .../test_stream_chat_usage_capture.py | 31 ++++++++++++++++ 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/backend/app/services/llm/qwen_client.py b/backend/app/services/llm/qwen_client.py index 39d179e..aa184f7 100644 --- a/backend/app/services/llm/qwen_client.py +++ b/backend/app/services/llm/qwen_client.py @@ -140,8 +140,14 @@ class QwenClient(BaseLLMClient): max_tokens: Optional[int] = None, temperature: Optional[float] = None, **kwargs - ) -> Generator[str, None, None]: - """Stream chat for the Qwen Client instance.""" + ) -> Generator[str, None, Optional[Dict[str, int]]]: + """Stream chat for the Qwen 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: # Keep provider-specific behavior explicit so debugging stays straightforward. payload = { @@ -150,7 +156,8 @@ class QwenClient(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 # Keep provider-specific behavior explicit so debugging stays straightforward. + "stream": True, # Keep provider-specific behavior explicit so debugging stays straightforward. + "stream_options": {"include_usage": True} } # Keep provider-specific behavior explicit so debugging stays straightforward. @@ -167,6 +174,9 @@ class QwenClient(BaseLLMClient): data = json.loads(data_str) choices = data.get("choices", []) if not choices: + if data.get("usage"): + # Trailing usage-only chunk — capture it, nothing to yield. + usage = data["usage"] continue # Keep provider-specific behavior explicit so debugging stays straightforward. delta = choices[0].get("delta", {}) content = delta.get("content", "") @@ -183,6 +193,8 @@ class QwenClient(BaseLLMClient): logger.error(f"Qwen流式调用失败: {e}") yield f"[ERROR: {str(e)}]" + return usage + async def async_stream_chat( self, messages: List[Dict[str, str]], @@ -299,8 +311,14 @@ class QwenVLClient(BaseLLMClient): max_tokens: Optional[int] = None, temperature: Optional[float] = None, **kwargs - ) -> Generator[str, None, None]: - """Stream chat for the Qwen V L Client instance.""" + ) -> Generator[str, None, Optional[Dict[str, int]]]: + """Stream chat for the Qwen V L 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, @@ -308,7 +326,8 @@ class QwenVLClient(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: @@ -323,6 +342,9 @@ class QwenVLClient(BaseLLMClient): data = json.loads(data_str) choices = data.get("choices", []) if not choices: + if data.get("usage"): + # Trailing usage-only chunk — capture it, nothing to yield. + usage = data["usage"] continue # Keep provider-specific behavior explicit so debugging stays straightforward. delta = choices[0].get("delta", {}) content = delta.get("content", "") @@ -335,6 +357,8 @@ class QwenVLClient(BaseLLMClient): logger.error(f"QwenVL流式调用失败: {e}") yield f"[ERROR: {str(e)}]" + return usage + def get_available_models(self) -> List[str]: """Return available models for the Qwen V L 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 index 346ece5..8318f70 100644 --- a/backend/tests/observability/test_stream_chat_usage_capture.py +++ b/backend/tests/observability/test_stream_chat_usage_capture.py @@ -81,3 +81,34 @@ def test_deepseek_stream_chat_without_usage_chunk_returns_none(): assert chunks == ["Hi"] assert returned_usage is None + + +from app.services.llm.qwen_client import QwenClient, QwenVLClient + + +def test_qwen_stream_chat_returns_usage_from_trailing_chunk(): + """QwenClient.stream_chat() must return the trailing usage dict.""" + config = LLMConfig(provider=LLMProvider.QWEN, model="qwen3.5-flash", api_key="k", base_url="http://x/v1") + client = QwenClient(config) + usage = {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14} + client._client = _mock_streaming_client(_sse_lines("Bonjour", usage=usage)) + + chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "hi"}])) + + assert chunks == ["Bonjour"] + assert returned_usage == usage + sent_payload = client._client.stream.call_args.kwargs["json"] + assert sent_payload["stream_options"] == {"include_usage": True} + + +def test_qwen_vl_stream_chat_returns_usage_from_trailing_chunk(): + """QwenVLClient.stream_chat() must return the trailing usage dict.""" + config = LLMConfig(provider=LLMProvider.QWEN_VL, model="qwen3-vl-plus", api_key="k", base_url="http://x/v1") + client = QwenVLClient(config) + usage = {"prompt_tokens": 20, "completion_tokens": 6, "total_tokens": 26} + client._client = _mock_streaming_client(_sse_lines("Describing image", usage=usage)) + + chunks, returned_usage = _drain(client.stream_chat([{"role": "user", "content": "describe"}])) + + assert chunks == ["Describing image"] + assert returned_usage == usage