litellm/tests/integration/streaming/test_stream_contracts.py

import asyncio
import json
import threading
import uuid
from typing import Final

import pytest
from hypothesis import Phase, example, given, settings, strategies as st
from openai import OpenAI

from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.wire import Reply, wire_server


def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
    value: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
    return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n"


def text_stream(identity: str) -> tuple[bytes, ...]:
    usage: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}
    return (frame(identity, {"role": "assistant", "content": "Hello "}), frame(identity, {"content": "雪 café"}), frame(identity, {}, finish="stop"), b"data: " + json.dumps(usage).encode() + b"\n\n", b"data: [DONE]\n\n")


@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage")
def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage() -> None:
    import litellm

    body: Final = b"".join(text_stream("stream-partition-control"))

    @settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
    @example(cuts=tuple(range(1, len(body))))
    @example(cuts=())
    @given(cuts=st.lists(st.integers(min_value=1, max_value=len(body) - 1), max_size=35, unique=True).map(tuple))
    def check(cuts: tuple[int, ...]) -> None:
        boundaries: Final = (0, *sorted(cuts), len(body))
        pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:]))
        with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire:
            stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "partition control"}], stream=True, stream_options={"include_usage": True}, timeout=5, num_retries=0)
            try:
                chunks: Final = tuple(stream)
            finally:
                asyncio.run(stream.aclose())
            assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
            assert {chunk.id for chunk in chunks} == {"stream-partition-control"}
            assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["stop"]
            usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
            assert len(usages) == 1
            assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
            assert len(wire.drain()) == 1

    check()


@pytest.mark.covers("other.streaming.tools.fragmented_calls_keep_independent_arguments")
def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None:
    import litellm

    identity: Final = "stream-tools-control"
    deltas: Final = (
        {"role": "assistant", "tool_calls": [{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}}, {"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}}]},
        {"tool_calls": [{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}}, {"index": 0, "function": {"arguments": '{"x":1,'}}]},
        {"tool_calls": [{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}}, {"index": 1, "function": {"arguments": '"y":4}'}}]},
    )
    frames: Final = (*tuple(frame(identity, delta) for delta in deltas), frame(identity, {}, finish="tool_calls"), b"data: [DONE]\n\n")
    with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire:
        stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "tool control"}], stream=True, timeout=5, num_retries=0)
        try:
            chunks: Final = tuple(stream)
        finally:
            asyncio.run(stream.aclose())
        events: Final = tuple((choice.index, tool) for chunk in chunks for choice in chunk.choices for tool in (choice.delta.tool_calls or ()))
        for index, name, call_id, arguments in ((0, "add", "call-add", {"x": 1, "y": 2}), (1, "multiply", "call-multiply", {"x": 3, "y": 4})):
            selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index))
            assert "".join(tool.id or "" for tool in selected) == call_id
            assert "".join(tool.function.name or "" for tool in selected) == name
            assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments
        assert {tool.index for _, tool in events} == {0, 1}
        assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["tool_calls"]
        assert len(wire.drain()) == 1


@pytest.mark.covers("other.streaming.usage.client_visibility_preserves_persisted_accounting")
def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        for include in (None, False, True):
            identity: Final = "stream-usage-" + uuid.uuid4().hex
            with wire_server(lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire:
                model: Final = scenario.model(api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
                with OpenAI(api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0) as client:
                    stream: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": identity}], stream=True, **({} if include is None else {"stream_options": {"include_usage": include}}))
                    with stream:
                        chunks: Final = tuple(stream)
                assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
                assert {chunk.id for chunk in chunks} == {identity}
                usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
                assert len(usages) == (1 if include else 0)
                if include:
                    assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
                requests: Final = wire.drain()
                assert len(requests) == 1
                assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True
                rows: Final = eventually(lambda identity=identity: read_rows('SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70)
                assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4
                assert float(rows[0]["spend"]) == pytest.approx(0.019)


@pytest.mark.covers("other.streaming.fallback.empty_leading_chunk_then_disconnect_streams_fallback_with_usage_and_spend")
def test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bills_the_fallback(gateway: Gateway) -> None:
    identity: Final = "stream-empty-fallback-" + uuid.uuid4().hex
    empty_first: Final = b"data: " + json.dumps({"id": identity + "-primary", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 0, "total_tokens": 11}}).encode() + b"\n\n"
    with gateway.scenario() as scenario:
        with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=(empty_first, b":" + b"x" * 4_000_000 + b"\n\n", empty_first), abort_after=2)) as primary, wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as fallback:
            primary_model: Final = scenario.model(api_base=primary.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
            fallback_model: Final = scenario.model(api_base=fallback.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
            original_fallbacks: Final = object_value(gateway.get("/router/settings")["current_values"]).get("fallbacks")
            gateway.post("/config/update", {"router_settings": {"fallbacks": [{primary_model: [fallback_model]}]}})
            scenario.cleanups.callback(gateway.post, "/config/update", {"router_settings": {"fallbacks": original_fallbacks}})
            body: Final = {"model": primary_model, "messages": [{"role": "user", "content": identity}], "stream": True, "stream_options": {"include_usage": True}}
            with gateway.client.stream("POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response:
                lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data:"))
            assert response.status_code == 200, lines
            assert lines[-1] == "data: [DONE]", lines
            events: Final = tuple(json.loads(line.removeprefix("data:")) for line in lines[:-1])
            assert all("error" not in event for event in events), lines
            assert "".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"]) == "Hello 雪 café", lines
            usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
            assert (usages[-1]["prompt_tokens"], usages[-1]["completion_tokens"]) == (11, 4), lines
            assert tuple(json.loads(request.body)["messages"] for request in primary.drain()) == (body["messages"],)
            assert tuple(json.loads(request.body)["messages"] for request in fallback.drain()) == (body["messages"],)
            rows: Final = eventually(lambda: read_rows('SELECT spend, prompt_tokens, completion_tokens, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70)
            assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["status"]) == (11, 4, "success"), rows
            assert float(rows[0]["spend"]) == pytest.approx(0.019), rows


@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
    import litellm

    for truncated in (True, False):
        with wire_server(lambda request, truncated=truncated: Reply(content_type="text/event-stream", chunks=text_stream("stream-truncated"), abort_after=1 if truncated else None)) as wire:
            stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "truncation control"}], stream=True, timeout=5, num_retries=0)
            try:
                if truncated:
                    with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read") as failure:
                        tuple(stream)
                    assert isinstance(failure.value.original_exception, litellm.APIConnectionError)
                    assert failure.value.generated_content == "Hello "
                    assert failure.value.is_pre_first_chunk is False
                else:
                    chunks: Final = tuple(stream)
                    assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
                    assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices)
            finally:
                asyncio.run(stream.aclose())
            assert len(wire.drain()) == 1


@pytest.mark.covers("other.streaming.cancellation.closes_actual_provider_connection")
def test_client_cancellation_releases_the_actual_provider_connection() -> None:
    import litellm

    gate: Final = threading.Event()
    frames: Final = (frame("stream-cancel", {"role": "assistant", "content": "first"}), b":" + b"x" * 4_000_000 + b"\n\n", b"data: [DONE]\n\n")
    with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)) as wire:
        stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "cancellation control"}], stream=True, timeout=5, num_retries=0)
        try:
            first: Final = next(stream)
            assert first.choices[0].delta.content == "first"
        finally:
            try:
                asyncio.run(stream.aclose())
            finally:
                gate.set()
        assert wire.disconnected.get(timeout=5) == "/v1/chat/completions"
        assert len(wire.drain()) == 1