litellm/tests/integration/streaming/test_stream_contracts.py

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

import pytest
import yaml
from hypothesis import Phase, example, given, settings
from hypothesis import strategies as st
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, wire_server
from openai import OpenAI


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.messages_bridge.empty_choices_usage_chunk_completes_stream")
def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None:
    identity: Final = "messages-empty-choices-" + uuid.uuid4().hex
    metadata: Final = (
        b"data: "
        + json.dumps(
            {
                "id": identity,
                "object": "chat.completion.chunk",
                "created": 1,
                "model": "gpt-4o-mini",
                "choices": [],
                "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
            },
            ensure_ascii=False,
        ).encode()
        + b"\n\n"
    )
    frames: Final = (metadata, *text_stream(identity))
    with (
        wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
        gateway.scenario() as scenario,
    ):
        model: Final = scenario.model(model="azure/gpt-4o-mini", api_base=wire.url + "/v1")
        with gateway.client.stream(
            "POST",
            "/v1/messages",
            json={
                "model": model,
                "max_tokens": 64,
                "stream": True,
                "messages": [{"role": "user", "content": identity}],
            },
            headers={"Authorization": f"Bearer {gateway.key}"},
        ) as response:
            assert response.status_code == 200, response.read().decode()
            events: Final = tuple(
                json.loads(line.removeprefix("data: ")) for line in response.iter_lines() if line.startswith("data: ")
            )
    assert tuple(event["type"] for event in events) == (
        "message_start",
        "content_block_start",
        "content_block_delta",
        "content_block_delta",
        "content_block_stop",
        "message_delta",
        "message_stop",
    ), f"observed events: {events!r}"
    assert (
        "".join(event["delta"]["text"] for event in events if event["type"] == "content_block_delta") == "Hello 雪 café"
    )
    message_delta: Final = next(event for event in events if event["type"] == "message_delta")
    assert message_delta["usage"] == {"input_tokens": 11, "output_tokens": 4}
    requests: Final = wire.drain()
    assert len(requests) == 1
    outbound: Final = json.loads(requests[0].body)
    assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
        f"observed outbound body: {outbound!r}"
    )


@pytest.mark.covers("other.streaming.responses_bridge.empty_choices_chunks_complete_stream")
def test_responses_stream_completes_through_empty_choices_metadata_and_usage_chunks(gateway: Gateway) -> None:
    identity: Final = "responses-empty-choices-" + uuid.uuid4().hex
    metadata: Final = (
        b"data: "
        + json.dumps(
            {
                "id": identity,
                "object": "chat.completion.chunk",
                "created": 1,
                "model": "gpt-4o-mini",
                "choices": [],
                "prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
            },
            ensure_ascii=False,
        ).encode()
        + b"\n\n"
    )
    frames: Final = (metadata, *text_stream(identity))
    with (
        wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
        gateway.scenario() as scenario,
    ):
        model: Final = scenario.model(model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1")
        with gateway.client.stream(
            "POST",
            "/v1/responses",
            json={"model": model, "input": identity, "stream": True},
            headers={"Authorization": f"Bearer {gateway.key}"},
        ) as response:
            assert response.status_code == 200, response.read().decode()
            events: Final = tuple(
                json.loads(line.removeprefix("data: "))
                for line in response.iter_lines()
                if line.startswith("data: ") and line != "data: [DONE]"
            )
    assert (
        "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == "Hello 雪 café"
    ), f"observed events: {events!r}"
    assert tuple(event["type"] for event in events if event["type"] != "response.output_text.delta") == (
        "response.created",
        "response.in_progress",
        "response.output_item.added",
        "response.content_part.added",
        "response.output_text.done",
        "response.content_part.done",
        "response.output_item.done",
        "response.completed",
    ), f"observed events: {events!r}"
    assert events[-1]["type"] == "response.completed"
    assert events[-1]["response"]["usage"] == {
        "input_tokens": 11,
        "output_tokens": 4,
        "output_tokens_details": {"reasoning_tokens": 0, "text_tokens": 4},
        "total_tokens": 15,
    }
    requests: Final = wire.drain()
    assert len(requests) == 1
    outbound: Final = json.loads(requests[0].body)
    assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
        f"observed outbound body: {outbound!r}"
    )


