scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-18 04:44:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123
litellm/tests/integration/providers/test_gemini_long_decode_wire.py
Mateo Wang 7090588bdf
feat(gemini)!: auto-continue gemini-4-argon long decodes with continuationToken (#46240) 
* feat(gemini): continue MAX_TOKENS decodes with continuationToken across chat, streaming, and interactions

* fix(gemini): continue long decode on CONTINUATION and cancel interactions at /cancel

* feat(gemini): add public gemini-4-argon cost-map rows for AI Studio and Vertex AI

Add gemini/gemini-4-argon and vertex_ai/gemini-4-argon alongside the existing
early-access row, carrying the same limits and capabilities. The AI Studio row
mirrors the early-access entry; the Vertex row differs only in provider and
pricing source. Prices follow once Google publishes them.

Extend the mapped metadata test to cover all three ids, that each routes to its
provider, that its backup copy matches, and that the public rows carry the same
limits as the early-access row.

* feat(gemini): raise gemini-4-argon to the 1M long decode cap and report initial_prompt_tokens across hops

* chore(prices): build on the merged gemini-4-argon rows and only raise the long decode cap

* fix(gemini): follow continuation hops in a loop instead of recursion

The unary hop collectors and the stream iterator called themselves once per continuation request, which the recursion check in CI rejects and which would hit the interpreter's call depth on a very long output. They now loop, and a test follows more hops than the interpreter allows nested calls

* fix(gemini): bill Google Search grounding from every continuation request

Google sends search grounding once, on the request that ran the search, and does not repeat it on a continuation request. The merged response kept only the last request's grounding, so usage.web_search_requests and the per-search fee dropped every earlier search. Unary merges and the stream iterator now carry the grounding of every request in the chain

* fix(gemini): continue long gemini-4-argon answers on streams and Vertex AI

A continued stream stopped after the first chunk of its second request. The stream wrapper asks the iterator for a fresh start before every chunk, and that reset it to the first request's finished stream. The iterator now reads from the newest request's stream

On Vertex AI the continuation request went to the v1 API version, which rejects continuationToken with a 400. A request that carries the token now goes to v1beta1, which accepts it, including a token that v1 issued

On a stream, initial_prompt_tokens came from a usage chunk in the middle of the stream, which reports a smaller prompt count than the last chunk. It is now the count the first request's stream ends with, the same number a unary call reports

continuation_token is no longer forwarded on the Interactions API, which answers 400 for the field. It is dropped there as it was before this branch

* chore(gemini): name the type errors the dual-mode stream iterator trips

The stream iterator holds a sync or an async line iterator in one untyped attribute. Pointing it at a continuation request's stream gave the type checker a concrete union for it, so six existing lines that only ever run in one of the two modes started to report errors

This adds a rule-named suppression with a reason on each of those lines and changes no code. Splitting the class into a sync and an async iterator would remove them and is left for a follow-up

* fix(gemini): continue only when Google ends a request with CONTINUATION, and stop streaming mid-answer tokens

A long decode now sends the next request only when Google's finish reason is CONTINUATION, so an answer cut at the caller's own max_tokens returns in one request with its continuation_token instead of drawing a rejected extra request

A request for more than one choice is never continued, since Google rejects a continuation with more than one candidate

Chat-format streams no longer carry the tokens Google sends mid-answer. Only a stream that ends cut off carries a continuation_token, on its last content chunk

The response cache key now includes continuation_token, so a resend with a token is not served the earlier segment

* fix(gemini): keep one response id across a continued answer

A continued chat stream changed its chunk id each time the next request to Google started, and a unary answer came back under the last request's id. Every chunk and the unary answer now carry the first request's id, including the chunk that ends a cut-off stream

A request whose n or max_tokens arrives as a string is sent once and never continued, where it used to fail with a 500 on the comparison. An empty continuation_token no longer changes the cache key

* test(gemini): add integration cells for long decode continuation

Covers continued answers on chat, /v1/messages and /v1/responses, streamed and not, on the Gemini API and Vertex AI against a scripted Google, plus concurrency, outage and worker restart runs

Also moves the Google stand-in cancel route to /cancel, the URL this PR now calls for Interactions cancel, and checks the cancel test posts there exactly once

* fix(vertex_ai): fetch the access token again before each continuation request

A long Gemini 4 answer on Vertex AI can run for over half an hour of follow-up requests, and every one of them reused the bearer token fetched before the first request. Once that token expired Google answered 401 and the caller got a cut-off answer with finish_reason length

Each Google request in a chain now asks for the token again, which is a cache hit until the token nears expiry, on the streamed and unary paths, sync and async. A refresh that fails counts as a failed follow-up, so the caller still gets the answer so far with its continuation token. A caller's own Authorization in extra_headers still wins on every request, and the Gemini API key path is unchanged

* fix(vertex_ai): start each continuation stream with an empty parser buffer

When a Google stream ended on an unfinished frame, its leftover text stayed in the stream parser and every frame of the next continuation stream was appended to it. None of them could be decoded, so the whole continuation was billed and dropped without an error

Each continuation stream now starts with a fresh buffer, and a test covers a stream that ends on an unfinished frame, sync and async

* fix(vertex_ai): treat a continuation stream with no readable frame as a failed follow-up

A continuation stream that answered 200 and then sent Google's error frame first, sent a body that is not JSON, or dropped before its first frame ended the caller's stream with `stop` or an error and no token, so a cut answer looked complete or was lost. It now ends the way a refused follow-up does, with the answer so far, `finish_reason` length and that request's token

Every continuation stream request now also gets the caller's per-request timeout, which #46349 added to the first one, so a follow-up that never answers ends at that timeout the same way

* fix(vertex_ai): keep every follow-up's citations and URL context in a continued answer

A non-streamed answer that took several Google requests kept citationMetadata and urlContextMetadata from the last request only, and left every later request's offsets counting from that request's own text. So a URL read on the first request was missing from the answer, citations from earlier requests were dropped, and a later request's citations and search supports pointed at the wrong part of the joined answer

The joined answer now keeps every request's citations (Vertex AI's `citations` and AI Studio's `citationSources`) and URL context, and shifts each later request's citation and search support offsets by the UTF-8 bytes of the answer text sent before it, thought parts excluded, since Google counts these offsets in bytes

* fix(vertex_ai): answer when a candidate sends null citationMetadata or urlContextMetadata

A Gemini candidate whose citationMetadata or urlContextMetadata was JSON null made the call fail on every Gemini model, non-streamed with a 422 or 500 and streamed with an error event, because the per-request metadata helper handed the null back as is and the caller tried to extend a list with it. A continued answer also failed when a follow-up's citation or search support came with a null startIndex

A null field is now kept as one entry on a single request, as it was before this branch, a follow-up's null URL context adds nothing to a continued answer, and a null startIndex counts from the start of that request's text. Tests cover missing, null and empty fields on one request and on a continued answer, streamed and not

* fix(vertex_ai): pass through spans with no end index and point a continued stream's citations into the whole answer

A later request's citation or search support with no endIndex made the whole non-streamed answer fail with a 500, so a span Google sends without an end index now comes back as Google sent it, the way one request for the whole answer returns it. A span with an end index but no start index still counts from 0

A continued stream handed callers each later request's citation and search support offsets counted from that request's own text. The stream now adds the answer bytes of every earlier request, so the offsets point into the whole answer like the non-streamed path and one request for the whole answer

Adds the audit's wire tests for both, which fail 24 cells before this change

* test(vertex_ai): cover a later request whose cited span starts at its first byte with no start index

Google leaves out a zero startIndex, so the second request of a continued answer can cite its own opening with only an endIndex. These wire cells check that span lands where one request for the whole answer puts it, on Gemini and Vertex AI, chat and Responses, unary and streamed

* fix(vertex_ai): log a continued answer's raw response as one object

A non-streamed answer joined from several Google requests logged its raw response as a JSON list of every request's body. The OpenTelemetry logger reads that raw response as one Gemini body to set its `llm.gemini.*` attributes, so on a continued answer it hit a validation error and set none of them

It now logs the merged body, which carries only the last request's token, and a single request still logs Google's body unchanged. The new unit test fails with the change reverted

Also adds integration cells for a cut chat stream's token chunk, an answer joined from 41 requests, a caller that gave up on a non-streamed call, counts and caps that are not numbers or are sent twice, `n` sent as true, a partial answer cached after a refused continuation, the token in another provider's cache key, and the sync SDK path with its Vertex token refetch

* test(integration): check a continued answer logs one merged raw response

A non-streamed Gemini answer that took several Google requests now hands logging one merged response, so the OpenTelemetry span carries the llm.gemini attributes. These 12 cells fail on the merge base and on the commit before that fix, and pass on the branch

* test(vertex_ai): cover a follow-up tool call and its thought signature in long decode

* test(gemini): pin that a one-request answer logs Google's body unchanged
2026-10-15 22:54:23 +00:00

2106 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
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