diff --git a/tests/test_litellm/streaming_contract_fixtures.py b/tests/test_litellm/streaming_contract_fixtures.py new file mode 100644 index 00000000000..ea1c1095589 --- /dev/null +++ b/tests/test_litellm/streaming_contract_fixtures.py @@ -0,0 +1,500 @@ +import asyncio +import json +import threading +from collections.abc import AsyncIterator, Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Literal + +import anyio +import httpx +from pydantic import BaseModel, JsonValue, TypeAdapter + +from litellm.integrations.custom_logger import CustomLogger + +Provider = Literal["openai", "anthropic", "gemini", "mock", "openai_text"] +Surface = Literal["chat", "text"] +Mode = Literal["sync", "async"] +Options = Literal["omitted", "none", "empty", "hidden", "visible"] +Lifecycle = Literal["success", "failure", "close", "cancel"] +MODELS: Final = { + "openai": "gpt-6-astra", + "anthropic": "claude-sonnet-5", + "gemini": "gemini-3.8-flash", + "mock": "gpt-6-astra", + "openai_text": "davinci-002", +} +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +TOOL_ARGUMENTS: Final = '{"city":"Zürich 🐈","count":2}' +TOOL_ID: Final = "call_matrix_17" + + +@dataclass(frozen=True, slots=True) +class Case: + provider: Provider + mode: Mode + surface: Surface = "chat" + options: Options = "visible" + prompt_tokens: int | None = 11 + output_tokens: int = 7 + fragments: tuple[str, ...] = ("hello ", "world") + tool: bool = False + tool_preamble: str = "" + mock_payload: Literal["text", "response"] = "text" + lifecycle: Lifecycle = "success" + admission: int | None = None + split_bytes: int = 0 + sse_space: bool = True + close_yield: bool = False + admission_key: Literal["metadata", "litellm_metadata"] = "metadata" + reservation: Literal["count", "null", "empty", "null_count"] = "count" + model_name: str | None = None + + @property + def response_model(self) -> str: + return self.model_name or MODELS[self.provider] + + @property + def model(self) -> str: + prefix: Final = ( + "text-completion-openai" + if self.provider == "openai_text" + else ("openai" if self.provider == "mock" else self.provider) + ) + return prefix + "/" + self.response_model + + @property + def text(self) -> str: + return "".join(self.fragments) + + @property + def stream_options(self) -> dict[str, JsonValue]: + match self.options: + case "omitted": + return {} + case "none": + return {"stream_options": None} + case "empty": + return {"stream_options": {}} + case "hidden": + return {"stream_options": {"include_usage": False}} + case "visible": + return {"stream_options": {"include_usage": True}} + + +def usage(case: Case) -> dict[str, JsonValue]: + return { + "prompt_tokens": case.prompt_tokens, + "completion_tokens": case.output_tokens, + "total_tokens": (case.prompt_tokens or 0) + case.output_tokens, + } + + +def openai_chunk(case: Case, delta: dict[str, JsonValue], finish: str | None = None) -> dict[str, JsonValue]: + return { + "id": "chatcmpl-matrix", + "object": "chat.completion.chunk", + "created": 1, + "model": case.response_model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + "usage": None, + } + + +def openai_events(case: Case) -> tuple[dict[str, JsonValue], ...]: + content: Final = ( + tuple( + openai_chunk(case, {"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}) + for fragment in case.fragments + ) + if case.tool + else tuple(openai_chunk(case, {"content": fragment}) for fragment in case.fragments) + ) + initial: Final[dict[str, JsonValue]] = ( + { + "role": "assistant", + **({"content": case.tool_preamble} if case.tool_preamble else {}), + "tool_calls": [ + {"index": 0, "id": TOOL_ID, "type": "function", "function": {"name": "lookup", "arguments": ""}} + ], + } + if case.tool + else {"role": "assistant", "content": ""} + ) + trailing: Final = ( + ( + { + "id": "chatcmpl-matrix", + "object": "chat.completion.chunk", + "created": 1, + "model": case.response_model, + "choices": [], + "usage": usage(case), + }, + ) + if case.prompt_tokens is not None and case.options not in ("empty", "hidden") + else () + ) + return ( + openai_chunk(case, initial), + *content, + openai_chunk(case, {}, "tool_calls" if case.tool else "stop"), + *trailing, + ) + + +def anthropic_events(case: Case) -> tuple[dict[str, JsonValue], ...]: + block_index: Final = 1 if case.tool_preamble else 0 + block: Final[dict[str, JsonValue]] = ( + {"type": "tool_use", "id": TOOL_ID, "name": "lookup", "input": {}} + if case.tool + else {"type": "text", "text": ""} + ) + return ( + { + "type": "message_start", + "message": { + "id": "msg_matrix", + "type": "message", + "role": "assistant", + "model": case.response_model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": case.prompt_tokens, "output_tokens": min(1, case.output_tokens)}, + }, + }, + *( + ( + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": case.tool_preamble}, + }, + {"type": "content_block_stop", "index": 0}, + ) + if case.tool_preamble + else () + ), + {"type": "content_block_start", "index": block_index, "content_block": block}, + *( + { + "type": "content_block_delta", + "index": block_index, + "delta": ( + {"type": "input_json_delta", "partial_json": fragment} + if case.tool + else {"type": "text_delta", "text": fragment} + ), + } + for fragment in case.fragments + ), + {"type": "content_block_stop", "index": block_index}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use" if case.tool else "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": case.output_tokens}, + }, + {"type": "message_stop"}, + ) + + +def gemini_events(case: Case) -> tuple[dict[str, JsonValue], ...]: + parts: Final[tuple[dict[str, JsonValue], ...]] = ( + ( + *(({"text": case.tool_preamble},) if case.tool_preamble else ()), + {"functionCall": {"id": TOOL_ID, "name": "lookup", "args": JSON_OBJECT.validate_json(case.text)}}, + ) + if case.tool + else tuple({"text": fragment} for fragment in case.fragments) + ) + terminal: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"parts": [], "role": "model"}, "index": 0, "finishReason": "STOP"}], + "modelVersion": case.response_model, + "responseId": "gemini_matrix", + **( + { + "usageMetadata": { + "promptTokenCount": case.prompt_tokens, + "candidatesTokenCount": case.output_tokens, + "totalTokenCount": case.prompt_tokens + case.output_tokens, + } + } + if case.prompt_tokens is not None + else {} + ), + } + return ( + *( + { + "candidates": [{"content": {"parts": [part], "role": "model"}, "index": 0}], + "modelVersion": case.response_model, + "responseId": "gemini_matrix", + } + for part in parts + ), + terminal, + ) + + +def text_events(case: Case) -> tuple[dict[str, JsonValue], ...]: + chunks: Final = tuple( + { + "id": "cmpl-matrix", + "object": "text_completion", + "created": 1, + "model": case.response_model, + "choices": [ + { + "index": 0, + "text": fragment, + "logprobs": None, + "finish_reason": "stop" if index == len(case.fragments) else None, + } + ], + "usage": None, + } + for index, fragment in enumerate((*case.fragments, "")) + ) + return ( + *chunks, + *( + ( + { + "id": "cmpl-matrix", + "object": "text_completion", + "created": 1, + "model": case.response_model, + "choices": [], + "usage": usage(case), + }, + ) + if case.prompt_tokens is not None and case.options == "visible" + else () + ), + ) + + +def provider_events(case: Case) -> tuple[dict[str, JsonValue], ...]: + match case.provider: + case "anthropic": + return anthropic_events(case) + case "gemini": + return gemini_events(case) + case "openai_text": + return text_events(case) + case _: + return openai_events(case) + + +def nonstream_response(case: Case) -> dict[str, JsonValue]: + match case.provider: + case "anthropic": + return { + "id": "msg_matrix", + "type": "message", + "role": "assistant", + "model": case.response_model, + "content": ( + [ + { + "type": "tool_use", + "id": TOOL_ID, + "name": "lookup", + "input": JSON_OBJECT.validate_json(case.text), + } + ] + if case.tool + else [{"type": "text", "text": case.text}] + ), + "stop_reason": "tool_use" if case.tool else "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": case.prompt_tokens, "output_tokens": case.output_tokens}, + } + case "gemini": + return { + "candidates": [ + { + "content": { + "parts": ( + [ + { + "functionCall": { + "id": TOOL_ID, + "name": "lookup", + "args": JSON_OBJECT.validate_json(case.text), + } + } + ] + if case.tool + else [{"text": case.text}] + ), + "role": "model", + }, + "index": 0, + "finishReason": "STOP", + } + ], + "modelVersion": case.response_model, + "responseId": "gemini_matrix", + "usageMetadata": { + "promptTokenCount": case.prompt_tokens, + "candidatesTokenCount": case.output_tokens, + "totalTokenCount": (case.prompt_tokens or 0) + case.output_tokens, + }, + } + case "openai_text": + return { + "id": "cmpl-matrix", + "object": "text_completion", + "created": 1, + "model": case.response_model, + "choices": [{"index": 0, "text": case.text, "finish_reason": "stop", "logprobs": None}], + "usage": usage(case), + } + case _: + return { + "id": "chatcmpl-matrix", + "object": "chat.completion", + "created": 1, + "model": case.response_model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": (case.tool_preamble or None) if case.tool else case.text, + **( + { + "tool_calls": [ + { + "id": TOOL_ID, + "type": "function", + "function": {"name": "lookup", "arguments": case.text}, + } + ] + } + if case.tool + else {} + ), + }, + "finish_reason": "tool_calls" if case.tool else "stop", + } + ], + "usage": usage(case), + } + + +class Wire(httpx.SyncByteStream, httpx.AsyncByteStream): + def __init__(self, case: Case) -> None: + self.case: Final = case + self.closed: bool = False + self.waiting: Final = asyncio.Event() + self.release: Final = asyncio.Event() + self.requests: tuple[str, ...] = () + self.request_bodies: tuple[dict[str, JsonValue], ...] = () + self.frames: Final = tuple( + ( + ("event: " + str(event["type"]) + "\n" if case.provider == "anthropic" else "") + + ("data: " if case.sse_space else "data:") + + json.dumps(event, ensure_ascii=False) + + "\n\n" + ).encode() + for event in provider_events(case) + ) + + def respond(self, request: httpx.Request) -> httpx.Response: + self.requests += (str(request.url),) + body: Final = JSON_OBJECT.validate_json(request.content) + self.request_bodies += (body,) + if body.get("stream") is True or "streamGenerateContent" in str(request.url): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=self) + return httpx.Response(200, json=nonstream_response(self.case)) + + def pieces(self, frame: bytes) -> tuple[bytes, ...]: + width: Final = self.case.split_bytes + return tuple(frame[index : index + width] for index in range(0, len(frame), width)) if width else (frame,) + + def __iter__(self) -> Iterator[bytes]: + for index, frame in enumerate(self.frames): + if self.case.lifecycle == "failure" and index == 3: + raise httpx.ReadError("matrix provider disconnected after content") + yield from self.pieces(frame) + if self.case.provider in ("openai", "openai_text"): + yield (b"data: " if self.case.sse_space else b"data:") + b"[DONE]\n\n" + + async def __aiter__(self) -> AsyncIterator[bytes]: + for index, frame in enumerate(self.frames): + if self.case.lifecycle == "failure" and index == 3: + raise httpx.ReadError("matrix provider disconnected after content") + if self.case.lifecycle == "cancel" and index == 3: + self.waiting.set() + await self.release.wait() + for piece in self.pieces(frame): + yield piece + if self.case.provider in ("openai", "openai_text"): + yield (b"data: " if self.case.sse_space else b"data:") + b"[DONE]\n\n" + + def close(self) -> None: + self.closed = True + + async def aclose(self) -> None: + if self.case.close_yield: + await anyio.sleep(0) + self.closed = True + + +@dataclass(frozen=True, slots=True) +class Event: + outcome: Literal["success", "failure"] + mode: Mode + response: dict[str, JsonValue] | str + partial_usage: object + exception: str | None + + +class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[Event, ...] = () + self.arrived: Final = threading.Event() + self.lock: Final = threading.Lock() + + def record( + self, outcome: Literal["success", "failure"], mode: Mode, kwargs: Mapping[str, object], response: object + ) -> None: + payload: Final = ( + JSON_OBJECT.validate_json(response.model_dump_json()) if isinstance(response, BaseModel) else str(response) + ) + with self.lock: + self.events += ( + Event( + outcome, + mode, + payload, + kwargs.get("combined_usage_object"), + str(kwargs["exception"]) if "exception" in kwargs else None, + ), + ) + self.arrived.set() + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.record("success", "sync", kwargs, response_obj) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.record("success", "async", kwargs, response_obj) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.record("failure", "sync", kwargs, response_obj) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.record("failure", "async", kwargs, response_obj) diff --git a/tests/test_litellm/streaming_contract_matrix.md b/tests/test_litellm/streaming_contract_matrix.md new file mode 100644 index 00000000000..33f98abc924 --- /dev/null +++ b/tests/test_litellm/streaming_contract_matrix.md @@ -0,0 +1,81 @@ +# Streaming contract matrix candidate + +This suite evaluates streaming behavior through `completion`, `acompletion`, `text_completion`, and `atext_completion`. It uses actual provider parsing, response conversion, aggregation, and final SDK logging. Only provider HTTP traffic is simulated, using injected HTTP clients. `mock_response` cases exercise the SDK's own mock path + +Run from the repository root with existing test dependencies: + +```sh +LITELLM_LOCAL_MODEL_COST_MAP=True python -m pytest tests/test_litellm/test_streaming_contract_matrix.py -q -n 2 -o addopts= +LITELLM_LOCAL_MODEL_COST_MAP=True python -m pytest tests/test_litellm/test_streaming_contract_matrix.py --collect-only -q -o addopts= +``` + +The second command prints every exact parameterized case ID. A single case can be selected by appending `::test_name[parameters]` to the filename. Add `--runxfail` to reproduce a known failure without its exclusion + +No provider credentials, proxy, database, or Docker stack are needed. Placeholder credentials and injected HTTP transports handle requests; unexpected HTTP traffic is rejected. The local cost-map setting also avoids fetching registry data during import + +## Cases and observations + +| Cases | Coverage | +| --- | --- | +| `test_sync_without_running_loop` | Ordinary synchronous calls without an active event loop; all five provider paths and both SDK surfaces | +| `test_success` | OpenAI chat, Anthropic, Gemini, native OpenAI text; both SDK surfaces and modes; five usage-option states; prompt/output counts 0, 1, and larger | +| `test_mock_reservations` | Both surfaces and modes; all usage-option states; both metadata fields; absent, null, empty, null count, zero, one, and larger admission counts | +| `test_fragmented_text`, `test_mock_content` | Empty, ordinary, Unicode, and combining characters; provider content fragments; HTTP byte boundaries including inside UTF-8 and SSE frames | +| `test_sse_without_optional_space` | Standard SSE data fields without the optional space after the colon, across all four HTTP provider paths and both modes | +| `test_fragmented_tool`, `test_mock_tool`, `test_tool_with_text` | Tool ID/name/arguments, optional accompanying text, caller terminal reason, assembled callback content and usage | +| `test_midstream_failure` | OpenAI, Anthropic, Gemini; chat/text and sync/async; content observed before a transport failure; caller error and exactly one final failure event | +| `test_early_close`, `test_async_cancel` | Close after receiving content; async cancellation while awaiting more provider data followed by explicit close; no fabricated final response; cleanup before fixture teardown | +| `test_close_inside_cancelled_scope` | OpenAI/Gemini explicit close while AnyIO cancellation is active, with a transport-close await checkpoint | +| `test_exhaustion_cleanup` | Simulated transport closure after draining success through trailing usage | +| `test_native_text_bridge`, `test_stream_nonstream_parity` | Supported chat/text bridges and equivalent deterministic streaming/nonstreaming results | +| `test_tool_usage_options`, `test_mock_tool_zero_reservation` | Tool usage options with/without text, model-sensitive fallback, and admitted zero on structured mock tools | +| `test_model_specific_usage_fallback`, `test_unicode_prompt_usage`, `test_unicode_partial_usage` | Model-sensitive Unicode input/output fallback and exact partial usage after Unicode content is interrupted | +| `test_provider_usage_outranks_admission`, `test_absent_provider_usage` | Real provider counts outrank admission metadata; deterministic fallback when usage is absent | + +The callback observer is registered as a real `CustomLogger`. Tests never invoke its final callbacks directly. It records the emitted response and failure usage; tests wait up to five seconds for an event and observe a further 50 milliseconds for duplicates. Streaming assertions cover the observed final event count, outcome, content, and deterministic token totals. Parity cases compare nonstream caller content/usage and final event count/outcome; nonstream callback payload content is not asserted + +The observer uses matching SDK/consumer modes and one registered logger. Legacy string callback integrations, internal re-dispatch, and async consumption of a synchronous SDK response are outside this matrix + +OpenAI chat defaults to requesting usage for omitted/None options in this revision. Empty/false options suppress upstream usage, so those cases assert fallback counts instead of inventing provider counts that would not have arrived. Native OpenAI text requests usage only when explicitly enabled. Anthropic and Gemini usage remains available internally even when caller-visible usage is disabled. Legacy text responses can carry zero-valued usage placeholders while actual usage is hidden + +## Applicability limits + +The fixtures use one response choice and one modern function tool. The tool takes explicit city/count arguments; zero-argument functions are not modeled. Each tool stream supplies its function type; omission on every fragment is not modeled. Multiple choices, parallel tool calls, legacy `function_call`, reasoning/thinking, audio/images, annotations/citations, provider-specific metadata, and dollar-cost accounting are outside this evaluation. Token usage means prompt, completion, and total counts; token subcategories are not covered. Tool-name sanitization, exact placement/repetition of optional role fields, and the number of otherwise empty nonterminal chunks are not evaluated + +Completed-response cases verify `stop` and `tool_calls`. Token-limit truncation, safety/refusal terminal reasons, and provider-request rejection are not modeled; the failure lifecycle uses a transport interruption after content begins + +Content order is checked against the provider's emitted sequence. The suite does not alter private chunk timestamps, simulate clock jumps, or pass shuffled chunks directly to aggregation. Exact timestamp copying and model-router changes between chunks are outside this matrix + +Legacy text completion has no chat tool-call representation; tool cases use the chat surface. Gemini GenerateContent function arguments arrive as JSON objects, so its tool cases fragment HTTP bytes rather than inventing OpenAI-style argument-delta fields. Anthropic message usage is required; absence/null usage cases are limited to paths whose response contract permits them + +HTTPX is the selected transport; the alternate aiohttp backend is not evaluated. The reservation Cartesian matrix uses string mock responses. Structured mock responses have focused empty/Unicode and tool cases, including an admitted-zero tool case across all usage options; other structured reservation values are not fully crossed + +HTTP fixtures use SSE, including both permitted colon-space forms. Alternate accumulated-JSON transport modes are outside this matrix. The SSE whitespace rule follows the [event-stream specification](https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation) + +Anthropic text fixtures begin with an empty `content_block_start` and deliver text through deltas, matching the [documented streaming examples](https://platform.claude.com/docs/en/build-with-claude/streaming). Nonempty text in the start event is not covered. Model-sensitive fallback counts reflect the current registry and tokenizer dependencies; an intentional tokenizer change may require updating those expectations + +Synchronous calls use explicit close; cancellation requires an async wait. Cancellation cases measure cancellation followed by the SDK's explicit `aclose`, not automatic cleanup of an abandoned iterator. The text bridge does not expose the chat wrapper's close API, so early-close cases use chat. `mock_response` has no provider transport to fail or close and its deterministic string generator does not await provider data. Provider interruption cases therefore use actual HTTP provider paths + +Two focused cases also close within an already cancelled AnyIO scope. Their simulated transport yields during asynchronous close, so the assertions can detect cleanup interrupted by cancellation. This follows the [AnyIO finalization contract](https://anyio.readthedocs.io/en/stable/cancellation.html#finalization) + +Partial usage is asserted for async failure paths where the SDK supplies it. Sync failure callbacks currently supply no partial usage object. The suite does not invent a contract that requires one. Transport-close observations establish cleanup at the injected HTTP boundary, not live provider connection-pool behavior + +## Existing findings + +| Reference | Narrow excluded check | +| --- | --- | +| STREAM-001 | Anthropic transport closure on exhaustion/early close/cancel plus close; Gemini sync early-close transport closure | +| STREAM-002 | Exact usage when a provider explicitly reports zero and aggregation substitutes estimated usage | +| STREAM-003 | Final tool callback reason `stop` instead of `tool_calls` for Gemini and structured mock responses | +| STREAM-004 | Empty string `mock_response` sends a provider request | +| STREAM-005 | Exact cold native sync/async text OpenAI/Pydantic `MockValSer` serialization error | + +Each exclusion is reached only at the affected assertion or exact exception. Other assertions execute first where the stream can complete. STREAM-005 can interrupt the stream, so later assertions in that occurrence remain unexecuted. Its attribution to a particular package is unresolved. No production workaround or schema warming is performed + +These cases must not be reported as fully protected. Use the verbose test report or JUnit output to see exact affected combinations. The findings and local mutation evidence are provided separately from the implementation diff + +## Evaluation and integration + +The existing root unit-test target discovers `tests/test_litellm/test_*.py`, including this file. No CI schedules or new required gates are added. This is a candidate for evaluation, with local mutation validation and explicit known-behavior exclusions; measured local runtime is not a measured CircleCI runtime + +The local validation report records exact source revisions, environment, case IDs, resource samples, handpicked fault results, automated mutation identifiers/diffs, reviewed survivors, and the applicable denominator. It also records functions omitted by mutmut's decorated-function limitation. The test branch contains only the suite, fixtures, and this usage documentation diff --git a/tests/test_litellm/test_streaming_contract_matrix.py b/tests/test_litellm/test_streaming_contract_matrix.py new file mode 100644 index 00000000000..2962a5fab8a --- /dev/null +++ b/tests/test_litellm/test_streaming_contract_matrix.py @@ -0,0 +1,929 @@ +import asyncio +import time +from collections.abc import AsyncIterator, Iterator +from contextlib import asynccontextmanager, contextmanager +from dataclasses import dataclass +from typing import Final, Literal, cast + +import anyio +import httpx +import openai +import pytest +import respx +from pydantic import BaseModel, JsonValue + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai.common_utils import OpenAIError +from tests.test_litellm.streaming_contract_fixtures import ( + JSON_OBJECT, + TOOL_ARGUMENTS, + TOOL_ID, + Case, + Event, + Mode, + Options, + Provider, + Recorder, + Surface, + Wire, + nonstream_response, +) + + +@dataclass(frozen=True, slots=True) +class Session: + case: Case + wire: Wire + recorder: Recorder + client: openai.OpenAI | openai.AsyncOpenAI | HTTPHandler | AsyncHTTPHandler + + def kwargs(self, stream: bool = True) -> dict[str, object]: + return { + "model": self.case.model, + "api_key": "streaming-matrix-placeholder", + "stream": stream, + "client": self.client, + "num_retries": 0, + "max_tokens": 100, + **(self.case.stream_options if stream else {}), + **( + {"reasoning_effort": "none"} + if self.case.provider == "openai" and self.case.tool and self.case.response_model == "gpt-6-astra" + else {} + ), + **( + { + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}, "count": {"type": "integer"}}, + "required": ["city", "count"], + }, + }, + } + ] + } + if self.case.tool + else {} + ), + **( + { + "mock_response": ( + litellm.ModelResponse(**nonstream_response(self.case)) + if self.case.tool or self.case.mock_payload == "response" + else self.case.text + ) + } + if self.case.provider == "mock" + else {} + ), + **( + { + self.case.admission_key: { + "user_api_key_budget_reservation": ( + {"input_tokens": self.case.admission} + if self.case.reservation in ("count", "null_count") + else (None if self.case.reservation == "null" else {}) + ) + } + } + if self.case.admission is not None or self.case.reservation != "count" + else {} + ), + } + + async def call(self, stream: bool = True, prompt: str = "hello") -> object: + match self.case.mode, self.case.surface: + case "sync", "chat": + return litellm.completion(messages=[{"role": "user", "content": prompt}], **self.kwargs(stream)) + case "async", "chat": + return await litellm.acompletion(messages=[{"role": "user", "content": prompt}], **self.kwargs(stream)) + case "sync", "text": + return litellm.text_completion(prompt=prompt, **self.kwargs(stream)) + case "async", "text": + return await litellm.atext_completion(prompt=prompt, **self.kwargs(stream)) + + +@contextmanager +def registered_callbacks(recorder: Recorder) -> Iterator[None]: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(litellm, "callbacks", [recorder]) + for callback_list in ( + "success_callback", + "failure_callback", + "input_callback", + "_async_success_callback", + "_async_failure_callback", + "_async_input_callback", + ): + patch.setattr(litellm, callback_list, []) + yield + + +@contextmanager +def synchronous_session(case: Case) -> Iterator[Session]: + wire: Final = Wire(case) + recorder: Final = Recorder() + with ( + registered_callbacks(recorder), + httpx.Client(transport=httpx.MockTransport(wire.respond), trust_env=False) as client, + ): + sdk_client: Final = ( + openai.OpenAI(api_key="placeholder", http_client=client, max_retries=0) + if case.provider in ("openai", "openai_text", "mock") + else HTTPHandler(client=client) + ) + try: + yield Session(case, wire, recorder, sdk_client) + except OpenAIError as error: + if case.provider == "openai_text" and "'MockValSer' object is not an instance of 'SchemaSerializer'" in str( + error + ): + pytest.xfail("STREAM-005: cold native text usage fails in OpenAI/Pydantic serialization") + raise + + +@asynccontextmanager +async def session(case: Case) -> AsyncIterator[Session]: + if case.mode == "sync": + with synchronous_session(case) as run: + yield run + return + wire: Final = Wire(case) + recorder: Final = Recorder() + with registered_callbacks(recorder): + async with httpx.AsyncClient(transport=httpx.MockTransport(wire.respond), trust_env=False) as async_http_client: + if case.provider in ("openai", "openai_text", "mock"): + async_client: Final = openai.AsyncOpenAI( + api_key="placeholder", http_client=async_http_client, max_retries=0 + ) + try: + yield Session(case, wire, recorder, async_client) + except OpenAIError as error: + if ( + case.provider == "openai_text" + and "'MockValSer' object is not an instance of 'SchemaSerializer'" in str(error) + ): + pytest.xfail("STREAM-005: cold native text usage fails in OpenAI/Pydantic serialization") + raise + return + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + handler.client = async_http_client + yield Session(case, wire, recorder, handler) + + +@pytest.fixture(autouse=True) +def offline(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "telemetry", False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + with respx.mock(assert_all_called=False, assert_all_mocked=True): + yield + + +def snapshot(response: object) -> dict[str, JsonValue]: + assert isinstance(response, BaseModel), type(response) + return JSON_OBJECT.validate_json(response.model_dump_json()) + + +async def drain(response: object, mode: Mode) -> tuple[dict[str, JsonValue], ...]: + if mode == "sync": + return tuple(snapshot(chunk) for chunk in cast(Iterator[object], response)) + return tuple([snapshot(chunk) async for chunk in cast(AsyncIterator[object], response)]) + + +async def first_content(response: object, mode: Mode, surface: Surface, remaining: int = 5) -> dict[str, JsonValue]: + assert remaining > 0, "no content before lifecycle interruption" + chunk: Final = snapshot( + next(cast(Iterator[object], response)) if mode == "sync" else await anext(cast(AsyncIterator[object], response)) + ) + assert all(choice.get("finish_reason") is None for choice in objects(chunk["choices"])) + if text((chunk,), surface): + return chunk + return await first_content(response, mode, surface, remaining - 1) + + +async def final_events(recorder: Recorder) -> tuple[Event, ...]: + assert await asyncio.to_thread(recorder.arrived.wait, 5), "missing final SDK callback" + await asyncio.sleep(0.05) + return recorder.events + + +def objects(value: JsonValue) -> tuple[dict[str, JsonValue], ...]: + assert isinstance(value, list), value + assert all(isinstance(item, dict) for item in value), value + return cast(tuple[dict[str, JsonValue], ...], tuple(value)) + + +def field(value: JsonValue, key: str) -> JsonValue: + assert isinstance(value, dict), value + return value.get(key) + + +def text(chunks: tuple[dict[str, JsonValue], ...], surface: Surface) -> str: + return "".join( + str(part) + for chunk in chunks + for choice in objects(chunk["choices"]) + if (part := choice.get("text") if surface == "text" else field(choice.get("delta"), "content")) is not None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "openai_text")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +@pytest.mark.parametrize("options", ("omitted", "none", "empty", "hidden", "visible")) +@pytest.mark.parametrize("prompt_tokens", (0, 1, 37), ids=("p0", "p1", "p37")) +@pytest.mark.parametrize("output_tokens", (0, 1, 17), ids=("c0", "c1", "c17")) +async def test_success( + provider: Provider, mode: Mode, surface: Surface, options: Options, prompt_tokens: int, output_tokens: int +) -> None: + case: Final = Case( + provider=provider, + mode=mode, + surface=surface, + options=options, + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + ) + async with session(case) as run: + response: Final = await run.call() + chunks: Final = await drain(response, mode) + assert text(chunks, surface) == case.text + assert len(run.wire.requests) == 1, run.wire.requests + if provider in ("openai", "openai_text"): + body: Final = run.wire.request_bodies[0] + assert body["model"] == case.response_model and body["stream"] is True + assert body.get("stream_options") == ( + {"include_usage": True} + if provider == "openai" and options in ("omitted", "none") + else case.stream_options.get("stream_options") + ) + reasons: Final = tuple( + choice["finish_reason"] + for chunk in chunks + for choice in objects(chunk["choices"]) + if choice.get("finish_reason") is not None + ) + assert reasons == ("stop",), chunks + visible_usage: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage") is not None) + visible: Final = options == "visible" or ( + provider == "openai" and surface == "chat" and options in ("omitted", "none") + ) + if visible: + assert len(visible_usage) == 1, chunks + assert chunks[-1].get("usage") == visible_usage[0], "usage must be consumed after content and termination" + elif surface == "chat": + assert visible_usage == (), chunks + else: + assert all( + tuple(field(item, key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) == (0, 0, 0) + for item in visible_usage + ), chunks + recounted: Final = (provider == "openai" and options in ("empty", "hidden")) or ( + provider == "openai_text" and options != "visible" + ) + expected: Final = (8, 2, 10) if recounted else (prompt_tokens, output_tokens, prompt_tokens + output_tokens) + events: Final = await final_events(run.recorder) + assert len(events) == 1, events + assert events[0].outcome == "success" and events[0].mode == mode + assert isinstance(events[0].response, dict), events + assert_success_output(chunks, case, events[0]) + callback_usage: Final = tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + if (prompt_tokens == 0 or output_tokens == 0) and not recounted and callback_usage != expected: + pytest.xfail("STREAM-002: provider authoritative zero usage is recounted") + assert callback_usage == expected + if visible: + assert ( + tuple(field(visible_usage[0], key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + == expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_midstream_failure(provider: Provider, mode: Mode, surface: Surface) -> None: + case: Final = Case( + provider=provider, + mode=mode, + surface=surface, + fragments=("first ", "second ", "third ", "last"), + lifecycle="failure", + ) + async with session(case) as run: + response: Final = await run.call() + first: Final = await first_content(response, mode, surface) + assert text((first,), surface) == "first " + with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="matrix provider disconnected"): + await drain(response, mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1, events + assert events[0].outcome == "failure", events + assert "matrix provider disconnected" in str(events[0].exception), events + if mode == "async": + partial: Final = snapshot(events[0].partial_usage) + expected_partial: Final = {"openai": (8, 3, 11), "anthropic": (11, 2, 13), "gemini": (8, 4, 12)}[provider] + assert ( + tuple(partial[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + == expected_partial + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_early_close(provider: Provider, mode: Mode) -> None: + case: Final = Case(provider=provider, mode=mode, lifecycle="close") + async with session(case) as run: + response: Final = await run.call() + first: Final = await first_content(response, mode, "chat") + assert text((first,), "chat") == case.fragments[0] + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + assert isinstance(response, CustomStreamWrapper) + await response.aclose() + await asyncio.sleep(0.05) + assert run.recorder.events == (), "closing an unfinished stream fabricated a final callback" + if not run.wire.closed and (provider == "anthropic" or (provider == "gemini" and mode == "sync")): + pytest.xfail("STREAM-001: explicit close does not reach this provider transport") + assert run.wire.closed, "explicit close did not reach simulated transport" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +async def test_async_cancel(provider: Provider) -> None: + case: Final = Case( + provider=provider, mode="async", lifecycle="cancel", fragments=("first ", "second ", "third ", "last") + ) + async with session(case) as run: + response: Final = await run.call() + first: Final = await first_content(response, "async", "chat") + assert text((first,), "chat") == case.fragments[0] + consumer: Final = asyncio.create_task(drain(response, "async")) + await asyncio.wait_for(run.wire.waiting.wait(), timeout=5) + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + assert isinstance(response, CustomStreamWrapper) + await response.aclose() + await asyncio.sleep(0.05) + assert run.recorder.events == (), "cancellation fabricated a successful final response" + if provider == "anthropic" and not run.wire.closed: + pytest.xfail("STREAM-001: explicit close does not reach the Anthropic transport") + assert run.wire.closed, "cancel followed by explicit close did not reach transport" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "gemini")) +async def test_close_inside_cancelled_scope(provider: Provider) -> None: + case: Final = Case(provider=provider, mode="async", lifecycle="close", close_yield=True) + async with session(case) as run: + response: Final = await run.call() + first: Final = await first_content(response, "async", "chat") + assert text((first,), "chat") == case.fragments[0] + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + assert isinstance(response, CustomStreamWrapper) + with anyio.CancelScope() as scope: + scope.cancel() + await response.aclose() + assert run.wire.closed, "active cancellation interrupted explicit transport close" + await asyncio.sleep(0.05) + assert run.recorder.events == (), "closing an unfinished stream fabricated a final callback" + + +def assert_success_output( + chunks: tuple[dict[str, JsonValue], ...], + case: Case, + event: Event, + tool_usage: tuple[int, int, int] = (11, 7, 18), +) -> None: + assert event.outcome == "success" and event.mode == case.mode, event + assert isinstance(event.response, dict), event + assert len(objects(event.response["choices"])) == 1 + callback_choice: Final = objects(event.response["choices"])[0] + assert callback_choice.get("index") == 0 + assert event.response.get("model") == case.response_model + assert event.response.get("object") == "chat.completion" + assert all( + chunk.get("object") == ("text_completion" if case.surface == "text" else "chat.completion.chunk") + for chunk in chunks + ) + assert all(isinstance(chunk.get("created"), int) and cast(int, chunk["created"]) > 0 for chunk in chunks) + assert all(isinstance(chunk.get("id"), str) and chunk["id"] for chunk in chunks) + assert all(choice.get("index") == 0 for chunk in chunks for choice in objects(chunk["choices"])) + assert field(callback_choice.get("message"), "role") == "assistant" + assert all(chunk.get("model") == case.response_model for chunk in chunks) + assert len({str(chunk.get("id")) for chunk in chunks}) == 1 + assert event.response.get("id") == chunks[0].get("id") + if case.surface == "chat": + roles: Final = tuple( + field(choice.get("delta"), "role") + for chunk in chunks + for choice in objects(chunk["choices"]) + if field(choice.get("delta"), "role") is not None + ) + assert all(role == "assistant" for role in roles) + if case.text: + assert roles + callback_content: Final = field(callback_choice.get("message"), "content") + if case.tool: + assert (callback_content or "") == case.tool_preamble + assert text(chunks, case.surface) == case.tool_preamble + else: + assert callback_content == case.text + reasons: Final = tuple( + choice.get("finish_reason") + for chunk in chunks + for choice in objects(chunk["choices"]) + if choice.get("finish_reason") is not None + ) + assert reasons == ("tool_calls" if case.tool else "stop",), chunks + terminal_index: Final = next( + index + for index, chunk in enumerate(chunks) + if any(choice.get("finish_reason") is not None for choice in objects(chunk["choices"])) + ) + assert text(chunks[terminal_index + 1 :], case.surface) == "", "content arrived after terminal completion" + if not case.tool: + assert text(chunks, case.surface) == case.text + assert callback_choice.get("finish_reason") == reasons[0] + return + assert ( + tuple(field(event.response.get("usage"), key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + == tool_usage + ) + calls: Final = tuple( + call + for chunk in chunks + for choice in objects(chunk["choices"]) + for call in objects(field(choice.get("delta"), "tool_calls") or []) + ) + assert all(call.get("index") == 0 for call in calls) + assert all(call.get("type") in (None, "function") for call in calls) + assert tuple(call["id"] for call in calls if call.get("id")) == (TOOL_ID,), calls + assert tuple(field(call.get("function"), "name") for call in calls if field(call.get("function"), "name")) == ( + "lookup", + ) + args: Final = "".join(str(field(call.get("function"), "arguments") or "") for call in calls) + assert JSON_OBJECT.validate_json(args) == JSON_OBJECT.validate_json(case.text), args + if case.provider != "gemini": + assert args == case.text + callback_calls: Final = objects(field(callback_choice.get("message"), "tool_calls")) + assert len(callback_calls) == 1 + assert callback_calls[0].get("type") == "function" + assert callback_calls[0]["id"] == TOOL_ID + assert field(callback_calls[0].get("function"), "name") == "lookup" + callback_args: Final = field(callback_calls[0].get("function"), "arguments") + assert isinstance(callback_args, str) + assert JSON_OBJECT.validate_json(callback_args) == JSON_OBJECT.validate_json(case.text) + if case.provider != "gemini": + assert callback_args == case.text + if case.provider in ("gemini", "mock") and callback_choice.get("finish_reason") == "stop": + pytest.xfail("STREAM-003: Gemini/structured mock tool callback has stop instead of tool_calls") + assert callback_choice.get("finish_reason") == reasons[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "openai_text")) +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_sse_without_optional_space(provider: Provider, mode: Mode) -> None: + case: Final = Case(provider=provider, mode=mode, sse_space=False, fragments=("café 猫 ", "🐈 e\u0301")) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + assert isinstance(events[0].response, dict) + assert tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == (11, 7, 18) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("split_bytes", (0, 1, 13)) +@pytest.mark.parametrize("surface", ("chat", "text")) +@pytest.mark.parametrize( + "fragments", (("",), ("hello ", "world"), ("café ", "猫", " 🐈", " e\u0301")), ids=("empty", "ascii", "unicode") +) +async def test_fragmented_text( + provider: Provider, mode: Mode, split_bytes: int, fragments: tuple[str, ...], surface: Surface +) -> None: + case: Final = Case(provider=provider, mode=mode, surface=surface, fragments=fragments, split_bytes=split_bytes) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("split_bytes", (0, 1, 13)) +async def test_fragmented_tool(provider: Provider, mode: Mode, split_bytes: int) -> None: + case: Final = Case( + provider=provider, + mode=mode, + tool=True, + split_bytes=split_bytes, + fragments=(TOOL_ARGUMENTS[:9], TOOL_ARGUMENTS[9:17], TOOL_ARGUMENTS[17:]), + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_native_text_bridge(mode: Mode, surface: Surface) -> None: + case: Final = Case(provider="openai_text", mode=mode, surface=surface) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert run.wire.requests == ("https://api.openai.com/v1/completions",) + assert_success_output(chunks, case, events[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("admission", (0, 999)) +async def test_provider_usage_outranks_admission(provider: Provider, mode: Mode, admission: int) -> None: + case: Final = Case(provider=provider, mode=mode, admission=admission) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert isinstance(events[0].response, dict) + visible_usage: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage") is not None) + assert len(visible_usage) == 1 + assert tuple( + field(visible_usage[0], key) for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == (11, 7, 18) + assert tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == (11, 7, 18) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +@pytest.mark.parametrize("options", ("omitted", "none", "empty", "hidden", "visible")) +@pytest.mark.parametrize("admission_key", ("metadata", "litellm_metadata")) +@pytest.mark.parametrize( + "reservation", + ((None, "count"), (None, "null"), (None, "empty"), (None, "null_count"), (0, "count"), (1, "count"), (37, "count")), + ids=("absent", "null", "empty", "null-count", "zero", "one", "larger"), +) +async def test_mock_reservations( + mode: Mode, + surface: Surface, + options: Options, + admission_key: Literal["metadata", "litellm_metadata"], + reservation: tuple[int | None, Literal["count", "null", "empty", "null_count"]], +) -> None: + admitted, state = reservation + case: Final = Case( + provider="mock", + mode=mode, + surface=surface, + options=options, + admission=admitted, + admission_key=admission_key, + reservation=state, + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert run.wire.requests == () + assert_success_output(chunks, case, events[0]) + assert isinstance(events[0].response, dict) + expected: Final = (8, 2, 10) if admitted is None else (admitted, 20, admitted + 20) + assert ( + tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + == expected + ) + visible_usage: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage") is not None) + if options == "visible": + assert len(visible_usage) == 1 + assert ( + tuple(field(visible_usage[0], key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + == expected + ) + elif surface == "chat": + assert visible_usage == () + else: + assert all( + tuple(field(item, key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) == (0, 0, 0) + for item in visible_usage + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "mock", "openai_text")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_stream_nonstream_parity(provider: Provider, mode: Mode, surface: Surface) -> None: + case: Final = Case(provider=provider, mode=mode, surface=surface, admission=11 if provider == "mock" else None) + async with session(case) as streaming: + chunks: Final = await drain(await streaming.call(), mode) + events: Final = await final_events(streaming.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + async with session(case) as nonstream: + response: Final = snapshot(await nonstream.call(stream=False)) + choice: Final = objects(response["choices"])[0] + content: Final = choice.get("text") if surface == "text" else field(choice.get("message"), "content") + assert content == text(chunks, surface) == case.text + assert choice.get("finish_reason") == "stop" + nonstream_events: Final = await final_events(nonstream.recorder) + assert len(nonstream_events) == 1 + assert nonstream_events[0].outcome == "success" + assert isinstance(events[0].response, dict) + assert tuple( + field(response.get("usage"), key) for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +@pytest.mark.parametrize("payload", ("text", "response")) +@pytest.mark.parametrize("content", ("", "café 猫 🐈 e\u0301"), ids=("empty", "unicode")) +async def test_mock_content(mode: Mode, surface: Surface, payload: Literal["text", "response"], content: str) -> None: + case: Final = Case(provider="mock", mode=mode, surface=surface, fragments=(content,), mock_payload=payload) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + if payload == "text" and content == "" and run.wire.requests: + pytest.xfail("STREAM-004: an empty mock_response string reaches the provider transport") + assert run.wire.requests == (), "mock_response must not send a provider request" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_mock_tool(mode: Mode) -> None: + case: Final = Case(provider="mock", mode=mode, tool=True, fragments=(TOOL_ARGUMENTS,)) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert run.wire.requests == () + assert_success_output(chunks, case, events[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "mock")) +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_tool_with_text(provider: Provider, mode: Mode) -> None: + case: Final = Case( + provider=provider, + mode=mode, + tool=True, + tool_preamble="Let me check 猫. ", + fragments=(TOOL_ARGUMENTS[:9], TOOL_ARGUMENTS[9:17], TOOL_ARGUMENTS[17:]), + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("options", ("omitted", "none", "empty", "hidden", "visible")) +@pytest.mark.parametrize("model_name", ("gpt-4o-mini", "gpt-6-astra")) +@pytest.mark.parametrize("preamble", ("", "Let me check 猫. "), ids=("tool-only", "with-text")) +async def test_tool_usage_options(mode: Mode, options: Options, model_name: str, preamble: str) -> None: + case: Final = Case( + provider="openai", + mode=mode, + options=options, + model_name=model_name, + tool=True, + tool_preamble=preamble, + fragments=(TOOL_ARGUMENTS[:9], TOOL_ARGUMENTS[9:17], TOOL_ARGUMENTS[17:]), + ) + fallback_completion: Final = 42 if preamble else (36 if model_name == "gpt-4o-mini" else 35) + expected: Final = ( + (8, fallback_completion, 8 + fallback_completion) if options in ("empty", "hidden") else (11, 7, 18) + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0], tool_usage=expected) + assert len(run.wire.requests) == 1 + assert run.wire.request_bodies[0].get("stream_options") == ( + {"include_usage": True} if options in ("omitted", "none") else case.stream_options.get("stream_options") + ) + visible_usage: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage") is not None) + if options in ("empty", "hidden"): + assert visible_usage == () + else: + assert len(visible_usage) == 1 and chunks[-1].get("usage") == visible_usage[0] + assert ( + tuple(field(visible_usage[0], key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) + == expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("options", ("omitted", "none", "empty", "hidden", "visible")) +async def test_mock_tool_zero_reservation(mode: Mode, options: Options) -> None: + case: Final = Case( + provider="mock", + mode=mode, + options=options, + tool=True, + prompt_tokens=0, + admission=0, + fragments=(TOOL_ARGUMENTS,), + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert run.wire.requests == () + visible_usage: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage") is not None) + if options == "visible": + assert len(visible_usage) == 1 and chunks[-1].get("usage") == visible_usage[0] + assert tuple( + field(visible_usage[0], key) for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == (0, 7, 7) + else: + assert visible_usage == () + assert_success_output(chunks, case, events[0], tool_usage=(0, 7, 7)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "gemini")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_absent_provider_usage(provider: Provider, mode: Mode, surface: Surface) -> None: + case: Final = Case(provider=provider, mode=mode, surface=surface, prompt_tokens=None) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + assert isinstance(events[0].response, dict) + assert tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) == (8, 2, 10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_unicode_partial_usage(provider: Provider, surface: Surface) -> None: + case: Final = Case( + provider=provider, + mode="async", + surface=surface, + lifecycle="failure", + fragments=("café 猫 🐈 ", "e\u0301 Zürich ", "third ", "last"), + ) + async with session(case) as run: + response: Final = await run.call() + first: Final = await first_content(response, "async", surface) + assert text((first,), surface) == case.fragments[0] + with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="matrix provider disconnected"): + await drain(response, "async") + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert events[0].outcome == "failure" and events[0].mode == "async" + partial: Final = snapshot(events[0].partial_usage) + expected: Final = {"openai": (8, 14, 22), "anthropic": (11, 9, 20), "gemini": (8, 15, 23)}[provider] + assert tuple(partial[key] for key in ("prompt_tokens", "completion_tokens", "total_tokens")) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +@pytest.mark.parametrize("model_name,completion_tokens", (("gpt-4o-mini", 9), ("gpt-6-astra", 10))) +async def test_model_specific_usage_fallback( + mode: Mode, surface: Surface, model_name: str, completion_tokens: int +) -> None: + case: Final = Case( + provider="openai", + mode=mode, + surface=surface, + model_name=model_name, + prompt_tokens=None, + fragments=("café 猫 ", "🐈 e\u0301"), + ) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + assert isinstance(events[0].response, dict) + expected: Final = (8, completion_tokens, 8 + completion_tokens) + assert ( + tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + == expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name,prompt_tokens", (("gpt-4o-mini", 16), ("gpt-6-astra", 17))) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_unicode_prompt_usage(model_name: str, prompt_tokens: int, mode: Mode, surface: Surface) -> None: + case: Final = Case(provider="openai", mode=mode, surface=surface, model_name=model_name, prompt_tokens=None) + prompt: Final = "café 猫 🐈 e\u0301" + async with session(case) as run: + chunks: Final = await drain(await run.call(prompt=prompt), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + assert run.wire.request_bodies[0]["messages"] == [{"role": "user", "content": prompt}] + assert isinstance(events[0].response, dict) + expected: Final = (prompt_tokens, 2, prompt_tokens + 2) + for usage in (chunks[-1].get("usage"), events[0].response.get("usage")): + assert ( + tuple(field(usage, key) for key in ("prompt_tokens", "completion_tokens", "total_tokens")) == expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "openai_text")) +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("surface", ("chat", "text")) +async def test_exhaustion_cleanup(provider: Provider, mode: Mode, surface: Surface) -> None: + case: Final = Case(provider=provider, mode=mode, surface=surface) + async with session(case) as run: + chunks: Final = await drain(await run.call(), mode) + events: Final = await final_events(run.recorder) + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + if provider == "anthropic" and not run.wire.closed: + pytest.xfail("STREAM-001: Anthropic exhaustion does not close the simulated transport") + assert run.wire.closed, "exhaustion did not close the simulated transport before fixture teardown" + + +@pytest.mark.parametrize("provider", ("openai", "anthropic", "gemini", "mock", "openai_text")) +@pytest.mark.parametrize("surface", ("chat", "text")) +def test_sync_without_running_loop(provider: Provider, surface: Surface) -> None: + case: Final = Case(provider=provider, mode="sync", surface=surface, admission=11 if provider == "mock" else None) + with synchronous_session(case) as run: + response: Final = ( + litellm.completion(messages=[{"role": "user", "content": "hello"}], **run.kwargs()) + if surface == "chat" + else litellm.text_completion(prompt="hello", **run.kwargs()) + ) + chunks: Final = tuple(snapshot(chunk) for chunk in cast(Iterator[object], response)) + assert run.recorder.arrived.wait(5), "missing final SDK callback without an active event loop" + time.sleep(0.05) + events: Final = run.recorder.events + assert len(events) == 1 + assert_success_output(chunks, case, events[0]) + assert isinstance(events[0].response, dict) + expected: Final = (11, 20, 31) if provider == "mock" else (11, 7, 18) + assert ( + tuple( + field(events[0].response.get("usage"), key) + for key in ("prompt_tokens", "completion_tokens", "total_tokens") + ) + == expected + )