def provider_cost_object_stream(identity: str, total_cost: float) -> tuple[bytes, ...]:
    cost: Final = {
        "input_tokens_cost": 0.0001,
        "output_tokens_cost": 0.0002,
        "request_cost": 0.012,
        "total_cost": total_cost,
    }
    usage: Final = {
        "id": identity,
        "object": "chat.completion.chunk",
        "created": 1,
        "model": "sonar",
        "choices": [],
        "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15, "cost": cost},
    }
    return (
        frame(identity, {"role": "assistant", "content": "Hello "}),
        frame(identity, {"content": "from search"}),
        frame(identity, {}, finish="stop"),
        b"data: " + json.dumps(usage).encode() + b"\n\n",
        b"data: [DONE]\n\n",
    )


def sse_data_lines(text: str) -> tuple[str, ...]:
    return tuple(line.removeprefix("data: ") for line in text.splitlines() if line.startswith("data: "))


@pytest.mark.covers("other.streaming.usage.provider_cost_object_completes_stream_and_bills_total_cost")
def test_perplexity_stream_with_cost_breakdown_object_completes_and_bills_total_cost(gateway: Gateway) -> None:
    identity: Final = "stream-cost-object-" + uuid.uuid4().hex
    total_cost: Final = 0.0123
    with (
        gateway.scenario() as scenario,
        wire_server(
            lambda request: Reply(
                content_type="text/event-stream", chunks=provider_cost_object_stream(identity, total_cost)
            )
        ) as wire,
    ):
        model: Final = scenario.model(model="perplexity/sonar", api_base=wire.url + "/v1")
        with gateway.client.stream(
            "POST",
            "/v1/chat/completions",
            json={
                "model": model,
                "messages": [{"role": "user", "content": identity}],
                "stream": True,
                "stream_options": {"include_usage": True},
            },
            headers={"Authorization": f"Bearer {gateway.key}"},
        ) as response:
            text: Final = response.read().decode()
        assert response.status_code == 200, text
        lines: Final = sse_data_lines(text)
        assert lines[-1] == "[DONE]", text
        events: Final = tuple(json.loads(line) for line in lines[:-1])
        assert [event for event in events if "error" in event] == [], text
        assert (
            "".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"])
            == "Hello from search"
        ), text
        assert [
            choice.get("finish_reason")
            for event in events
            for choice in event["choices"]
            if choice.get("finish_reason")
        ] == ["stop"], text
        usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
        assert len(usages) == 1, text
        assert (usages[0]["prompt_tokens"], usages[0]["completion_tokens"], usages[0]["total_tokens"]) == (11, 4, 15), (
            text
        )
        requests: Final = wire.drain()
        assert len(requests) == 1
        outbound: Final = json.loads(requests[0].body)
        assert outbound["model"] == "sonar" and outbound["stream"] is True, outbound
        assert outbound["messages"] == [{"role": "user", "content": identity}], outbound
        rows: Final = eventually(
            lambda: 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"], rows[0]["completion_tokens"]) == (11, 4)
        assert float(rows[0]["spend"]) == pytest.approx(total_cost)


@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,
    tmp_path: Path,
) -> 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 (
        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,
    ):
        config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
        config["model_list"] = [
            {
                "model_name": name,
                "litellm_params": {
                    "model": "openai/gpt-4o-mini",
                    "api_key": "synthetic-fallback-key",
                    "api_base": server.url + "/v1",
                    "input_cost_per_token": 0.001,
                    "output_cost_per_token": 0.002,
                },
            }
            for name, server in (("primary", primary), ("fallback", fallback))
        ]
        config["router_settings"] = {
            "num_retries": 0,
            "disable_cooldowns": True,
            "fallbacks": [{"primary": ["fallback"]}],
        }
        path: Final = tmp_path / "fallbacks.yaml"
        path.write_text(yaml.safe_dump(config))
        with owned_proxy(gateway, tmp_path, {}, config=path) as candidate:
            body: Final = {
                "model": "primary",
                "messages": [{"role": "user", "content": identity}],
                "stream": True,
                "stream_options": {"include_usage": True},
            }
            with candidate.client.stream(
                "POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {candidate.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()
                if request.target.endswith("/chat/completions")
            ) == (body["messages"],)
            assert tuple(
                json.loads(request.body)["messages"]
                for request in fallback.drain()
                if request.target.endswith("/chat/completions")
            ) == (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