AIRegulation-DocAnalysis/backend/tests/observability/test_stream_chat_usage_capture.py

"""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


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