scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-25 05:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 125
litellm/tests/integration/providers/test_gemini_long_decode_wire.py
yuneng-jiang fc27dfdc7d
test(integration): give each owned proxy a private database and Redis by default (#47213) 
* test(integration): give each owned proxy a private database and Redis by default

owned_proxy and owned_proxy_process now boot on a scratch database and a
per-proxy redis-server unless the caller opts into the worker's state with
state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips
DATABASE_URL. Under the replica profile the private database gets the
writer and reader roles so the reader path stays covered.

Gateway carries the database_url its proxy writes to, and the Scenario
cleanup checks and delete_key_if_present read from it, so they keep
checking something on a private database.

The 189 test files with owned proxy sites that depend on the worker's
database or Redis opt in with a shares_worker_state module marker, which a
module-scoped conftest fixture turns into the shared state. The 7 helper
modules with such sites opt in per call, since conftest fixtures and other
test files reach them.

* test(integration): reach session fixtures with the shared marker and keep the worker retry default

The module-scoped fixture that applied shares_worker_state ran after
session-scoped fixtures, so proxies those booted went private in marked
files. The marker is now applied around each marked item's setup, call
and teardown, which covers fixtures of every scope.

Owned proxies on the worker database inherited the worker gateway's
router_settings.num_retries of 0 from it. A private proxy now gets the same
default in its config copy unless the config sets num_retries itself.

Five providers files read the owned proxy's spend rows or response cache
through the worker's database or Redis and opt into sharing.
2026-10-18 03:37:41 +00:00

2108 lines
100 KiB
Python
Raw Blame History

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
pytestmark: Final = pytest.mark.shares_worker_state
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"<html>not json</html>", 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
Reference in a new issue View git blame Copy permalink