from __future__ import annotations import ast import contextlib import itertools import json import threading import uuid from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final, Literal, TypeAlias import anthropic import httpx import openai import pytest import yaml from _pytest.mark.structures import ParameterSet from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, string_value from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, posted_targets from integration._support.otlp_sink import Span, owned_sinks, recorded_spans from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration.providers._gemini_long_decode_support import ( API_KEY, BACKEND, CAP, CHECKPOINT, CUT, CUT_TOKEN, ENDPOINTS, FIRST_HALF, HOP_TOKEN, NO_CACHE, NOTHING_EXTRA, PLAIN, PROMPT, PROVIDERS, REFUSAL, RESUMED, SECOND_HALF, SPEND_COLUMNS, Answer, Endpoint, Hop, Provider, Seen, answered, answers, assert_billed, assert_one_id, at, carries_token, complete, cost_of, cut_off, drops_before_answering, drops_mid_body, ends_cleanly, frame, hop_reply, items_of, number, output_text, path_of, payloads, post, received, refuses, register, register_vertex, request_body, seen, single, spend_rows, sse, stream_frames, streams, summed, target, text_of, upstream, ) from pydantic import BaseModel, JsonValue, TypeAdapter Client = Literal["sdk_sync", "sdk_async", "httpx"] Failure = Literal["400", "500", "drop", "notjson", "errorframe"] Dropped = Literal["first", "continued"] Unread = Literal["dropped", "empty", "unfinished", "silent"] DeploymentTimeout = Literal["timeout", "stream_timeout"] Absent = Literal["missing", "null", "empty"] Respond: TypeAlias = Callable[[Request], Reply] _CLIENTS: Final[tuple[Client, ...]] = ("sdk_sync", "sdk_async", "httpx") _FAILURES: Final[tuple[Failure, ...]] = ("400", "500", "drop", "notjson", "errorframe") _UNREAD: Final[tuple[Unread, ...]] = ("dropped", "empty", "unfinished", "silent") _DEPLOYMENT_TIMEOUTS: Final[tuple[DeploymentTimeout, ...]] = ("timeout", "stream_timeout") _JOINED: Final = FIRST_HALF.text + SECOND_HALF.text _CHAIN_USAGE: Final = summed(FIRST_HALF, SECOND_HALF) _CHAT_MESSAGES: Final = ({"role": "user", "content": PROMPT},) _USAGE_IN_STREAM: Final = {"include_usage": True} _NATIVE_BODY: Final[Mapping[str, JsonValue]] = { "contents": [{"role": "user", "parts": [{"text": PROMPT}]}], "generationConfig": {"maxOutputTokens": CAP}, } def _mode(stream: bool) -> str: return "stream" if stream else "sync" def _cells(providers: Sequence[Provider], endpoints: Sequence[Endpoint]) -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") for provider, endpoint, stream in itertools.product(providers, endpoints, (False, True)) ) def _provider_modes() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") for provider, stream in itertools.product(PROVIDERS, (False, True)) ) def _chain_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, endpoint, stream, client, id=f"{provider}-{endpoint}-{_mode(stream)}-{client}") for provider, endpoint, stream, client in itertools.product(PROVIDERS, ENDPOINTS, (False, True), _CLIENTS) ) def _openai(gateway: Gateway) -> openai.OpenAI: return openai.OpenAI( base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0, http_client=httpx.Client(trust_env=False, timeout=60), ) def _async_openai(gateway: Gateway) -> openai.AsyncOpenAI: return openai.AsyncOpenAI( base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0, http_client=httpx.AsyncClient(trust_env=False, timeout=60), ) def _anthropic(gateway: Gateway) -> anthropic.Anthropic: return anthropic.Anthropic( base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, http_client=httpx.Client(trust_env=False, timeout=60), ) def _async_anthropic(gateway: Gateway) -> anthropic.AsyncAnthropic: return anthropic.AsyncAnthropic( base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, http_client=httpx.AsyncClient(trust_env=False, timeout=60), ) def _dumped(endpoint: Endpoint, stream: bool, models: Sequence[BaseModel]) -> Seen: texts: Final = tuple(model.model_dump_json(by_alias=True, exclude_unset=True) for model in models) return seen(endpoint, stream, tuple(JSON_OBJECT.validate_json(text) for text in texts), "\n".join(texts)) def _chat_sync(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _openai(gateway) if not stream: reply: Final = client.chat.completions.create( model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, extra_body=dict(NO_CACHE) ) return _dumped("chat", False, (reply,)) chunks: Final = tuple( client.chat.completions.create( model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, stream=True, stream_options=_USAGE_IN_STREAM, extra_body=dict(NO_CACHE), ) ) return _dumped("chat", True, chunks) async def _chat_async(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _async_openai(gateway) if not stream: reply: Final = await client.chat.completions.create( model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, extra_body=dict(NO_CACHE) ) return _dumped("chat", False, (reply,)) chunks: Final = [ chunk async for chunk in await client.chat.completions.create( model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, stream=True, stream_options=_USAGE_IN_STREAM, extra_body=dict(NO_CACHE), ) ] return _dumped("chat", True, chunks) def _messages_sync(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _anthropic(gateway) if not stream: reply: Final = client.messages.create( model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, extra_body=dict(NO_CACHE) ) return _dumped("messages", False, (reply,)) events: Final = tuple( client.messages.create( model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, stream=True, extra_body=dict(NO_CACHE) ) ) return _dumped("messages", True, events) async def _messages_async(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _async_anthropic(gateway) if not stream: reply: Final = await client.messages.create( model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, extra_body=dict(NO_CACHE) ) return _dumped("messages", False, (reply,)) events: Final = [ event async for event in await client.messages.create( model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, stream=True, extra_body=dict(NO_CACHE) ) ] return _dumped("messages", True, events) def _responses_sync(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _openai(gateway) if not stream: reply: Final = client.responses.create( model=model, input=PROMPT, max_output_tokens=CAP, extra_body=dict(NO_CACHE) ) return _dumped("responses", False, (reply,)) events: Final = tuple( client.responses.create( model=model, input=PROMPT, max_output_tokens=CAP, stream=True, extra_body=dict(NO_CACHE) ) ) return _dumped("responses", True, events) async def _responses_async(gateway: Gateway, model: str, stream: bool) -> Seen: client: Final = _async_openai(gateway) if not stream: reply: Final = await client.responses.create( model=model, input=PROMPT, max_output_tokens=CAP, extra_body=dict(NO_CACHE) ) return _dumped("responses", False, (reply,)) events: Final = [ event async for event in await client.responses.create( model=model, input=PROMPT, max_output_tokens=CAP, stream=True, extra_body=dict(NO_CACHE) ) ] return _dumped("responses", True, events) def _over_httpx(gateway: Gateway, endpoint: Endpoint, model: str, stream: bool) -> Seen: response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text return observed async def _observe(gateway: Gateway, endpoint: Endpoint, client: Client, model: str, stream: bool) -> Seen: match (endpoint, client): case ("chat", "sdk_sync"): return _chat_sync(gateway, model, stream) case ("chat", "sdk_async"): return await _chat_async(gateway, model, stream) case ("messages", "sdk_sync"): return _messages_sync(gateway, model, stream) case ("messages", "sdk_async"): return await _messages_async(gateway, model, stream) case ("responses", "sdk_sync"): return _responses_sync(gateway, model, stream) case ("responses", "sdk_async"): return await _responses_async(gateway, model, stream) case _: return _over_httpx(gateway, endpoint, model, stream) def _identities(observed: Seen, response_id: str) -> frozenset[str]: return frozenset({single(observed.ids), response_id}) @pytest.mark.parametrize(("provider", "endpoint", "stream", "client"), _chain_cells()) async def test_answer_google_ends_early_is_continued_into_one_joined_answer( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, client: Client ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = await _observe(gateway, endpoint, client, model, stream) assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed assert observed.finishes == (complete(endpoint),), observed assert observed.usage == _CHAIN_USAGE, observed assert observed.initial_prompt_tokens == (FIRST_HALF.prompt_tokens if endpoint == "chat" else None), observed assert observed.tokens == () and observed.errors == (), observed assert CHECKPOINT not in observed.raw and HOP_TOKEN not in observed.raw, observed.raw assert_one_id(endpoint, stream, observed, first_id) opening, continuing = received(wire, provider, stream, (False, True)) assert at(opening, "generationConfig", "max_output_tokens") == CAP, opening assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_answer_cut_at_the_callers_cap_returns_from_one_request_and_chat_carries_its_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: response_id: Final = f"cut-{uuid.uuid4().hex}" with wire_server(upstream(answers(response_id, CUT))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.streamed) == (CUT.text, CUT.text), observed assert observed.finishes == (cut_off(endpoint),), observed assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed assert (observed.usage, observed.initial_prompt_tokens) == (CUT.usage, None), observed assert CHECKPOINT not in observed.raw, observed.raw assert observed.tokens == ((CUT_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert (CUT_TOKEN in observed.raw) == carries_token(endpoint, stream), observed.raw assert observed.token_carriers == (observed.content_carriers[-1:] if endpoint == "chat" else ()), observed assert_one_id(endpoint, stream, observed, response_id) (only,) = received(wire, provider, stream, (False,)) assert at(only, "generationConfig", "max_output_tokens") == CAP, only assert_billed(model, endpoint, CUT.usage, _identities(observed, response_id)) _EMPTY_VERTEX_LISTS: Final[Mapping[str, JsonValue]] = MappingProxyType( { "vertex_ai_grounding_metadata": [], "vertex_ai_url_context_metadata": [], "vertex_ai_safety_ratings": [], "vertex_ai_safety_results": [], "vertex_ai_citation_metadata": [], } ) def _vertex_lists(chunk: Mapping[str, JsonValue]) -> dict[str, JsonValue]: return {key: value for key, value in chunk.items() if key.startswith("vertex_ai_")} @pytest.mark.parametrize("provider", PROVIDERS) def test_cut_chat_stream_hands_its_token_on_a_chunk_with_empty_vertex_ai_lists_and_no_reasoning_count( gateway: Gateway, provider: Provider ) -> None: with wire_server(upstream(answers(f"cut-{uuid.uuid4().hex}", CUT))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) chunks: Final = payloads(post(gateway, "chat", request_body("chat", model, True)).text, True) carriers: Final = [chunk for chunk in chunks if "continuation_token" in json.dumps(chunk)] assert [_vertex_lists(chunk) for chunk in carriers] == [dict(_EMPTY_VERTEX_LISTS)], chunks assert [_vertex_lists(chunk) for chunk in chunks if chunk not in carriers] == [{}] * (len(chunks) - 1), chunks metered: Final = [chunk for chunk in chunks if isinstance(chunk.get("usage"), dict)] assert [at(chunk, "usage", "completion_tokens_details") for chunk in metered] == [ {"text_tokens": CUT.output_tokens} ], chunks received(wire, provider, True, (False,)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_resend_with_a_continuation_token_reaches_google_and_returns_the_rest( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: resumed_id: Final = f"resumed-{uuid.uuid4().hex}" respond: Final = upstream(answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(resumed_id, RESUMED)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, stream, {**NO_CACHE, "continuation_token": CUT_TOKEN}) observed: Final = answered(endpoint, stream, post(gateway, endpoint, body)) assert (observed.answer, observed.streamed) == (RESUMED.text, RESUMED.text), observed assert observed.finishes == (complete(endpoint),), observed assert (observed.usage, observed.tokens) == (RESUMED.usage, ()), observed assert CHECKPOINT not in observed.raw, observed.raw assert_one_id(endpoint, stream, observed, resumed_id) (only,) = received(wire, provider, stream, (True,)) assert only["continuationToken"] == CUT_TOKEN and "continuation_token" not in json.dumps(only), only assert_billed(model, endpoint, RESUMED.usage, _identities(observed, resumed_id)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_answer_google_finishes_itself_stays_one_request_with_no_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: response_id: Final = f"plain-{uuid.uuid4().hex}" with wire_server(upstream(answers(response_id, PLAIN))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.streamed) == (PLAIN.text, PLAIN.text), observed assert observed.finishes == (complete(endpoint),), observed assert (observed.usage, observed.initial_prompt_tokens) == (PLAIN.usage, None), observed assert (observed.tokens, observed.token_carriers) == ((), ()), observed assert CHECKPOINT not in observed.raw, observed.raw assert_one_id(endpoint, stream, observed, response_id) (only,) = received(wire, provider, stream, (False,)) assert "continuationToken" not in only, only assert_billed(model, endpoint, PLAIN.usage, _identities(observed, response_id)) def _unavailable_mid_stream(stream: bool) -> Reply: if not stream: return refuses(503)(stream) error: Final = {"error": {"code": 503, "message": REFUSAL, "status": "UNAVAILABLE"}} return Reply(content_type="text/event-stream", body=b"data: " + json.dumps(error).encode() + b"\r\n\r\n") def _failing(failure: Failure) -> Answer: match failure: case "400": return refuses(400) case "500": return refuses(500) case "drop": return drops_before_answering case "notjson": return lambda _stream: Reply(body=b"not json", content_type="text/html") case "errorframe": return _unavailable_mid_stream def _failed_cells() -> tuple[ParameterSet, ...]: reached: Final = ( *itertools.product(PROVIDERS, ("chat",), (False, True)), *itertools.product(("gemini",), ("messages", "responses"), (False, True)), ) return tuple( pytest.param(provider, endpoint, stream, failure, id=f"{provider}-{endpoint}-{_mode(stream)}-{failure}") for (provider, endpoint, stream), failure in itertools.product(reached, _FAILURES) ) @pytest.mark.parametrize(("provider", "endpoint", "stream", "failure"), _failed_cells()) def test_continuation_request_that_fails_to_open_returns_the_answer_so_far_with_its_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, failure: Failure ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" with ( wire_server(upstream(answers(first_id, FIRST_HALF), _failing(failure))) as wire, gateway.scenario() as scenario, ): model: Final = register(gateway, scenario, provider, wire) response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed assert observed.finishes == (cut_off(endpoint),), observed assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed assert (observed.usage, observed.initial_prompt_tokens) == (FIRST_HALF.usage, None), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert (HOP_TOKEN in observed.raw) == carries_token(endpoint, stream) and observed.errors == (), observed.raw assert CHECKPOINT not in observed.raw and REFUSAL not in observed.raw, observed.raw assert len(observed.token_carriers) == (1 if endpoint == "chat" and stream else 0), observed assert observed.token_carriers[:1] >= observed.content_carriers[-1:], observed assert_one_id(endpoint, stream, observed, first_id) opening, continuing = received(wire, provider, stream, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) _SECOND_TOKEN: Final = "token-after-the-second-request" _MIDDLE: Final = Hop(("then the ", "middle, "), "CONTINUATION", _SECOND_TOKEN, 27, 5) _LAST: Final = Hop(("and then ", "the end."), "STOP", None, 32, 4) def _in_three_requests(first_id: str, last: Answer) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") if token is None: return hop_reply(first_id, FIRST_HALF, streams(request)) if token == HOP_TOKEN: return hop_reply(f"middle-{uuid.uuid4().hex}", _MIDDLE, streams(request)) assert token == _SECOND_TOKEN, request return last(streams(request)) return respond def _assert_each_request_sent_the_token_before_it(wire: Wire, provider: Provider, stream: bool) -> None: opening, second, third = received(wire, provider, stream, (False, True, True)) assert (second, third) == ( {**opening, "continuationToken": HOP_TOKEN}, {**opening, "continuationToken": _SECOND_TOKEN}, ), (opening, second, third) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_answer_google_ends_early_twice_is_joined_from_three_requests_each_sent_the_token_before_it( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" joined: Final = FIRST_HALF.text + _MIDDLE.text + _LAST.text usage: Final = summed(FIRST_HALF, _MIDDLE, _LAST) respond: Final = _in_three_requests(first_id, answers(f"last-{uuid.uuid4().hex}", _LAST)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.streamed) == (joined, joined), observed assert (observed.finishes, observed.usage) == ((complete(endpoint),), usage), observed assert observed.initial_prompt_tokens == (FIRST_HALF.prompt_tokens if endpoint == "chat" else None), observed assert observed.tokens == () and observed.errors == (), observed assert not any(token in observed.raw for token in (CHECKPOINT, HOP_TOKEN, _SECOND_TOKEN)), observed.raw assert_one_id(endpoint, stream, observed, first_id) _assert_each_request_sent_the_token_before_it(wire, provider, stream) assert_billed(model, endpoint, usage, _identities(observed, first_id)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_third_request_that_fails_to_open_returns_both_earlier_parts_with_the_second_requests_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" so_far: Final = FIRST_HALF.text + _MIDDLE.text usage: Final = summed(FIRST_HALF, _MIDDLE) with wire_server(_in_three_requests(first_id, refuses(500))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.streamed) == (so_far, so_far), observed assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), usage), observed assert observed.tokens == ((_SECOND_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert (_SECOND_TOKEN in observed.raw) == carries_token(endpoint, stream), observed.raw assert not any(leaked in observed.raw for leaked in (CHECKPOINT, HOP_TOKEN, REFUSAL)), observed.raw assert observed.errors == (), observed assert_one_id(endpoint, stream, observed, first_id) _assert_each_request_sent_the_token_before_it(wire, provider, stream) assert_billed(model, endpoint, usage, _identities(observed, first_id)) _CUTS: Final = 40 def _part(index: int) -> Hop: if index == _CUTS: return Hop((f"part {index} ", "ends it."), "STOP", None, 20, 3) return Hop((f"part {index} ", "goes on, "), "CONTINUATION", f"token-after-part-{index}", 20, 3) _PARTS: Final = tuple(_part(index) for index in range(_CUTS + 1)) _PART_AFTER: Final[Mapping[str | None, Hop]] = MappingProxyType( {None: _PARTS[0], **{earlier.token: later for earlier, later in itertools.pairwise(_PARTS)}} ) def _cut_until_the_last_part(first_id: str) -> Respond: def respond(request: Request) -> Reply: token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") assert token is None or isinstance(token, str), request response_id: Final = first_id if token is None else f"later-{uuid.uuid4().hex}" return hop_reply(response_id, _PART_AFTER[token], streams(request)) return respond @pytest.mark.parametrize(("provider", "stream"), _provider_modes()) def test_answer_google_ends_early_forty_times_is_joined_from_forty_one_requests_and_billed_once( gateway: Gateway, provider: Provider, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" joined: Final = "".join(part.text for part in _PARTS) usage: Final = summed(*_PARTS) with wire_server(_cut_until_the_last_part(first_id)) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = answered("chat", stream, post(gateway, "chat", request_body("chat", model, stream))) assert (observed.answer, observed.finishes, observed.usage) == (joined, ("stop",), usage), observed assert observed.tokens == () and observed.errors == (), observed assert_one_id("chat", stream, observed, first_id) sent: Final = received(wire, provider, stream, (False, *(True,) * _CUTS)) assert [body.get("continuationToken") for body in sent] == [None, *(part.token for part in _PARTS[:-1])] assert_billed(model, "chat", usage, _identities(observed, first_id)) _TOKEN_PATH: Final = "/_oauth/token" _EXPIRED: Final = "The scripted Google refused an access token that had expired" _REVOKED: Final = "The scripted token endpoint revoked this service account" _CALLERS_AUTHORIZATION: Final = "Bearer the-callers-own-credential" def _vertex_cells(endpoints: Sequence[Endpoint]) -> tuple[ParameterSet, ...]: return tuple( pytest.param(endpoint, stream, id=f"vertex_ai-{endpoint}-{_mode(stream)}") for endpoint, stream in itertools.product(endpoints, (False, True)) ) def _token_reply(token: str, lifetime: int) -> Reply: return Reply(body=json.dumps({"access_token": token, "token_type": "Bearer", "expires_in": lifetime}).encode()) def _issuing_tokens(lifetime: int, google: Respond) -> Respond: issued: Final = itertools.count() def respond(request: Request) -> Reply: if request.target == _TOKEN_PATH: return _token_reply(f"issued-{next(issued)}", lifetime) return google(request) return respond def _tokens_that_expire_by_the_next_request(google: Respond) -> Respond: issued: Final = itertools.count() unspent: Final[SimpleQueue[str]] = SimpleQueue() def respond(request: Request) -> Reply: if request.target == _TOKEN_PATH: token: Final = f"issued-{next(issued)}" unspent.put(token) return _token_reply(token, 0) valid: Final = frozenset(f"Bearer {unspent.get()}" for _ in range(unspent.qsize())) if request.headers.get("authorization") in valid: return google(request) error: Final = {"error": {"code": 401, "message": _EXPIRED, "status": "UNAUTHENTICATED"}} return Reply(status=401, body=json.dumps(error).encode()) return respond def _tokens_revoked_once_google_is_asked(google: Respond) -> Respond: issued: Final = itertools.count() asked: Final = threading.Event() def respond(request: Request) -> Reply: if request.target != _TOKEN_PATH: asked.set() return google(request) if not asked.is_set(): return _token_reply(f"issued-{next(issued)}", 0) refusal: Final = {"error": "invalid_grant", "error_description": _REVOKED} return Reply(status=400, body=json.dumps(refusal).encode()) return respond def _google_calls(requests: Sequence[Request], stream: bool, continued: tuple[bool, ...]) -> tuple[Request, ...]: calls: Final = tuple(request for request in requests if request.target != _TOKEN_PATH) assert [(call.method, call.target) for call in calls] == [ ("POST", target("vertex_ai", stream, resumed)) for resumed in continued ], requests return calls def _chain_through(token_endpoint: Callable[[Respond], Respond], first_id: str) -> Respond: return token_endpoint(upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF))) @pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(ENDPOINTS)) def test_vertex_continuation_request_sent_after_the_access_token_expired_carries_a_new_one_and_completes( gateway: Gateway, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" with ( wire_server(_chain_through(_tokens_that_expire_by_the_next_request, first_id)) as wire, gateway.scenario() as scenario, ): model: Final = register_vertex(scenario, wire, wire.url) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed assert (observed.finishes, observed.usage) == ((complete(endpoint),), _CHAIN_USAGE), observed assert observed.tokens == () and observed.errors == (), observed assert _EXPIRED not in observed.raw and HOP_TOKEN not in observed.raw, observed.raw assert_one_id(endpoint, stream, observed, first_id) _google_calls(wire.drain(), stream, (False, True)) assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) @pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(ENDPOINTS)) def test_vertex_access_token_that_cannot_be_renewed_before_the_continuation_returns_the_answer_so_far_with_its_token( gateway: Gateway, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" with ( wire_server(_chain_through(_tokens_revoked_once_google_is_asked, first_id)) as wire, gateway.scenario() as scenario, ): model: Final = register_vertex(scenario, wire, wire.url) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), FIRST_HALF.usage), observed assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert (HOP_TOKEN in observed.raw) == carries_token(endpoint, stream) and observed.errors == (), observed.raw assert _REVOKED not in observed.raw and CHECKPOINT not in observed.raw, observed.raw assert_one_id(endpoint, stream, observed, first_id) _google_calls(wire.drain(), stream, (False,)) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) @pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(("chat",))) def test_vertex_access_token_still_fresh_is_fetched_once_and_sent_on_every_request_of_a_continued_answer( gateway: Gateway, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" with ( wire_server(_chain_through(partial(_issuing_tokens, 3600), first_id)) as wire, gateway.scenario() as scenario, ): model: Final = register_vertex(scenario, wire, wire.url) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.finishes) == (_JOINED, (complete(endpoint),)), observed requests: Final = wire.drain() calls: Final = _google_calls(requests, stream, (False, True)) assert [request.target for request in requests].count(_TOKEN_PATH) == 1, requests assert len({call.headers["authorization"] for call in calls}) == 1, requests @pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(("chat",))) def test_vertex_deployments_own_authorization_header_is_sent_on_every_request_of_a_continued_answer( gateway: Gateway, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" with ( wire_server(_chain_through(partial(_issuing_tokens, 0), first_id)) as wire, gateway.scenario() as scenario, ): model: Final = register_vertex( scenario, wire, wire.url, extra_headers={"Authorization": _CALLERS_AUTHORIZATION} ) observed: Final = _over_httpx(gateway, endpoint, model, stream) assert (observed.answer, observed.finishes) == (_JOINED, (complete(endpoint),)), observed calls: Final = _google_calls(wire.drain(), stream, (False, True)) assert [call.headers["authorization"] for call in calls] == [_CALLERS_AUTHORIZATION] * 2, calls _SILENCE_SECONDS: Final = 5 def _silent(release: threading.Event) -> Answer: def hold(_stream: bool) -> Reply: assert release.wait(timeout=90), "The silent continuation request was never released" return Reply(drop_connection=True) return hold @pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) def test_non_streamed_continuation_request_left_unanswered_returns_the_answer_so_far_at_the_callers_timeout( gateway: Gateway, provider: Provider, endpoint: Endpoint ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" release: Final = threading.Event() with ( wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario, ): cleanup.callback(release.set) model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, False, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) observed: Final = answered(endpoint, False, post(gateway, endpoint, body)) assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, False) else ()), observed assert_one_id(endpoint, False, observed, first_id) opening, continuing = received(wire, provider, False, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) _GIVE_UP_SECONDS: Final = 2 def _held(release: threading.Event, answer: Answer) -> Answer: def hold(stream: bool) -> Reply: assert release.wait(timeout=90), "The first request was never released" return answer(stream) return hold @pytest.mark.parametrize("provider", PROVIDERS) def test_non_streamed_request_its_caller_gave_up_on_is_still_continued_and_billed_for_the_whole_answer( gateway: Gateway, provider: Provider ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" release: Final = threading.Event() respond: Final = upstream( _held(release, answers(first_id, FIRST_HALF)), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) ) with wire_server(respond) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario: cleanup.callback(release.set) model: Final = register(gateway, scenario, provider, wire) with pytest.raises(httpx.ReadTimeout): gateway.client.post( path_of("chat"), json=request_body("chat", model, False), headers={"Authorization": f"Bearer {gateway.key}"}, timeout=_GIVE_UP_SECONDS, ) release.set() assert_billed(model, "chat", _CHAIN_USAGE, frozenset({first_id})) opening, continuing = received(wire, provider, False, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) @pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) def test_streamed_continuation_request_left_unanswered_ends_the_stream_with_the_answer_so_far_at_the_callers_timeout( gateway: Gateway, provider: Provider, endpoint: Endpoint ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" release: Final = threading.Event() with ( wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario, ): cleanup.callback(release.set) model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, True, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) response: Final = post(gateway, endpoint, body) observed: Final = answered(endpoint, True, response) assert ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed assert_one_id(endpoint, True, observed, first_id) opening, continuing = received(wire, provider, True, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) @pytest.mark.parametrize( ("provider", "endpoint", "setting"), tuple(itertools.product(PROVIDERS, ENDPOINTS, _DEPLOYMENT_TIMEOUTS)), ) def test_streamed_continuation_request_left_unanswered_ends_at_the_deployments_timeout_with_the_answer_so_far( gateway: Gateway, provider: Provider, endpoint: Endpoint, setting: DeploymentTimeout ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" release: Final = threading.Event() with ( wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario, ): cleanup.callback(release.set) model: Final = register(gateway, scenario, provider, wire, settings={setting: _SILENCE_SECONDS}) response: Final = post(gateway, endpoint, request_body(endpoint, model, True, NO_CACHE)) observed: Final = answered(endpoint, True, response) assert ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed assert_one_id(endpoint, True, observed, first_id) opening, continuing = received(wire, provider, True, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_third_request_left_unanswered_returns_both_earlier_parts_and_the_second_token_at_the_callers_timeout( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" so_far: Final = FIRST_HALF.text + _MIDDLE.text usage: Final = summed(FIRST_HALF, _MIDDLE) release: Final = threading.Event() with ( wire_server(_in_three_requests(first_id, _silent(release))) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario, ): cleanup.callback(release.set) model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, stream, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) response: Final = post(gateway, endpoint, body) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.streamed) == (so_far, so_far), observed assert (observed.finishes, observed.usage, observed.errors) == ((cut_off(endpoint),), usage, ()), observed assert observed.tokens == ((_SECOND_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert_one_id(endpoint, stream, observed, first_id) _assert_each_request_sent_the_token_before_it(wire, provider, stream) assert_billed(model, endpoint, usage, _identities(observed, first_id)) def _dropping(dropped: Dropped) -> Callable[[Request], Reply]: if dropped == "first": return upstream(drops_mid_body(f"dropped-{uuid.uuid4().hex}", PLAIN)) return upstream( answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), drops_mid_body(f"dropped-{uuid.uuid4().hex}", SECOND_HALF) ) def _delivered_before_the_drop(dropped: Dropped) -> str: if dropped == "first": return PLAIN.pieces[0] return FIRST_HALF.text + SECOND_HALF.pieces[0] def _dropped_cells() -> tuple[ParameterSet, ...]: reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "messages"), ("gemini", "responses")) return tuple( pytest.param(provider, endpoint, dropped, id=f"{provider}-{endpoint}-{dropped}") for (provider, endpoint), dropped in itertools.product(reached, ("first", "continued")) ) @pytest.mark.parametrize(("provider", "endpoint", "dropped"), _dropped_cells()) def test_stream_that_drops_mid_body_reaches_the_caller_as_an_error_with_no_finish_and_no_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, dropped: Dropped ) -> None: with wire_server(_dropping(dropped)) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = post(gateway, endpoint, request_body(endpoint, model, True)) assert response.status_code == 200, response.text observed: Final = seen(endpoint, True, payloads(response.text, True), response.text) assert observed.errors != (), response.text assert observed.streamed == _delivered_before_the_drop(dropped), observed assert (observed.finishes, observed.tokens, observed.token_carriers) == ((), (), ()), observed assert CHECKPOINT not in response.text and HOP_TOKEN not in response.text, response.text bodies: Final = received(wire, provider, True, (False,) if dropped == "first" else (False, True)) assert [body.get("continuationToken") for body in bodies] == [None, HOP_TOKEN][: len(bodies)], bodies _UNFINISHED_FRAME: Final = b'data: {"candidates": [{"content": {"parts": [{"text": "never finished\n\n' def _ends_on_an_unfinished_frame(response_id: str) -> Answer: finished: Final = tuple(f"data: {json.dumps(item)}\n\n".encode() for item in stream_frames(response_id, FIRST_HALF)) return lambda _stream: Reply(content_type="text/event-stream", chunks=(*finished, _UNFINISHED_FRAME)) @pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) def test_stream_that_ends_on_an_unfinished_frame_is_still_continued_into_one_joined_answer( gateway: Gateway, provider: Provider, endpoint: Endpoint ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" respond: Final = upstream( _ends_on_an_unfinished_frame(first_id), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) ) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = _over_httpx(gateway, endpoint, model, True) assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed assert (observed.finishes, observed.usage) == ((complete(endpoint),), _CHAIN_USAGE), observed assert observed.tokens == () and observed.errors == (), observed assert_one_id(endpoint, True, observed, first_id) opening, continuing = received(wire, provider, True, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) def _opens_with_no_readable_frame(unread: Unread) -> Answer: match unread: case "dropped": return lambda stream: hop_reply(f"dropped-{uuid.uuid4().hex}", SECOND_HALF, stream, abort_after=0) case "empty": return lambda _stream: Reply(content_type="text/event-stream", chunks=()) case "unfinished": return lambda _stream: Reply(content_type="text/event-stream", chunks=(_UNFINISHED_FRAME,)) case "silent": return lambda _stream: Reply( content_type="text/event-stream", chunks=(b"\n",), gate_after_first=threading.Event() ) def _unread_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, endpoint, unread, id=f"{provider}-{endpoint}-{unread}") for provider, endpoint, unread in itertools.product(PROVIDERS, ENDPOINTS, _UNREAD) ) @pytest.mark.parametrize(("provider", "endpoint", "unread"), _unread_cells()) def test_continuation_stream_that_opens_but_never_sends_a_readable_frame_ends_with_the_answer_so_far_and_its_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, unread: Unread ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" respond: Final = upstream(answers(first_id, FIRST_HALF), _opens_with_no_readable_frame(unread)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, True, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) response: Final = post(gateway, endpoint, body) observed: Final = answered(endpoint, True, response) assert ends_cleanly(endpoint, response.text), response.text assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), FIRST_HALF.usage), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed assert observed.errors == () and CHECKPOINT not in response.text, response.text assert_one_id(endpoint, True, observed, first_id) opening, continuing = received(wire, provider, True, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) _MALFORMED: Final[tuple[ParameterSet, ...]] = ( pytest.param(12345, False, id="int-sync"), pytest.param(12345, True, id="int-stream"), pytest.param([CUT_TOKEN], False, id="list-sync"), pytest.param({"token": CUT_TOKEN}, False, id="object-sync"), pytest.param(True, False, id="bool-sync"), ) _WRONG_TYPE: Final = "Invalid value at 'continuation_token' (TYPE_STRING)" def _strict_about_the_token(resumed_id: str) -> Callable[[Request], Reply]: refused: Final = {"error": {"code": 400, "message": _WRONG_TYPE, "status": "INVALID_ARGUMENT"}} def respond(request: Request) -> Reply: body: Final = JSON_OBJECT.validate_json(request.body) if "continuationToken" not in body: return hop_reply(f"restarted-{uuid.uuid4().hex}", PLAIN, streams(request)) if not isinstance(body["continuationToken"], str): return Reply(status=400, body=json.dumps(refused).encode()) return hop_reply(resumed_id, RESUMED, streams(request)) return respond @pytest.mark.parametrize(("token", "stream"), _MALFORMED) def test_continuation_token_of_the_wrong_type_is_refused_and_never_restarts_the_answer( gateway: Gateway, token: JsonValue, stream: bool ) -> None: with wire_server(_strict_about_the_token(f"resumed-{uuid.uuid4().hex}")) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) refused: Final = post( gateway, "chat", request_body("chat", model, stream, {**NO_CACHE, "continuation_token": token}) ) assert refused.status_code == 400, refused.text assert "continuation_token" in refused.text and PLAIN.text not in refused.text, refused.text sent: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain()) assert len(sent) <= 1 and all(body.get("continuationToken") == token for body in sent), sent healthy: Final = answered("chat", stream, post(gateway, "chat", request_body("chat", model, stream))) assert (healthy.answer, healthy.finishes) == (PLAIN.text, ("stop",)), healthy @pytest.mark.parametrize("provider", PROVIDERS) def test_five_kilobyte_continuation_token_reaches_google_byte_for_byte(gateway: Gateway, provider: Provider) -> None: token: Final = f"{uuid.uuid4().hex}-" + "t" * 5000 resumed_id: Final = f"resumed-{uuid.uuid4().hex}" with wire_server(_strict_about_the_token(resumed_id)) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body("chat", model, False, {**NO_CACHE, "continuation_token": token}) observed: Final = answered("chat", False, post(gateway, "chat", body)) assert (observed.answer, observed.finishes, observed.tokens) == (RESUMED.text, ("stop",), ()), observed (only,) = received(wire, provider, False, (True,)) assert only["continuationToken"] == token, only def test_continuation_token_sent_twice_in_one_body_reaches_google_once(gateway: Gateway) -> None: tokens: Final = (f"first-{uuid.uuid4().hex}", f"second-{uuid.uuid4().hex}") with wire_server(_strict_about_the_token(f"resumed-{uuid.uuid4().hex}")) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) whole: Final = json.dumps(request_body("chat", model, False)) raw: Final = f'{whole[:-1]}, "continuation_token": "{tokens[0]}", "continuation_token": "{tokens[1]}"}}' response: Final = gateway.client.request( "POST", "/v1/chat/completions", content=raw.encode(), headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) observed: Final = answered("chat", False, response) assert observed.answer == RESUMED.text, observed (only,) = received(wire, "gemini", False, (True,)) assert only["continuationToken"] in tokens, only def _absent(kind: Absent) -> Mapping[str, JsonValue]: match kind: case "missing": return NO_CACHE case "null": return {**NO_CACHE, "continuation_token": None} case "empty": return {**NO_CACHE, "continuation_token": ""} def _absent_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(kind, provider, stream, id=f"{kind}-{provider}-{_mode(stream)}") for kind, provider, stream in itertools.product(("missing", "null", "empty"), PROVIDERS, (False, True)) ) @pytest.mark.parametrize(("kind", "provider", "stream"), _absent_cells()) def test_missing_null_and_empty_continuation_token_all_send_a_request_with_no_token( gateway: Gateway, kind: Absent, provider: Provider, stream: bool ) -> None: response_id: Final = f"plain-{uuid.uuid4().hex}" with wire_server(upstream(answers(response_id, PLAIN))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = answered( "chat", stream, post(gateway, "chat", request_body("chat", model, stream, _absent(kind))) ) assert (observed.answer, observed.finishes, observed.tokens) == (PLAIN.text, ("stop",), ()), observed assert_one_id("chat", stream, observed, response_id) (only,) = received(wire, provider, stream, (False,)) assert "continuationToken" not in only and "continuation_token" not in json.dumps(only), only _DRAFTS: Final = (("Draft A", "token-of-draft-a"), ("Draft B", "token-of-draft-b")) def _two_drafts(request: Request) -> Reply: assert not streams(request), request.target candidates: Final[list[JsonValue]] = [ { "content": {"role": "model", "parts": [{"text": text}]}, "index": index, "finishReason": "CONTINUATION", "continuationToken": token, } for index, (text, token) in enumerate(_DRAFTS) ] body: Final = { "responseId": f"drafts-{uuid.uuid4().hex}", "candidates": candidates, "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 14, "totalTokenCount": 34}, "modelVersion": BACKEND, } return Reply(body=json.dumps(body).encode()) @pytest.mark.parametrize("provider", PROVIDERS) def test_request_for_two_candidates_is_never_continued_and_each_choice_keeps_its_own_token( gateway: Gateway, provider: Provider ) -> None: with wire_server(_two_drafts) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = post(gateway, "chat", request_body("chat", model, False, {**NO_CACHE, "n": 2})) assert response.status_code == 200, response.text choices: Final = items_of(JSON_OBJECT.validate_json(response.content).get("choices")) assert [ ( at(choice, "message", "content"), at(choice, "finish_reason"), at(choice, "message", "provider_specific_fields", "continuation_token"), ) for choice in choices ] == [(text, "length", token) for text, token in _DRAFTS], response.text (only,) = received(wire, provider, False, (False,)) assert at(only, "generationConfig", "candidate_count") == 2, only def _uncapped_cells() -> tuple[ParameterSet, ...]: reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "responses")) return tuple( pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") for (provider, endpoint), stream in itertools.product(reached, (False, True)) ) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _uncapped_cells()) def test_request_with_no_output_cap_is_never_continued_and_returns_googles_cut_with_its_token( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) observed: Final = answered( endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream, cap=None)) ) assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed assert observed.finishes == (cut_off(endpoint),), observed assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed assert observed.token_carriers == (observed.content_carriers[-1:] if endpoint == "chat" else ()), observed assert CHECKPOINT not in observed.raw, observed.raw (only,) = received(wire, provider, stream, (False,)) assert at(only, "generationConfig", "max_output_tokens") is None, only _NOT_A_NUMBER: Final[tuple[tuple[str, Mapping[str, JsonValue], int | None], ...]] = ( ("n-string", {"n": "1"}, CAP), ("n-list", {"n": [1]}, CAP), ("n-empty", {"n": ""}, CAP), ("n-5kb", {"n": "1" * 5120}, CAP), ("max_tokens-string", {"max_tokens": "300"}, None), ("max_tokens-list", {"max_tokens": [300]}, None), ("max_tokens-empty", {"max_tokens": ""}, None), ("max_tokens-5kb", {"max_tokens": "3" * 5120}, None), ) _NOT_A_NUMBER_CELLS: Final[tuple[ParameterSet, ...]] = tuple( pytest.param(extra, cap, stream, id=f"{name}-{_mode(stream)}") for (name, extra, cap), stream in itertools.product(_NOT_A_NUMBER, (False, True)) ) def _halves() -> Respond: return upstream( answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) ) @pytest.mark.parametrize(("extra", "cap", "stream"), _NOT_A_NUMBER_CELLS) def test_count_or_cap_that_is_not_a_number_makes_one_request_and_is_never_continued( gateway: Gateway, extra: Mapping[str, JsonValue], cap: int | None, stream: bool ) -> None: with wire_server(_halves()) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) body: Final = request_body("chat", model, stream, {**NO_CACHE, **extra}, cap=cap) observed: Final = answered("chat", stream, post(gateway, "chat", body)) assert (observed.answer, observed.errors) == (FIRST_HALF.text, ()), observed (only,) = received(wire, "gemini", stream, (False,)) assert "continuationToken" not in only, only def _twice_cells() -> tuple[ParameterSet, ...]: pairs: Final = ( ("max_tokens-number-last", f'"max_tokens": "300", "max_tokens": {CAP}', True), ("max_tokens-string-last", f'"max_tokens": {CAP}, "max_tokens": "300"', False), ("n-one-last", f'"max_tokens": {CAP}, "n": 2, "n": 1', True), ("n-string-last", f'"max_tokens": {CAP}, "n": 1, "n": "1"', False), ) return tuple( pytest.param(pair, continued, stream, id=f"{name}-{_mode(stream)}") for (name, pair, continued), stream in itertools.product(pairs, (False, True)) ) @pytest.mark.parametrize(("pairs", "continued", "stream"), _twice_cells()) def test_count_or_cap_sent_twice_in_one_body_is_read_from_its_last_value( gateway: Gateway, pairs: str, continued: bool, stream: bool ) -> None: with wire_server(_halves()) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) whole: Final = json.dumps(request_body("chat", model, stream, cap=None)) response: Final = gateway.client.request( "POST", path_of("chat"), content=f"{whole[:-1]}, {pairs}}}".encode(), headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) observed: Final = answered("chat", stream, response) assert (observed.answer, observed.errors) == (_JOINED if continued else FIRST_HALF.text, ()), observed received(wire, "gemini", stream, (False, True) if continued else (False,)) @pytest.mark.parametrize(("provider", "stream"), _provider_modes()) def test_count_of_true_is_read_as_one_answer_and_continued(gateway: Gateway, provider: Provider, stream: bool) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body("chat", model, stream, {**NO_CACHE, "n": True}) observed: Final = answered("chat", stream, post(gateway, "chat", body)) assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed assert (observed.tokens, observed.errors) == ((), ()), observed opening, continuing = received(wire, provider, stream, (False, True)) assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) assert_billed(model, "chat", _CHAIN_USAGE, _identities(observed, first_id)) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) def test_cache_hit_twin_of_a_continued_answer_replays_it_and_bills_the_chain_once( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" replayed: Final = (_JOINED, (complete(endpoint),), _CHAIN_USAGE) respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") first: Final = answered(endpoint, stream, post(gateway, endpoint, body)) assert (first.answer, first.finishes, first.usage) == replayed, first assert len(wire.drain()) == 2 assert_billed(model, endpoint, _CHAIN_USAGE, _identities(first, first_id)) second: Final = answered(endpoint, stream, post(gateway, endpoint, body)) assert (second.answer, second.finishes, second.usage) == replayed and second.tokens == (), second assert CHECKPOINT not in second.raw and HOP_TOKEN not in second.raw, second.raw assert wire.drain() == () rows: Final = spend_rows(model, expected=2) (served_from_cache,) = [row for row in rows if str(row["cache_hit"]) == "True"] (sent_to_google,) = [row for row in rows if str(row["cache_hit"]) != "True"] assert number(sent_to_google["spend"]) == pytest.approx(cost_of(_CHAIN_USAGE)), rows assert (served_from_cache["status"], number(served_from_cache["spend"])) == ("success", 0), rows def _rest_for(token: str) -> Hop: return Hop((f"the rest for {token} ", "is here."), "STOP", None, 27, 5) def _resumes_by_token(request: Request) -> Reply: token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") hop: Final = _rest_for(token) if isinstance(token, str) else PLAIN return hop_reply(f"answer-{uuid.uuid4().hex}", hop, streams(request)) def test_request_with_a_token_is_never_served_the_cached_answer_of_one_without_it_or_with_another( gateway: Gateway, ) -> None: mine: Final = f"mine-{uuid.uuid4().hex}" other: Final = f"other-{uuid.uuid4().hex}" with wire_server(_resumes_by_token) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) body: Final = request_body("chat", model, False, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") fresh: Final = answered("chat", False, post(gateway, "chat", body)) spend_rows(model) cached: Final = answered("chat", False, post(gateway, "chat", body)) assert (fresh.answer, cached.answer, len(wire.drain())) == (PLAIN.text, PLAIN.text, 1) resumed: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": mine})) assert resumed.answer == _rest_for(mine).text, resumed (sent,) = received(wire, "gemini", False, (True,)) assert sent["continuationToken"] == mine, sent spend_rows(model, expected=3) replayed: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": mine})) assert (replayed.answer, wire.drain()) == (_rest_for(mine).text, ()), replayed another: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": other})) assert another.answer == _rest_for(other).text, another (crossed,) = received(wire, "gemini", False, (True,)) assert crossed["continuationToken"] == other, crossed @pytest.mark.parametrize("absent", ("", None), ids=("empty", "null")) def test_empty_or_null_continuation_token_is_served_the_cached_answer_of_the_request_without_one( gateway: Gateway, absent: str | None ) -> None: with wire_server(_resumes_by_token) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, "gemini", wire) body: Final = request_body("chat", model, False, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") fresh: Final = answered("chat", False, post(gateway, "chat", body)) spend_rows(model) assert (fresh.answer, len(wire.drain())) == (PLAIN.text, 1), fresh twin: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": absent})) assert (twin.answer, wire.drain()) == (PLAIN.text, ()), twin rows: Final = spend_rows(model, expected=2) assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, True], rows def _refused_continuation_cells() -> tuple[ParameterSet, ...]: reached: Final = ( *itertools.product(PROVIDERS, ("chat",), (False, True)), *itertools.product(("gemini",), ("messages", "responses"), (False, True)), ) return tuple( pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") for provider, endpoint, stream in reached ) @pytest.mark.parametrize(("provider", "endpoint", "stream"), _refused_continuation_cells()) def test_answer_so_far_after_a_refused_continuation_is_cached_and_replayed_without_asking_google_again( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" so_far: Final = (FIRST_HALF.text, (cut_off(endpoint),), FIRST_HALF.usage) with wire_server(upstream(answers(first_id, FIRST_HALF), refuses(500))) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) body: Final = request_body(endpoint, model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") first: Final = answered(endpoint, stream, post(gateway, endpoint, body)) assert (first.answer, first.finishes, first.usage) == so_far, first assert first.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), first assert len(wire.drain()) == 2 assert_billed(model, endpoint, FIRST_HALF.usage, _identities(first, first_id)) replayed: Final = answered(endpoint, stream, post(gateway, endpoint, body)) assert (replayed.answer, replayed.finishes, replayed.usage) == so_far, replayed assert wire.drain() == () rows: Final = spend_rows(model, expected=2) assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, True], rows _OPENAI_ANSWER: Final = "An answer from another provider." def _openai_answering(request: Request) -> Reply: body: Final = JSON_OBJECT.validate_json(request.body) return chat_reply( f"chatcmpl-{uuid.uuid4().hex}", string_value(body["model"]), _OPENAI_ANSWER, stream=body.get("stream") is True ) def _register_openai(scenario: Scenario, wire: Wire) -> str: return scenario.model(model="openai/gpt-5.6", api_base=f"{wire.url}/v1") @pytest.mark.parametrize("stream", (False, True), ids=_mode) def test_continuation_token_on_another_providers_request_misses_the_cached_answer_of_the_request_without_it( gateway: Gateway, stream: bool ) -> None: with wire_server(answering_model_discovery(_openai_answering)) as wire, gateway.scenario() as scenario: model: Final = _register_openai(scenario, wire) body: Final = request_body("chat", model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") tokened: Final = {**body, "continuation_token": CUT_TOKEN} assert answered("chat", stream, post(gateway, "chat", body)).answer == _OPENAI_ANSWER spend_rows(model) assert answered("chat", stream, post(gateway, "chat", tokened)).answer == _OPENAI_ANSWER spend_rows(model, expected=2) assert answered("chat", stream, post(gateway, "chat", tokened)).answer == _OPENAI_ANSWER rows: Final = spend_rows(model, expected=3) assert posted_targets(wire) == ("/v1/chat/completions",) * 2 assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, False, True], rows def _openai_stream_reporting_cached_prompt_tokens(request: Request) -> Reply: chunk: Final[dict[str, JsonValue]] = { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion.chunk", "created": 1, "model": string_value(JSON_OBJECT.validate_json(request.body)["model"]), } usage: Final[dict[str, JsonValue]] = { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8, "prompt_tokens_details": {"cached_tokens": 2}, } frames: Final[tuple[dict[str, JsonValue], ...]] = ( {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": _OPENAI_ANSWER}}]}, {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}, ) return Reply( content_type="text/event-stream", chunks=(*(f"data: {json.dumps(item)}\n\n".encode() for item in frames), b"data: [DONE]\n\n"), ) def _spend_metadata(model: str) -> dict[str, JsonValue]: spend_rows(model) (row,) = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)) metadata: Final = row["metadata"] return JSON_OBJECT.validate_json(metadata) if isinstance(metadata, str) else JSON_OBJECT.validate_python(metadata) def test_streamed_answer_from_another_provider_logs_an_empty_initial_prompt_count_the_caller_never_sees( gateway: Gateway, ) -> None: with ( wire_server(answering_model_discovery(_openai_stream_reporting_cached_prompt_tokens)) as wire, gateway.scenario() as scenario, ): model: Final = _register_openai(scenario, wire) response: Final = post(gateway, "chat", request_body("chat", model, True)) assert answered("chat", True, response).answer == _OPENAI_ANSWER assert "initial_prompt_tokens" not in response.text, response.text logged: Final = at(_spend_metadata(model), "usage_object", "prompt_tokens_details") assert isinstance(logged, dict) and logged.get("initial_prompt_tokens", "absent") is None, logged def _native_path(model: str, stream: bool) -> str: return f"/v1beta/models/{model}:{'streamGenerateContent' if stream else 'generateContent'}" def _native(gateway: Gateway, model: str, stream: bool, body: Mapping[str, JsonValue]) -> httpx.Response: return gateway.request("POST", _native_path(model, stream), dict(body), params={"alt": "sse"} if stream else None) def _native_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") for provider, stream in itertools.product(PROVIDERS, (False, True)) ) @pytest.mark.parametrize(("provider", "stream"), _native_cells()) def test_native_route_hands_googles_continuation_fields_to_the_caller_from_one_request( gateway: Gateway, provider: Provider, stream: bool ) -> None: respond: Final = upstream( answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) ) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = _native(gateway, model, stream, _NATIVE_BODY) assert response.status_code == 200, response.text candidates: Final = tuple(at(item, "candidates", 0) for item in payloads(response.text, stream)) assert "".join(text_of(at(item, "content", "parts", 0, "text")) for item in candidates) == FIRST_HALF.text assert tuple(at(item, "continuationToken") for item in candidates) == ( (CHECKPOINT, HOP_TOKEN) if stream else (HOP_TOKEN,) ), response.text assert at(candidates[-1], "finishReason") == "CONTINUATION", response.text (only,) = received(wire, provider, stream, (False,)) assert "continuationToken" not in only, only (row,) = spend_rows(model) assert (row["status"], row["prompt_tokens"], row["completion_tokens"]) == ("success", *FIRST_HALF.usage), row assert number(row["spend"]) == pytest.approx(cost_of(FIRST_HALF.usage)), row @pytest.mark.parametrize(("provider", "stream"), _native_cells()) def test_native_route_passes_the_callers_continuation_token_to_google( gateway: Gateway, provider: Provider, stream: bool ) -> None: respond: Final = upstream( answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(f"resumed-{uuid.uuid4().hex}", RESUMED) ) with wire_server(respond) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) response: Final = _native(gateway, model, stream, {**_NATIVE_BODY, "continuationToken": CUT_TOKEN}) assert response.status_code == 200, response.text candidates: Final = tuple(at(item, "candidates", 0) for item in payloads(response.text, stream)) assert "".join(text_of(at(item, "content", "parts", 0, "text")) for item in candidates) == RESUMED.text (sent,) = wire.drain() assert sent.target in (target(provider, stream), target(provider, stream, continued=True)), sent.target assert JSON_OBJECT.validate_json(sent.body)["continuationToken"] == CUT_TOKEN, sent.body Searched = Literal["first", "each"] _SEARCHES: Final = ("scripted search one", "scripted search two") def _grounded(reply: dict[str, JsonValue], searches: Sequence[str]) -> dict[str, JsonValue]: candidate: Final = at(reply, "candidates", 0) assert isinstance(candidate, dict), reply grounding: Final[dict[str, JsonValue]] = { "webSearchQueries": [*searches], "groundingChunks": [{"web": {"uri": "https://example.com/source", "title": "example.com"}}], } return {**reply, "candidates": [{**candidate, "groundingMetadata": grounding}]} def _searching(response_id: str, hop: Hop, searches: Sequence[str]) -> Answer: def reply(stream: bool) -> Reply: if not searches: return hop_reply(response_id, hop, stream) if stream: opening, closing = stream_frames(response_id, hop) return sse((opening, _grounded(closing, searches))) whole: Final = frame(response_id, hop.text, hop.usage, hop.finish, hop.token) return Reply(body=json.dumps(_grounded(whole, searches)).encode()) return reply def _searches_counted(response: httpx.Response, stream: bool) -> JsonValue: metered: Final = tuple(item for item in payloads(response.text, stream) if isinstance(item.get("usage"), dict)) assert metered, response.text return at(metered[-1], "usage", "prompt_tokens_details", "web_search_requests") def _search_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, stream, searched, id=f"{provider}-{_mode(stream)}-{searched}") for provider, stream, searched in itertools.product(PROVIDERS, (False, True), ("first", "each")) ) @pytest.mark.parametrize(("provider", "stream", "searched"), _search_cells()) def test_google_searches_across_a_continued_answer_are_all_counted_and_billed_as_one_request_would_be( gateway: Gateway, provider: Provider, stream: bool, searched: Searched ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" opening_searches: Final = _SEARCHES if searched == "first" else _SEARCHES[:1] chained: Final = upstream( _searching(first_id, FIRST_HALF, opening_searches), _searching(f"second-{uuid.uuid4().hex}", SECOND_HALF, _SEARCHES[len(opening_searches) :]), ) in_one_request: Final = upstream(_searching(f"plain-{uuid.uuid4().hex}", PLAIN, _SEARCHES)) with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) twin: Final = register(gateway, scenario, provider, twin_wire) searching: Final[Mapping[str, JsonValue]] = {**NO_CACHE, "tools": [{"googleSearch": {}}]} response: Final = post(gateway, "chat", request_body("chat", model, stream, searching)) observed: Final = answered("chat", stream, response) assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed assert _searches_counted(response, stream) == len(_SEARCHES), response.text opening, continuing = received(wire, provider, stream, (False, True)) assert "tools" in opening and continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) alone: Final = post(gateway, "chat", request_body("chat", twin, stream, searching)) assert answered("chat", stream, alone).answer == PLAIN.text, alone.text assert _searches_counted(alone, stream) == len(_SEARCHES), alone.text (chain_row,) = spend_rows(model) (twin_row,) = spend_rows(twin) search_fee: Final = number(twin_row["spend"]) - cost_of(PLAIN.usage) assert search_fee > 0, twin_row assert number(chain_row["spend"]) == pytest.approx(cost_of(_CHAIN_USAGE) + search_fee), (chain_row, twin_row) Spelling = Literal["citations", "citationSources"] Earlier = Literal["multibyte", "thought"] _SPELLINGS: Final[tuple[Spelling, ...]] = ("citations", "citationSources") _EARLIER: Final[tuple[Earlier, ...]] = ("multibyte", "thought") _CITED_OPENING: Final = Hop(("Grüße aus 東京 🌸, ", "wo der Fuji steht, "), "CONTINUATION", HOP_TOKEN, 20, 7) _CITED_CLOSING: Final = Hop(("dann der ", "zweite Teil über Zürich."), "STOP", None, 27, 5) _CITED_WHOLE: Final = Hop((_CITED_OPENING.text, _CITED_CLOSING.text), "STOP", None, 20, 12) _THOUGHT: Final = "Ich plane die Antwort über 東京 zuerst." @dataclass(frozen=True, slots=True) class _Source: phrase: str uri: str @dataclass(frozen=True, slots=True) class _Carried: answer: str annotations: tuple[str, ...] citations: tuple[str, ...] retrieved: tuple[str, ...] _OPENING_SOURCE: Final = _Source("der Fuji", "https://example.com/fuji") _CLOSING_SOURCE: Final = _Source("über Zürich", "https://example.com/zurich") def _byte_span(text: str, phrase: str) -> dict[str, JsonValue]: # Google measures citation and grounding segment offsets in UTF-8 bytes (ai.google.dev/api/generate-content, # Segment and CitationSource, read 2026-10-14) start: Final = text.encode().index(phrase.encode()) return {"startIndex": start, "endIndex": start + len(phrase.encode())} def _cited_metadata(text: str, sources: Sequence[_Source], spelling: Spelling) -> dict[str, JsonValue]: return { "citationMetadata": {spelling: [{**_byte_span(text, source.phrase), "uri": source.uri} for source in sources]}, "urlContextMetadata": { "urlMetadata": [ {"retrievedUrl": source.uri, "urlRetrievalStatus": "URL_RETRIEVAL_STATUS_SUCCESS"} for source in sources ] }, "groundingMetadata": { "groundingChunks": [{"web": {"uri": source.uri, "title": source.phrase}} for source in sources], "groundingSupports": [ { "segment": {**_byte_span(text, source.phrase), "text": source.phrase}, "groundingChunkIndices": [index], } for index, source in enumerate(sources) ], }, } def _opening_metadata(spelling: Spelling) -> dict[str, JsonValue]: return _cited_metadata(_CITED_OPENING.text, (_OPENING_SOURCE,), spelling) def _decorated( reply: dict[str, JsonValue], thought: str | None, metadata: Mapping[str, JsonValue] ) -> dict[str, JsonValue]: candidate: Final = at(reply, "candidates", 0) assert isinstance(candidate, dict), reply thoughts: Final[tuple[JsonValue, ...]] = ({"text": thought, "thought": True},) if thought else () content: Final[dict[str, JsonValue]] = { "role": "model", "parts": [*thoughts, *items_of(at(candidate, "content", "parts"))], } return {**reply, "candidates": [{**candidate, "content": content, **metadata}]} def _citing(response_id: str, hop: Hop, metadata: Mapping[str, JsonValue], thought: str | None) -> Answer: def reply(stream: bool) -> Reply: if stream: opening, closing = stream_frames(response_id, hop) return sse((_decorated(opening, thought, NOTHING_EXTRA), _decorated(closing, None, metadata))) whole: Final = frame(response_id, hop.text, hop.usage, hop.finish, hop.token) return Reply(body=json.dumps(_decorated(whole, thought, metadata)).encode()) return reply def _canonical(values: Iterable[JsonValue]) -> tuple[str, ...]: return tuple(sorted(json.dumps(value, sort_keys=True) for value in values)) def _entries(listed: JsonValue, *keys: str) -> Iterator[JsonValue]: return itertools.chain.from_iterable( items_of(at(entry, key)) for entry, key in itertools.product(items_of(listed), keys) ) def _chat_carried(payloads_seen: Sequence[dict[str, JsonValue]], message: str) -> _Carried: return _Carried( answer="".join(text_of(at(payload, "choices", 0, message, "content")) for payload in payloads_seen), annotations=_canonical( itertools.chain.from_iterable( items_of(at(payload, "choices", 0, message, "annotations")) for payload in payloads_seen ) ), citations=_canonical( itertools.chain.from_iterable( _entries(payload.get("vertex_ai_citation_metadata"), "citations", "citationSources") for payload in payloads_seen ) ), retrieved=_canonical( itertools.chain.from_iterable( _entries(payload.get("vertex_ai_url_context_metadata"), "urlMetadata") for payload in payloads_seen ) ), ) def _output_parts(response: JsonValue) -> Iterator[JsonValue]: messages: Final = (item for item in items_of(at(response, "output")) if at(item, "type") == "message") return itertools.chain.from_iterable(items_of(at(item, "content")) for item in messages) def _responses_carried(response: JsonValue) -> _Carried: return _Carried( answer=output_text(response), annotations=_canonical( itertools.chain.from_iterable(items_of(at(part, "annotations")) for part in _output_parts(response)) ), citations=(), retrieved=(), ) def _completed(events: Sequence[dict[str, JsonValue]]) -> JsonValue: (completed,) = (event.get("response") for event in events if event.get("type") == "response.completed") return completed def _carried(endpoint: Endpoint, stream: bool, response: httpx.Response) -> _Carried: assert response.status_code == 200, response.text seen_payloads: Final = payloads(response.text, stream) match (endpoint, stream): case ("chat", False): return _chat_carried(seen_payloads, "message") case ("chat", True): return _chat_carried(seen_payloads, "delta") case ("responses", False): return _responses_carried(seen_payloads[0]) case ("responses", True): return _responses_carried(_completed(seen_payloads)) case _: raise AssertionError(f"No citations are read from {endpoint}") def _citation_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param( provider, endpoint, stream, spelling, earlier, id=f"{provider}-{endpoint}-{_mode(stream)}-{spelling}-{earlier}", ) for provider, endpoint, stream, spelling, earlier in itertools.product( PROVIDERS, ("chat", "responses"), (False, True), _SPELLINGS, _EARLIER ) ) @pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling", "earlier"), _citation_cells()) def test_citations_and_url_context_of_a_continued_answer_match_what_one_request_for_the_whole_answer_returns( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling, earlier: Earlier ) -> None: thought: Final = _THOUGHT if earlier == "thought" else None chained: Final = upstream( _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, _opening_metadata(spelling), thought), _citing( f"second-{uuid.uuid4().hex}", _CITED_CLOSING, _cited_metadata(_CITED_CLOSING.text, (_CLOSING_SOURCE,), spelling), None, ), ) in_one_request: Final = upstream( _citing( f"plain-{uuid.uuid4().hex}", _CITED_WHOLE, _cited_metadata(_CITED_WHOLE.text, (_OPENING_SOURCE, _CLOSING_SOURCE), spelling), thought, ) ) with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) twin: Final = register(gateway, scenario, provider, twin_wire) continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) assert alone.answer == _CITED_WHOLE.text and len(alone.annotations) == 2, alone assert endpoint != "chat" or (len(alone.citations), len(alone.retrieved)) == (2, 2), alone assert continued == alone, (continued, alone) assert len(wire.drain()) == 2 def _unplaced_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param(provider, endpoint, stream, spelling, id=f"{provider}-{endpoint}-{_mode(stream)}-{spelling}") for provider, endpoint, stream, spelling in itertools.product( PROVIDERS, ("chat", "responses"), (False, True), _SPELLINGS ) ) def _cited_uris(carried: _Carried) -> frozenset[JsonValue]: return frozenset(JSON_OBJECT.validate_json(cited).get("uri") for cited in carried.citations) @pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling"), _unplaced_cells()) def test_later_request_citing_a_source_with_no_offsets_returns_what_one_request_for_the_whole_answer_returns( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling ) -> None: # Google marks every CitationSource field optional, endIndex included (ai.google.dev/api/generate-content, # CitationSource, read 2026-10-14) unplaced: Final[JsonValue] = {"uri": _CLOSING_SOURCE.uri} placed: Final[JsonValue] = {**_byte_span(_CITED_OPENING.text, _OPENING_SOURCE.phrase), "uri": _OPENING_SOURCE.uri} chained: Final = upstream( _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, {"citationMetadata": {spelling: [placed]}}, None), _citing(f"second-{uuid.uuid4().hex}", _CITED_CLOSING, {"citationMetadata": {spelling: [unplaced]}}, None), ) in_one_request: Final = upstream( _citing(f"plain-{uuid.uuid4().hex}", _CITED_WHOLE, {"citationMetadata": {spelling: [placed, unplaced]}}, None) ) with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) twin: Final = register(gateway, scenario, provider, twin_wire) continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) assert alone.answer == continued.answer == _CITED_WHOLE.text, (continued, alone) assert stream or endpoint != "chat" or _OPENING_SOURCE.uri in _cited_uris(alone), alone assert _cited_uris(continued) == _cited_uris(alone), (continued, alone) assert len(wire.drain()) == 2 _FIRST_BYTE_SOURCE: Final = _Source("dann der", "https://example.com/dann") def _cited_from_first_byte(text: str, source: _Source, spelling: Spelling) -> dict[str, JsonValue]: assert text.startswith(source.phrase), text end: Final = len(source.phrase.encode()) return { "citationMetadata": {spelling: [{"endIndex": end, "uri": source.uri}]}, "urlContextMetadata": { "urlMetadata": [{"retrievedUrl": source.uri, "urlRetrievalStatus": "URL_RETRIEVAL_STATUS_SUCCESS"}] }, "groundingMetadata": { "groundingChunks": [{"web": {"uri": source.uri, "title": source.phrase}}], "groundingSupports": [{"segment": {"endIndex": end, "text": source.phrase}, "groundingChunkIndices": [0]}], }, } @pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling"), _unplaced_cells()) def test_later_request_span_from_its_first_byte_with_no_start_index_points_where_one_request_for_the_whole_answer_does( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling ) -> None: chained: Final = upstream( _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, _opening_metadata(spelling), None), _citing( f"second-{uuid.uuid4().hex}", _CITED_CLOSING, _cited_from_first_byte(_CITED_CLOSING.text, _FIRST_BYTE_SOURCE, spelling), None, ), ) in_one_request: Final = upstream( _citing( f"plain-{uuid.uuid4().hex}", _CITED_WHOLE, _cited_metadata(_CITED_WHOLE.text, (_OPENING_SOURCE, _FIRST_BYTE_SOURCE), spelling), None, ) ) with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire) twin: Final = register(gateway, scenario, provider, twin_wire) continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) assert alone.answer == _CITED_WHOLE.text and len(alone.annotations) == 2, alone assert endpoint != "chat" or (len(alone.citations), len(alone.retrieved)) == (2, 2), alone assert continued == alone, (continued, alone) assert len(wire.drain()) == 2 OddMetadata = Literal["citations_null", "citations_empty", "url_context_null", "url_context_empty", "both_missing"] _ODD_METADATA: Final[Mapping[OddMetadata, Mapping[str, JsonValue]]] = MappingProxyType( { "citations_null": {"citationMetadata": None}, "citations_empty": {"citationMetadata": {}}, "url_context_null": {"urlContextMetadata": None}, "url_context_empty": {"urlContextMetadata": {}}, "both_missing": NOTHING_EXTRA, } ) _OLDER_BACKEND: Final = "gemini-3.8-flash" _ODD_ANSWERS: Final[tuple[tuple[str, tuple[Hop, ...]], ...]] = ( (BACKEND, (PLAIN,)), (BACKEND, (FIRST_HALF, SECOND_HALF)), (_OLDER_BACKEND, (PLAIN,)), (_OLDER_BACKEND, (FIRST_HALF, SECOND_HALF)), ) def _odd(response_id: str, hop: Hop, backend: str, metadata: Mapping[str, JsonValue]) -> Answer: def reply(stream: bool) -> Reply: frames: Final = ( stream_frames(response_id, hop) if stream else (frame(response_id, hop.text, hop.usage, hop.finish, hop.token),) ) odd: Final = tuple({**_decorated(item, None, metadata), "modelVersion": backend} for item in frames) return sse(odd) if stream else Reply(body=json.dumps(odd[0]).encode()) return reply def _odd_cells() -> tuple[ParameterSet, ...]: return tuple( pytest.param( provider, endpoint, stream, backend, hops, odd, id=f"{provider}-{endpoint}-{_mode(stream)}-{backend}-{len(hops)}-requests-{odd}", ) for provider, endpoint, stream, (backend, hops), odd in itertools.product( PROVIDERS, ENDPOINTS, (False, True), _ODD_ANSWERS, _ODD_METADATA ) ) @pytest.mark.parametrize(("provider", "endpoint", "stream", "backend", "hops", "odd"), _odd_cells()) def test_null_empty_or_missing_citation_and_url_context_metadata_still_returns_the_whole_answer( gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, backend: str, hops: tuple[Hop, ...], odd: OddMetadata, ) -> None: replies: Final = tuple(_odd(f"odd-{uuid.uuid4().hex}", hop, backend, _ODD_METADATA[odd]) for hop in hops) with wire_server(upstream(*replies)) as wire, gateway.scenario() as scenario: model: Final = register(gateway, scenario, provider, wire, backend=backend) response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) observed: Final = answered(endpoint, stream, response) assert not stream or ends_cleanly(endpoint, response.text), response.text assert observed.answer == "".join(hop.text for hop in hops) and observed.errors == (), observed assert observed.finishes == (complete(endpoint),), observed assert len(wire.drain()) == len(hops) _OTHER_BODY: Final[Mapping[str, JsonValue]] = { **_NATIVE_BODY, "contents": [{"role": "user", "parts": [{"text": "Write the short answer."}]}], } def _pass_through_config(directory: Path, wire_url: str) -> Path: base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) path: Final = directory / "gemini-pass-through.yaml" path.write_text( yaml.safe_dump({**base, "environment_variables": {"GEMINI_API_BASE": wire_url, "GEMINI_API_KEY": API_KEY}}) ) return path def _scripted(script: Sequence[tuple[Mapping[str, JsonValue], Reply]]) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: sent: Final = JSON_OBJECT.validate_json(request.body) matching: Final = tuple(reply for body, reply in script if body == sent) assert len(matching) == 1, request return matching[0] return respond def _sent_by_google(reply: Reply, stream: bool) -> tuple[dict[str, JsonValue], ...]: return payloads(b"".join(reply.chunks or (reply.body,)).decode(), stream) def _pass_through(gateway: Gateway, stream: bool, body: Mapping[str, JsonValue]) -> httpx.Response: return gateway.request( "POST", f"/gemini{_native_path(BACKEND, stream)}", dict(body), params={"alt": "sse"} if stream else None, headers={"x-goog-api-key": gateway.key}, ) def _pass_through_row(response: httpx.Response) -> dict[str, JsonValue]: (row,) = eventually( lambda: read_rows(f"{SPEND_COLUMNS} WHERE request_id=%s", (response.headers["x-litellm-call-id"],)), lambda found: len(found) >= 1, seconds=70, ) return row @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @pytest.mark.parametrize("stream", (False, True), ids=_mode) def test_gemini_pass_through_route_relays_googles_continuation_fields_as_sent_and_bills_each_request_on_its_own( gateway: Gateway, tmp_path: Path, stream: bool ) -> None: script: Final = ( (_NATIVE_BODY, hop_reply("pass-through-opening", FIRST_HALF, stream)), ({**_NATIVE_BODY, "continuationToken": HOP_TOKEN}, hop_reply("pass-through-rest", SECOND_HALF, stream)), (_OTHER_BODY, hop_reply("pass-through-other", PLAIN, stream)), ) with wire_server(_scripted(script)) as wire: with owned_proxy_process(gateway, tmp_path, {}, config=_pass_through_config(tmp_path, wire.url)) as owned: responses: Final = tuple(_pass_through(owned.gateway, stream, body) for body, _ in script) assert [response.status_code for response in responses] == [200] * len(script), [ response.text for response in responses ] assert [payloads(response.text, stream) for response in responses] == [ _sent_by_google(reply, stream) for _, reply in script ] requests: Final = wire.drain() assert [(request.method, request.target.split("?", 1)[0]) for request in requests] == [ ("POST", _native_path(BACKEND, stream)) ] * len(script), requests assert [JSON_OBJECT.validate_json(request.body) for request in requests] == [body for body, _ in script] rows: Final = tuple(_pass_through_row(response) for response in responses) assert [ (row["call_type"], row["status"], row["prompt_tokens"], row["completion_tokens"]) for row in rows ] == [("pass_through_endpoint", "success", *hop.usage) for hop in (FIRST_HALF, SECOND_HALF, PLAIN)], rows opening_spend, rest_spend, other_spend = (number(row["spend"]) for row in rows) assert opening_spend == pytest.approx(other_spend), rows assert min(opening_spend, rest_spend) > 0, rows _LOGGED_CANDIDATES: Final = TypeAdapter(list[JsonValue]) @dataclass(frozen=True, slots=True) class _Traced: gateway: Gateway sink: str def _traced_config(directory: Path, sink: str) -> Path: base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) settings: Final = JSON_OBJECT.validate_python(base.get("litellm_settings", {})) path: Final = directory / "gemini-long-decode-otel.yaml" path.write_text( yaml.safe_dump( { **base, "litellm_settings": {**settings, "callbacks": ["otel"]}, "callback_settings": { "otel": {"exporter": "http/json", "endpoint": sink, "use_simple_processor": True} }, } ) ) return path @pytest.fixture(scope="module") def traced(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Traced]: directory: Final = tmp_path_factory.mktemp("long-decode-otel") with owned_sinks(directory) as sinks: overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} config: Final = _traced_config(directory, sinks.operator) with owned_proxy_process(session_gateway, directory, overrides, config=config) as owned: yield _Traced(owned.gateway, sinks.operator) def _raw_request_spans(sink: str, provider: Provider, marker: str) -> tuple[Span, ...]: return tuple( span for span in recorded_spans(sink)[1] if span["name"] == "raw_gen_ai_request" and marker in str(span["attributes"].get(f"llm.{provider}.contents")) ) def _logged(span: Span, provider: Provider, field: str) -> str: return string_value(span["attributes"][f"llm.{provider}.{field}"]) @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @pytest.mark.parametrize("last", (SECOND_HALF, CUT), ids=("finished", "cut")) @pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) def test_continued_answers_logged_raw_response_is_one_google_response_holding_the_whole_answer( traced: _Traced, provider: Provider, endpoint: Endpoint, last: Hop ) -> None: first_id: Final = f"first-{uuid.uuid4().hex}" marker: Final = f"{PROMPT} {uuid.uuid4().hex}" respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"last-{uuid.uuid4().hex}", last)) with wire_server(respond) as wire, traced.gateway.scenario() as scenario: model: Final = register(traced.gateway, scenario, provider, wire) response: Final = post(traced.gateway, endpoint, request_body(endpoint, model, False, prompt=marker)) assert response.status_code == 200, response.text (span,) = eventually( lambda: _raw_request_spans(traced.sink, provider, marker), lambda found: len(found) >= 1, seconds=60, return_last_on_timeout=True, ) assert span["attributes"].get(f"llm.{provider}.responseId") == first_id, span (candidate,) = _LOGGED_CANDIDATES.validate_python(ast.literal_eval(_logged(span, provider, "candidates"))) parts: Final = items_of(at(candidate, "content", "parts")) assert "".join(text_of(at(part, "text")) for part in parts) == FIRST_HALF.text + last.text, candidate ending: Final = (at(candidate, "finishReason"), at(candidate, "continuationToken")) assert ending == (last.finish, last.token), candidate usage: Final = JSON_OBJECT.validate_python(ast.literal_eval(_logged(span, provider, "usageMetadata"))) assert (at(usage, "promptTokenCount"), at(usage, "candidatesTokenCount")) == summed(FIRST_HALF, last), usage