AIRegulation/AIRegulation-DocAnalysis
7
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: capture streaming token usage in QwenClient and QwenVLClient

Browse Source 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-23 13:19:23 +08:00
co-authored by Copilot
parent 81a6d54fff
commit f2bd0deeb3
2 changed files with 61 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/services/llm/qwen_client.py
+30 -6
View File
@@ -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