diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 163fbd64a9d..655131e3c69 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -42,6 +42,7 @@ from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" +_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) _DECODED_LITERAL: Final = TypeAdapter(object) @@ -399,8 +400,10 @@ class Cache: for param in kwargs: if param in scope_excluded_params: continue - if param in combined_kwargs: + if param in combined_kwargs or param in _ALWAYS_KEYED_PROVIDER_PARAMS: param_value: str | None = self._get_param_value(param, kwargs) + if param in _ALWAYS_KEYED_PROVIDER_PARAMS and not param_value: + continue if param_value is not None: cache_key += f"{param}: {param_value}" elif not is_litellm_owned_kwarg(param): diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ff7a232e83a..cc1f6fef019 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -244,6 +244,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { # Gemini / Vertex AI "STOP": "stop", "MAX_TOKENS": "length", + "CONTINUATION": "length", "SAFETY": "content_filter", "RECITATION": "content_filter", "FINISH_REASON_UNSPECIFIED": "stop", diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 6b1713482b0..232d9d8eb42 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -73,7 +73,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped to ``config``). """ - return ("safetySettings", "toolConfig", "cachedContent", "labels") + return ("safetySettings", "toolConfig", "cachedContent", "labels", "continuationToken") @abstractmethod def map_generate_content_optional_params( diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index f0e6bb12d39..a58037f691f 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -369,13 +369,13 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[str, dict]: - """POST /{api_version}/interactions/{interaction_id}:cancel (if supported)""" + """POST /{api_version}/interactions/{interaction_id}/cancel""" resolved_api_base: Final = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}:cancel", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}/cancel", {}, ) diff --git a/litellm/llms/vertex_ai/gemini/long_decode.py b/litellm/llms/vertex_ai/gemini/long_decode.py new file mode 100644 index 00000000000..81d6c04d0e8 --- /dev/null +++ b/litellm/llms/vertex_ai/gemini/long_decode.py @@ -0,0 +1,354 @@ +import json +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate, chain +from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys + +import httpx + +from litellm._logging import verbose_logger +from litellm.exceptions import Timeout +from litellm.types.llms.vertex_ai import ( + Candidates, + Citation, + CitationMetadata, + GenerateContentResponseBody, + GroundingMetadata, + GroundingSegment, + GroundingSupport, + HttpxPartType, + RequestBody, + UrlContextMetadata, + UsageMetadata, +) +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +from ..common_utils import VertexAIError + +_SUMMED_USAGE_COUNTS: Final = ( + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "cachedContentTokenCount", + "thoughtsTokenCount", + "toolUsePromptTokenCount", + "responseTokenCount", +) + +_CONCATENATED_USAGE_DETAILS: Final = ( + "promptTokensDetails", + "cacheTokensDetails", + "candidatesTokensDetails", + "responseTokensDetails", + "toolUsePromptTokensDetails", +) + +HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout) + +HOP_END_FINISH_REASON: Final = "CONTINUATION" +CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON}) + +_VERTEX_V1_PATH: Final = "/v1/projects/" +_VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/" + +LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...] +AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]] +SyncHopPost: TypeAlias = Callable[[RequestBody], httpx.Response] +_Metadata: Final = TypeVar("_Metadata", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class AsyncStreamResume: + request_body: RequestBody + open_hop: AsyncHopPost + + +@dataclass(frozen=True, slots=True) +class SyncStreamResume: + request_body: RequestBody + open_hop: SyncHopPost + + +StreamResume: TypeAlias = AsyncStreamResume | SyncStreamResume + + +def _int_or_none(value: object) -> int | None: + return value if isinstance(value, int) else None + + +def continuation_cap_of(request_body: RequestBody) -> int | None: + generation_config: Final = request_body.get("generationConfig") + if generation_config is None or "max_output_tokens" not in generation_config: + return None + if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1: + return None + return _int_or_none(generation_config["max_output_tokens"]) + + +def client_continuation_token(candidate: Candidates) -> str | None: + if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS: + return None + return candidate.get("continuationToken") + + +def _first_candidate(body: GenerateContentResponseBody) -> Candidates | None: + candidates: Final = body.get("candidates") + return candidates[0] if candidates else None + + +def output_tokens_of(usage: UsageMetadata) -> int: + return usage.get("candidatesTokenCount", 0) + usage.get("thoughtsTokenCount", 0) + + +def next_hop_token( + body: GenerateContentResponseBody, + hop_usage: UsageMetadata, + produced_before: int, + output_cap: int | None, +) -> str | None: + if output_cap is None: + return None + candidate: Final = _first_candidate(body) + if candidate is None or candidate.get("finishReason") != HOP_END_FINISH_REASON: + return None + token: Final = candidate.get("continuationToken") + hop_tokens: Final = output_tokens_of(hop_usage) + if not token or hop_tokens <= 0 or produced_before + hop_tokens >= output_cap: + return None + return token + + +def hop_request(request_body: RequestBody, token: str) -> RequestBody: + return {**request_body, "continuationToken": token} + + +def url_for(url: str, request_body: RequestBody) -> str: + if "continuationToken" not in request_body: + return url + return url.replace(_VERTEX_V1_PATH, _VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN, 1) + + +def sum_usage_metadata(first: UsageMetadata, *rest: UsageMetadata) -> UsageMetadata: + usages: Final = (first, *rest) + counts: Final = { + key: sum(usage.get(key, 0) for usage in usages) + for key in _SUMMED_USAGE_COUNTS + if any(key in usage for usage in usages) + } + details: Final = { + key: list(chain.from_iterable(usage.get(key, ()) for usage in usages)) + for key in _CONCATENATED_USAGE_DETAILS + if any(key in usage for usage in usages) + } + return cast(UsageMetadata, {**usages[-1], **counts, **details}) # cast-ok: a TypedDict cannot take computed keys + + +def _candidate_parts(candidate: Candidates) -> Sequence[HttpxPartType]: + content: Final = candidate.get("content") + return (content.get("parts") or ()) if content is not None else () + + +def per_hop(metadata: _Metadata | tuple[_Metadata, ...]) -> tuple[_Metadata, ...]: + return metadata if isinstance(metadata, tuple) else (metadata,) + + +def _candidate_grounding(candidate: Candidates) -> tuple[GroundingMetadata, ...]: + grounding: Final = candidate.get("groundingMetadata") + return () if grounding is None else per_hop(grounding) + + +def _answer_text(part: HttpxPartType) -> str: + text: Final = part.get("text") + return text if isinstance(text, str) and not part.get("thought") else "" + + +def answer_bytes(candidate: Candidates) -> int: + return sum(len(text.encode()) for text in map(_answer_text, _candidate_parts(candidate))) + + +def _shifted_support(support: GroundingSupport, offset: int) -> GroundingSupport: + segment: Final = support.get("segment") + end: Final = None if segment is None else _int_or_none(segment.get("endIndex")) + if segment is None or end is None: + return support + shifted: Final[GroundingSegment] = { + **segment, + "startIndex": (segment.get("startIndex") or 0) + offset, + "endIndex": end + offset, + } + return {**support, "segment": shifted} + + +def _shifted_grounding(grounding: GroundingMetadata, offset: int) -> GroundingMetadata: + if not offset or "groundingSupports" not in grounding: + return grounding + return { + **grounding, + "groundingSupports": tuple(_shifted_support(support, offset) for support in grounding["groundingSupports"]), + } + + +def _hop_grounding(candidate: Candidates, offset: int) -> tuple[GroundingMetadata, ...]: + return tuple(_shifted_grounding(grounding, offset) for grounding in _candidate_grounding(candidate)) + + +def _shifted_citation(citation: Citation, offset: int) -> Citation: + end: Final = _int_or_none(citation.get("endIndex")) + if end is None: + return citation + return {**citation, "startIndex": (citation.get("startIndex") or 0) + offset, "endIndex": end + offset} + + +def _shifted_citation_metadata(metadata: CitationMetadata, offset: int) -> CitationMetadata: + citations: Final = metadata.get("citations") + sources: Final = metadata.get("citationSources") + with_citations: Final[CitationMetadata] = ( + metadata + if citations is None + else {**metadata, "citations": tuple(_shifted_citation(c, offset) for c in citations)} + ) + return ( + with_citations + if sources is None + else {**with_citations, "citationSources": tuple(_shifted_citation(c, offset) for c in sources)} + ) + + +def _hop_citations(candidate: Candidates, offset: int) -> tuple[CitationMetadata, ...]: + metadata: Final = candidate.get("citationMetadata") + if metadata is None: + return () + return tuple(_shifted_citation_metadata(item, offset) for item in per_hop(metadata)) + + +def with_spans_shifted(candidate: Candidates, offset: int) -> Candidates: + if not offset: + return candidate + grounding: Final = _hop_grounding(candidate, offset) + citations: Final = _hop_citations(candidate, offset) + with_grounding: Final[Candidates] = {**candidate, "groundingMetadata": grounding} if grounding else candidate + return {**with_grounding, "citationMetadata": citations} if citations else with_grounding + + +def _hop_url_context(candidate: Candidates) -> tuple[UrlContextMetadata, ...]: + url_context: Final = candidate.get("urlContextMetadata") + return () if url_context is None else per_hop(url_context) + + +def _with_hop_metadata(candidate: Candidates, hop_candidates: Sequence[Candidates]) -> Candidates: + offsets: Final = tuple(accumulate(map(answer_bytes, hop_candidates[:-1]), initial=0)) + grounding: Final = tuple(chain.from_iterable(map(_hop_grounding, hop_candidates, offsets))) + with_grounding: Final[Candidates] = {**candidate, "groundingMetadata": grounding} if grounding else candidate + citations: Final = tuple(chain.from_iterable(map(_hop_citations, hop_candidates, offsets))) + with_citations: Final[Candidates] = ( + {**with_grounding, "citationMetadata": citations} if citations else with_grounding + ) + url_context: Final = tuple(chain.from_iterable(map(_hop_url_context, hop_candidates))) + return {**with_citations, "urlContextMetadata": url_context} if url_context else with_citations + + +def merge_hops(hops: LongDecodeHops) -> GenerateContentResponseBody: + last: Final = hops[-1] + if len(hops) == 1: + return last + hop_candidates: Final = tuple(c for c in map(_first_candidate, hops) if c is not None) + merged_first: Final[Candidates] = { + **(_first_candidate(last) or Candidates()), + "content": {"role": "model", "parts": list(chain.from_iterable(map(_candidate_parts, hop_candidates)))}, + } + merged: Final[GenerateContentResponseBody] = { + **last, + "candidates": [_with_hop_metadata(merged_first, hop_candidates), *(last.get("candidates") or [])[1:]], + "usageMetadata": sum_usage_metadata(*(hop.get("usageMetadata", UsageMetadata()) for hop in hops)), + } + first_response_id: Final = hops[0].get("responseId") + return {**merged, "responseId": first_response_id} if first_response_id else merged + + +def initial_prompt_tokens_of(hops: LongDecodeHops) -> int | None: + if len(hops) == 1: + return None + return hops[0].get("usageMetadata", UsageMetadata()).get("promptTokenCount") + + +def with_initial_prompt_tokens(usage: Usage, initial_prompt_tokens: int | None) -> Usage: + if initial_prompt_tokens is None: + return usage + details: Final = usage.prompt_tokens_details + marked: Final = ( + details.model_copy(update={"initial_prompt_tokens": initial_prompt_tokens}) + if isinstance(details, PromptTokensDetailsWrapper) + else PromptTokensDetailsWrapper(initial_prompt_tokens=initial_prompt_tokens) + ) + return usage.model_copy(update={"prompt_tokens_details": marked}) + + +def parse_body(response: httpx.Response) -> GenerateContentResponseBody: + try: + return GenerateContentResponseBody(**response.json()) + except Exception as e: + raise VertexAIError( + message=( + f"Error converting to valid response block={e}. " + "File an issue if litellm error - https://github.com/BerriAI/litellm/issues" + ), + status_code=422, + headers=response.headers, + ) + + +def hops_original_response( + hops: LongDecodeHops, merged: GenerateContentResponseBody, first_response: httpx.Response +) -> str: + return first_response.text if len(hops) == 1 else json.dumps(merged) + + +def log_failed_hop(error: Exception | str) -> None: + verbose_logger.warning( + "Gemini long decode: a continuation hop failed, returning the answer so far with its continuation_token: %s", + error, + ) + + +def _unary_next_hop_token(hops: LongDecodeHops, output_cap: int | None) -> str | None: + last: Final = hops[-1] + return next_hop_token( + last, + last.get("usageMetadata", UsageMetadata()), + sum(output_tokens_of(hop.get("usageMetadata", UsageMetadata())) for hop in hops[:-1]), + output_cap, + ) + + +async def collect_hops_async( + first: GenerateContentResponseBody, + request_body: RequestBody, + post: AsyncHopPost, +) -> LongDecodeHops: + output_cap: Final = continuation_cap_of(request_body) + hops: LongDecodeHops = (first,) # rebind-ok: grows by one hop per loop pass, CI bans the recursive form + while (token := _unary_next_hop_token(hops, output_cap)) is not None: + try: + body = parse_body(await post(hop_request(request_body, token))) + except HOP_FAILURES as e: + log_failed_hop(e) + return hops + hops = (*hops, body) + return hops + + +def collect_hops_sync( + first: GenerateContentResponseBody, + request_body: RequestBody, + post: SyncHopPost, +) -> LongDecodeHops: + output_cap: Final = continuation_cap_of(request_body) + hops: LongDecodeHops = (first,) # rebind-ok: grows by one hop per loop pass, CI bans the recursive form + while (token := _unary_next_hop_token(hops, output_cap)) is not None: + try: + body = parse_body(post(hop_request(request_body, token))) + except HOP_FAILURES as e: + log_failed_hop(e) + return hops + hops = (*hops, body) + return hops diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 909ca526233..57327731863 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1233,6 +1233,9 @@ def transform_request_body( else: data["serviceTier"] = service_tier + if continuation_token := optional_params.pop("continuation_token", None): + data["continuationToken"] = continuation_token + # Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty if labels and custom_llm_provider != LlmProviders.GEMINI: data["labels"] = labels diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index be808e05202..bed7bc1bd74 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -4,7 +4,7 @@ import json import re import time -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import partial from types import MappingProxyType @@ -65,6 +65,7 @@ from litellm.types.llms.vertex_ai import ( GenerateContentResponseBody, HttpxPartType, LogprobsResult, + RequestBody, ToolConfig, Tools, UsageMetadata, @@ -96,6 +97,31 @@ from ..common_utils import ( ) from ..vertex_llm_base import VertexBase from .grounding_requests import calculate_grounding_requests +from .long_decode import ( + HOP_FAILURES, + AsyncStreamResume, + LongDecodeHops, + StreamResume, + SyncStreamResume, + answer_bytes, + client_continuation_token, + collect_hops_async, + collect_hops_sync, + continuation_cap_of, + hop_request, + hops_original_response, + initial_prompt_tokens_of, + log_failed_hop, + merge_hops, + next_hop_token, + output_tokens_of, + parse_body, + per_hop, + sum_usage_metadata, + url_for, + with_initial_prompt_tokens, + with_spans_shifted, +) from .transformation import ( async_transform_request_body, gemini_convert_messages_with_history, @@ -2189,7 +2215,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): citation_metadata: Final[list] = [] if "groundingMetadata" in candidate: - if isinstance(candidate["groundingMetadata"], list): + if isinstance(candidate["groundingMetadata"], (list, tuple)): grounding_metadata.extend(candidate["groundingMetadata"]) else: grounding_metadata.append(candidate["groundingMetadata"]) @@ -2198,11 +2224,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings.append(candidate["safetyRatings"]) if "citationMetadata" in candidate: - citation_metadata.append(candidate["citationMetadata"]) + citation_metadata.extend(per_hop(candidate["citationMetadata"])) if "urlContextMetadata" in candidate: - # Add URL context metadata to grounding metadata - url_context_metadata.append(cast(dict, candidate["urlContextMetadata"])) + url_context_metadata.extend(cast(tuple[dict, ...], per_hop(candidate["urlContextMetadata"]))) return ( grounding_metadata, @@ -2475,17 +2500,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): reasoning_content = "\n".join(reasoning_content_parts) chat_completion_message["reasoning_content"] = reasoning_content - # Store thoughtSignatures in provider_specific_fields - if thought_signatures is not None: - thought_signature_fields = chat_completion_message.get("provider_specific_fields") or {} - thought_signature_fields["thought_signatures"] = thought_signatures - chat_completion_message["provider_specific_fields"] = thought_signature_fields - - # Store server-side tool invocations in provider_specific_fields - if server_side_tool_invocations is not None: - tool_invocation_fields = chat_completion_message.get("provider_specific_fields") or {} - tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations - chat_completion_message["provider_specific_fields"] = tool_invocation_fields + message_provider_fields = { + key: value + for key, value in ( + ("thought_signatures", thought_signatures), + ("server_side_tool_invocations", server_side_tool_invocations), + ("continuation_token", client_continuation_token(candidate)), + ) + if value is not None + } + if message_provider_fields: + chat_completion_message["provider_specific_fields"] = message_provider_fields if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( @@ -2562,18 +2587,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): additional_args={"complete_input_dict": request_data}, ) - ## RESPONSE OBJECT - try: - completion_response: Final = GenerateContentResponseBody(**raw_response.json()) - except Exception as e: - raise VertexAIError( - message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", - status_code=422, - headers=raw_response.headers, - ) - return self._transform_google_generate_content_to_openai_model_response( - completion_response=completion_response, + completion_response=parse_body(raw_response), model_response=model_response, model=model, logging_obj=logging_obj, @@ -2761,7 +2776,9 @@ async def make_call( data: str, model: str, messages: list, - logging_obj, + logging_obj: LoggingClass, + request_body: RequestBody, + hop_headers: Callable[[], Awaitable[Mapping[str, str]]], timeout: float | httpx.Timeout | None = None, ): if gemini_client is not None: @@ -2773,7 +2790,12 @@ async def make_call( try: response: Final = await client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout + url_for(api_base, request_body), + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, ) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -2790,12 +2812,23 @@ async def make_call( headers=response.headers, ) + async def open_hop(body: RequestBody) -> httpx.Response: + return await client.post( + url_for(api_base, body), + headers=dict(await hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + completion_stream: Final = ModelResponseIterator( streaming_response=response.aiter_lines(), sync_stream=False, logging_obj=logging_obj, response_headers=response.headers, response=response, + resume=AsyncStreamResume(request_body=request_body, open_hop=open_hop), ) # LOGGING logging_obj.post_call( @@ -2816,7 +2849,9 @@ def make_sync_call( data: str, model: str, messages: list, - logging_obj, + logging_obj: LoggingClass, + request_body: RequestBody, + hop_headers: Callable[[], Mapping[str, str]], timeout: float | httpx.Timeout | None = None, ): if gemini_client is not None: @@ -2825,7 +2860,12 @@ def make_sync_call( client = HTTPHandler() # Create a new client if none provided response: Final = client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout + url_for(api_base, request_body), + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, + timeout=timeout, ) if response.status_code != 200 and response.status_code != 201: @@ -2835,12 +2875,23 @@ def make_sync_call( headers=response.headers, ) + def open_hop(body: RequestBody) -> httpx.Response: + return client.post( + url_for(api_base, body), + headers=dict(hop_headers()), + data=json.dumps(body), + stream=True, + logging_obj=logging_obj, + timeout=timeout, + ) + completion_stream: Final = ModelResponseIterator( streaming_response=response.iter_lines(), sync_stream=True, logging_obj=logging_obj, response_headers=response.headers, response=response, + resume=SyncStreamResume(request_body=request_body, open_hop=open_hop), ) # LOGGING @@ -2858,6 +2909,42 @@ class VertexLLM(VertexBase): def __init__(self) -> None: super().__init__() + async def _fresh_headers_async( + self, + headers: Mapping[str, str], + extra_headers: Mapping[str, str] | None, + vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, + vertex_project: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> Mapping[str, str]: + if custom_llm_provider == "gemini": + return headers + try: + token, _ = await self._ensure_access_token_async( + credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider=custom_llm_provider + ) + except Exception as e: + raise VertexAIError(status_code=401, message=str(e), headers=None) from e + return {**headers, "Authorization": f"Bearer {token}", **(extra_headers or {})} + + def _fresh_headers( + self, + headers: Mapping[str, str], + extra_headers: Mapping[str, str] | None, + vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, + vertex_project: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> Mapping[str, str]: + if custom_llm_provider == "gemini": + return headers + try: + token, _ = self._ensure_access_token( + credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider=custom_llm_provider + ) + except Exception as e: + raise VertexAIError(status_code=401, message=str(e), headers=None) from e + return {**headers, "Authorization": f"Bearer {token}", **(extra_headers or {})} + async def async_streaming( self, model: str, @@ -2944,6 +3031,15 @@ class VertexLLM(VertexBase): model=model, messages=messages, logging_obj=logging_obj, + request_body=request_body, + hop_headers=partial( + self._fresh_headers_async, + headers, + extra_headers, + vertex_credentials, + vertex_project, + custom_llm_provider, + ), timeout=timeout, ), model=model, @@ -2964,7 +3060,7 @@ class VertexLLM(VertexBase): ], # if it's vertex_ai or gemini (google ai studio) timeout: float | httpx.Timeout | None, encoding: object, - logging_obj, + logging_obj: LoggingClass, stream, optional_params: dict, litellm_params: dict, @@ -3033,39 +3129,42 @@ class VertexLLM(VertexBase): }, ) - try: - response: Final = await client.post( - api_base, - headers=headers, - json=cast(dict, request_body), - logging_obj=logging_obj, - ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise VertexAIError( - status_code=error_code, - message=err.response.text, - headers=err.response.headers, - ) - except httpx.TimeoutException: - raise VertexAIError( - status_code=408, - message="Timeout error occurred.", - headers=None, - ) + fresh_headers: Final = partial( + self._fresh_headers_async, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider + ) - return VertexGeminiConfig().transform_response( + async def post(body: RequestBody) -> httpx.Response: + try: + response: Final = await client.post( + url_for(api_base, body), + headers=dict(await fresh_headers()), + json=dict(body), + logging_obj=logging_obj, + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + raise VertexAIError( + status_code=err.response.status_code, + message=err.response.text, + headers=err.response.headers, + ) + except httpx.TimeoutException: + raise VertexAIError( + status_code=408, + message="Timeout error occurred.", + headers=None, + ) + return response + + first_response: Final = await post(request_body) + return _transform_long_decode_hops( model=model, - raw_response=response, + hops=await collect_hops_async(parse_body(first_response), request_body, post), + first_response=first_response, model_response=model_response, logging_obj=logging_obj, - api_key="", - request_data=cast(dict, request_body), + request_body=request_body, messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, ) def completion( @@ -3078,7 +3177,7 @@ class VertexLLM(VertexBase): "vertex_ai", "vertex_ai_beta", "gemini" ], # if it's vertex_ai or gemini (google ai studio) encoding: object, - logging_obj, + logging_obj: LoggingClass, optional_params: dict, acompletion: bool, timeout: float | httpx.Timeout | None, @@ -3207,6 +3306,10 @@ class VertexLLM(VertexBase): }, ) + fresh_headers: Final = partial( + self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider + ) + ## SYNC STREAMING CALL ## if stream is True: request_data_str: Final = json.dumps(data) @@ -3221,6 +3324,8 @@ class VertexLLM(VertexBase): messages=messages, logging_obj=logging_obj, headers=headers, + request_body=data, + hop_headers=fresh_headers, timeout=timeout, ), model=model, @@ -3241,37 +3346,68 @@ class VertexLLM(VertexBase): else: client = client - try: - response: Final = client.post(url=url, headers=headers, json=data, logging_obj=logging_obj) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise VertexAIError( - status_code=error_code, - message=err.response.text, - headers=err.response.headers, - ) - except httpx.TimeoutException: - raise VertexAIError( - status_code=408, - message="Timeout error occurred.", - headers=None, - ) + def post(body: RequestBody) -> httpx.Response: + try: + response: Final = client.post( + url=url_for(url, body), headers=dict(fresh_headers()), json=dict(body), logging_obj=logging_obj + ) + response.raise_for_status() + except httpx.HTTPStatusError as err: + raise VertexAIError( + status_code=err.response.status_code, + message=err.response.text, + headers=err.response.headers, + ) + except httpx.TimeoutException: + raise VertexAIError( + status_code=408, + message="Timeout error occurred.", + headers=None, + ) + return response - return VertexGeminiConfig().transform_response( + first_response: Final = post(data) + return _transform_long_decode_hops( model=model, - raw_response=response, + hops=collect_hops_sync(parse_body(first_response), data, post), + first_response=first_response, model_response=model_response, logging_obj=logging_obj, - optional_params=optional_params, - litellm_params=litellm_params, - api_key="", - request_data=data, + request_body=data, messages=messages, - encoding=encoding, ) +def _transform_long_decode_hops( + model: str, + hops: LongDecodeHops, + first_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LoggingClass, + request_body: RequestBody, + messages: Sequence[AllMessageValues], +) -> ModelResponse: + body: Final = merge_hops(hops) + logging_obj.post_call( + input=messages, + api_key="", + original_response=hops_original_response(hops, body, first_response), + additional_args={"complete_input_dict": request_body}, + ) + merged: Final = VertexGeminiConfig().transform_google_generate_content_to_openai_model_response( + completion_response=body, + model_response=model_response, + model=model, + logging_obj=logging_obj, + raw_response=first_response, + ) + merged_usage: Final = getattr(merged, "usage", None) + if not isinstance(merged_usage, Usage): + return merged + marked: Final = with_initial_prompt_tokens(merged_usage, initial_prompt_tokens_of(hops)) + return merged if marked is merged_usage else merged.model_copy(update={"usage": marked}) + + class ModelResponseIterator: def __init__( self, @@ -3280,6 +3416,7 @@ class ModelResponseIterator: logging_obj: LoggingClass, response_headers: dict[str, str] | None = None, response: httpx.Response | None = None, + resume: StreamResume | None = None, ): from litellm.litellm_core_utils.prompt_templates.common_utils import ( check_is_function_call, @@ -3295,6 +3432,18 @@ class ModelResponseIterator: self.is_function_call = check_is_function_call(logging_obj) self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False + self._resume = resume + self._output_cap = continuation_cap_of(resume.request_body) if resume is not None else None + self._pending_hop_token: str | None = None + self._unanswered_hop_token: str | None = None + self._first_response_id: str | None = None + self._usage_offset: UsageMetadata = UsageMetadata() + self._hop_usage: UsageMetadata = UsageMetadata() + self._initial_prompt_tokens: int | None = None + self._grounding_offset: tuple[Mapping[str, object], ...] = () + self._hop_grounding: tuple[Mapping[str, object], ...] = () + self._answer_offset: int = 0 + self._hop_answer_bytes: int = 0 @property def accumulated_json(self) -> str: @@ -3406,6 +3555,20 @@ class ModelResponseIterator: citation_metadata, ) + def _marked_stream_usage(self, usage: Usage | None) -> Usage | None: + if usage is None or not self._usage_offset: + return usage + return with_initial_prompt_tokens(usage, self._initial_prompt_tokens) + + def _grounding_through_this_hop( + self, grounding_metadata: Sequence[Mapping[str, object]] + ) -> Sequence[Mapping[str, object]]: + if grounding_metadata: + self._hop_grounding = tuple(grounding_metadata) + if not self._grounding_offset: + return grounding_metadata + return (*self._grounding_offset, *self._hop_grounding) + def _apply_stream_usage_metadata( self, processed_chunk: GenerateContentResponseBody, @@ -3420,7 +3583,7 @@ class ModelResponseIterator: cache_creation_tokens=_cache_creation_tokens(self.logging_obj), ) - VertexGeminiConfig.set_grounding_usage_counters(usage, grounding_metadata) + VertexGeminiConfig.set_grounding_usage_counters(usage, self._grounding_through_this_hop(grounding_metadata)) traffic_type: Final = processed_chunk.get("usageMetadata", {}).get("trafficType") if traffic_type: @@ -3442,11 +3605,14 @@ class ModelResponseIterator: # Detect mid-stream error chunks (e.g. 429 RESOURCE_EXHAUSTED). # Vertex AI can return errors as HTTP 200 but with an "error" field in the SSE body. self._check_streaming_error(chunk) + self._unanswered_hop_token = None from litellm.types.utils import ModelResponseStream - processed_chunk: Final = GenerateContentResponseBody(**chunk) - response_id: Final = processed_chunk.get("responseId") + processed_chunk: Final = self._with_answer_offset( + self._with_usage_offset(GenerateContentResponseBody(**chunk)) + ) + response_id: Final = self._first_response_id or processed_chunk.get("responseId") served: Final = _served_model_name(processed_chunk.get("modelVersion")) model_response = ModelResponseStream( choices=[], @@ -3477,9 +3643,12 @@ class ModelResponseIterator: citation_metadata, ) = self._apply_stream_candidates(_candidates, model_response) - usage: Final = self._apply_stream_usage_metadata(processed_chunk, model_response, grounding_metadata) + usage: Final = self._marked_stream_usage( + self._apply_stream_usage_metadata(processed_chunk, model_response, grounding_metadata) + ) setattr(model_response, "usage", usage) + self._plan_resume(processed_chunk, model_response) model_response.hidden_params["is_finished"] = False return model_response @@ -3487,6 +3656,108 @@ class ModelResponseIterator: except json.JSONDecodeError: raise ValueError(f"Failed to decode JSON from chunk: {chunk}") + def _with_usage_offset(self, chunk: GenerateContentResponseBody) -> GenerateContentResponseBody: + if "usageMetadata" not in chunk: + return chunk + self._hop_usage = chunk["usageMetadata"] + if not self._usage_offset: + return chunk + return {**chunk, "usageMetadata": sum_usage_metadata(self._usage_offset, self._hop_usage)} + + def _with_answer_offset(self, chunk: GenerateContentResponseBody) -> GenerateContentResponseBody: + candidates: Final = chunk.get("candidates") + if self._output_cap is None or not candidates: + return chunk + self._hop_answer_bytes += answer_bytes(candidates[0]) + if not self._answer_offset: + return chunk + return {**chunk, "candidates": [with_spans_shifted(candidates[0], self._answer_offset), *candidates[1:]]} + + def _plan_resume(self, processed_chunk: GenerateContentResponseBody, model_response: "ModelResponseStream") -> None: + if self._output_cap is None: + return + token: Final = next_hop_token( + processed_chunk, + hop_usage=self._hop_usage, + produced_before=output_tokens_of(self._usage_offset), + output_cap=self._output_cap, + ) + if token is None: + return + self._pending_hop_token = token + self._first_response_id = self._first_response_id or processed_chunk.get("responseId") + for choice in model_response.choices: + delta_fields: Mapping[str, object] = choice.delta.provider_specific_fields or {} + choice.finish_reason = None + choice.delta.provider_specific_fields = { + key: value for key, value in delta_fields.items() if key != "continuation_token" + } or None + + def _start_hop(self, response: httpx.Response, token: str) -> None: + if self._initial_prompt_tokens is None: + self._initial_prompt_tokens = self._hop_usage.get("promptTokenCount") + self._usage_offset = sum_usage_metadata(self._usage_offset, self._hop_usage) + self._hop_usage = UsageMetadata() + self._grounding_offset = (*self._grounding_offset, *self._hop_grounding) + self._hop_grounding = () + self._answer_offset += self._hop_answer_bytes + self._hop_answer_bytes = 0 + self._json_buffer = JSONFragmentAccumulator() + self._unanswered_hop_token = token + self.response = response + + def _truncated_stream_end(self, token: str) -> "ModelResponseStream": + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + id=self._first_response_id, + choices=[ + StreamingChoices( + index=0, + finish_reason="length", + delta=Delta(provider_specific_fields={"continuation_token": token}), + ) + ], + ) + + def _end_of_unanswered_hop( + self, reason: Exception | str = "the continuation stream ended without a readable frame" + ) -> Optional["ModelResponseStream"]: + token: Final = self._unanswered_hop_token + if token is None: + return None + self._unanswered_hop_token = None + log_failed_hop(reason) + return self._truncated_stream_end(token) + + async def _open_next_hop_async(self, token: str) -> Optional["ModelResponseStream"]: + if not isinstance(self._resume, AsyncStreamResume): + return None + await self.aclose() + try: + hop_response: Final = await self._resume.open_hop(hop_request(self._resume.request_body, token)) + except HOP_FAILURES as e: + log_failed_hop(e) + return self._truncated_stream_end(token) + self._start_hop(hop_response, token) + self.streaming_response = hop_response.aiter_lines() + self.async_response_iterator = self.streaming_response + return None + + def _open_next_hop_sync(self, token: str) -> Optional["ModelResponseStream"]: + if not isinstance(self._resume, SyncStreamResume): + return None + self.close() + try: + hop_response: Final = self._resume.open_hop(hop_request(self._resume.request_body, token)) + except HOP_FAILURES as e: + log_failed_hop(e) + return self._truncated_stream_end(token) + self._start_hop(hop_response, token) + self.streaming_response = hop_response.iter_lines() + self.response_iterator = self.streaming_response + return None + # Sync iterator def __iter__(self): self.response_iterator = self.streaming_response @@ -3555,58 +3826,92 @@ class ModelResponseIterator: except Exception: raise - def __next__(self): + def __next__(self) -> Optional["ModelResponseStream"]: try: - chunk: Final = self.response_iterator.__next__() - except StopIteration: - if self.chunk_type == "accumulated_json" and self._json_buffer: - result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True) - if result is not None: - return result - raise StopIteration - except ValueError as e: - raise RuntimeError(f"Error receiving chunk from stream: {e}") + return self._next_chunk() + except HOP_FAILURES as e: + if (unanswered_end := self._end_of_unanswered_hop(e)) is None: + raise + self.close() + return unanswered_end - try: - return self._common_chunk_parsing_logic(chunk=chunk) - except StopIteration: - raise StopIteration - except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + def _next_chunk(self) -> Optional["ModelResponseStream"]: + while True: + try: + chunk = self.response_iterator.__next__() # pyright: ignore[reportAttributeAccessIssue] # sync stream + except StopIteration: + if self.chunk_type == "accumulated_json" and self._json_buffer: + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result + if (unanswered_end := self._end_of_unanswered_hop()) is not None: + return unanswered_end + if (token := self._pending_hop_token) is None: + raise StopIteration + self._pending_hop_token = None + if (truncated_end := self._open_next_hop_sync(token)) is not None: + return truncated_end + continue + except ValueError as e: + raise RuntimeError(f"Error receiving chunk from stream: {e}") + + try: + return self._common_chunk_parsing_logic(chunk=chunk) + except StopIteration: + raise StopIteration + except ValueError as e: + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") # Async iterator def __aiter__(self): - self.async_response_iterator = self.streaming_response.__aiter__() + self.async_response_iterator = self.streaming_response.__aiter__() # pyright: ignore[reportAttributeAccessIssue] # async stream return self - async def __anext__(self): + async def __anext__(self) -> Optional["ModelResponseStream"]: try: - chunk: Final = await self.async_response_iterator.__anext__() - except StopAsyncIteration: - if self.chunk_type == "accumulated_json" and self._json_buffer: - result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True) - if result is not None: - return result - raise StopAsyncIteration - except ValueError as e: - raise RuntimeError(f"Error receiving chunk from stream: {e}") + return await self._anext_chunk() + except HOP_FAILURES as e: + if (unanswered_end := self._end_of_unanswered_hop(e)) is None: + raise + await self.aclose() + return unanswered_end - try: - return self._common_chunk_parsing_logic(chunk=chunk) - except StopAsyncIteration: - raise StopAsyncIteration - except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + async def _anext_chunk(self) -> Optional["ModelResponseStream"]: + while True: + try: + chunk = await self.async_response_iterator.__anext__() + except StopAsyncIteration: + if self.chunk_type == "accumulated_json" and self._json_buffer: + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result + if (unanswered_end := self._end_of_unanswered_hop()) is not None: + return unanswered_end + if (token := self._pending_hop_token) is None: + raise StopAsyncIteration + self._pending_hop_token = None + if (truncated_end := await self._open_next_hop_async(token)) is not None: + return truncated_end + continue + except ValueError as e: + raise RuntimeError(f"Error receiving chunk from stream: {e}") + + try: + return self._common_chunk_parsing_logic(chunk=chunk) + except StopAsyncIteration: + raise StopAsyncIteration + except ValueError as e: + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") async def aclose(self) -> None: - iterator: Final = getattr( + iterator: Final = getattr( # pyright: ignore[reportUnknownVariableType] # sync or async line iterator self, "async_response_iterator", self.streaming_response, ) if iterator is not None and hasattr(iterator, "aclose"): try: - await iterator.aclose() + await iterator.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # hasattr guard except Exception as e: # noqa: BLE001 verbose_logger.debug("ModelResponseIterator.aclose: error closing iterator: %s", e) if self.response is not None: @@ -3616,10 +3921,10 @@ class ModelResponseIterator: verbose_logger.debug("ModelResponseIterator.aclose: error closing response: %s", e) def close(self) -> None: - iterator: Final = getattr(self, "response_iterator", self.streaming_response) + iterator: Final = getattr(self, "response_iterator", self.streaming_response) # pyright: ignore[reportUnknownVariableType] # sync or async line iterator if iterator is not None and hasattr(iterator, "close"): try: - iterator.close() + iterator.close() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # hasattr guard except Exception as e: # noqa: BLE001 verbose_logger.debug("ModelResponseIterator.close: error closing iterator: %s", e) if self.response is not None: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4c6ba4a6a83..308dfdeb62a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -74790,8 +74790,8 @@ "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, + "max_output_tokens": 1048576, + "max_tokens": 1048576, "mode": "chat", "output_cost_per_reasoning_token": 1e-05, "output_cost_per_token": 1e-05, diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index eabc360fd6d..784e3d4719f 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -517,7 +517,7 @@ async def cancel_interaction( """ Cancel an interaction by ID. - Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel + Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}/cancel """ from litellm.proxy.proxy_server import ( general_settings, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 116fdda0b5d..3dc7a56239d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from enum import Enum from typing import Any, Final, Literal, Protocol @@ -357,6 +358,7 @@ class RequestBody(TypedDict, total=False): cachedContent: str labels: dict[str, str] serviceTier: str + continuationToken: str # writable-ok: transform_request_body fills the body field by field class CachedContentRequestBody(TypedDict, total=False): @@ -397,8 +399,9 @@ class Citation(TypedDict): publicationDate: Date -class CitationMetadata(TypedDict): - citations: list[Citation] +class CitationMetadata(TypedDict, total=False): + citations: ReadOnly[Sequence[Citation]] + citationSources: ReadOnly[Sequence[Citation]] class SearchEntryPoint(TypedDict, total=False): @@ -406,10 +409,24 @@ class SearchEntryPoint(TypedDict, total=False): sdkBlob: str +class GroundingSegment(TypedDict, total=False): + partIndex: ReadOnly[int] + startIndex: ReadOnly[int] + endIndex: ReadOnly[int] + text: ReadOnly[str] + + +class GroundingSupport(TypedDict, total=False): + segment: ReadOnly[GroundingSegment] + groundingChunkIndices: ReadOnly[Sequence[int]] + confidenceScores: ReadOnly[Sequence[float]] + + class GroundingMetadata(TypedDict, total=False): webSearchQueries: list[str] searchEntryPoint: SearchEntryPoint groundingAttributions: list[dict] + groundingSupports: ReadOnly[Sequence[GroundingSupport]] class LogprobsCandidate(TypedDict): @@ -440,6 +457,7 @@ GeminiFinishReason = Literal[ "FINISH_REASON_UNSPECIFIED", "STOP", "MAX_TOKENS", + "CONTINUATION", "SAFETY", "RECITATION", "LANGUAGE", @@ -466,11 +484,12 @@ class Candidates(TypedDict, total=False): content: HttpxContentType finishReason: GeminiFinishReason safetyRatings: list[SafetyRatings] - citationMetadata: CitationMetadata - groundingMetadata: GroundingMetadata + citationMetadata: CitationMetadata | tuple[CitationMetadata, ...] + groundingMetadata: GroundingMetadata | tuple[GroundingMetadata, ...] finishMessage: str logprobsResult: LogprobsResult - urlContextMetadata: UrlContextMetadata + urlContextMetadata: UrlContextMetadata | tuple[UrlContextMetadata, ...] + continuationToken: ReadOnly[str] class PromptFeedback(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5f97c25d3d2..d3c884b4a1f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1801,6 +1801,9 @@ class PromptTokensDetailsWrapper( tool_use_tokens: int | None = None """Prompt tokens consumed by server-side tool use (e.g. Gemini grounding via googleSearch).""" + initial_prompt_tokens: int | None = None + """The first hop's prompt tokens on a Gemini long decode, the real context size when prompt_tokens sums every hop.""" + character_count: int | None = None """Character count sent to the model. Used for Vertex AI multimodal embeddings.""" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c6ba4a6a83..308dfdeb62a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -74790,8 +74790,8 @@ "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, + "max_output_tokens": 1048576, + "max_tokens": 1048576, "mode": "chat", "output_cost_per_reasoning_token": 1e-05, "output_cost_per_token": 1e-05, diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index f7af39bcdfb..d2b67c5438e 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -567,10 +567,10 @@ class Provider: Route("/triton/embeddings", triton_embeddings, methods=["POST"]), Route("/vector_stores/{vector_store_id}/search", self.vector_store_search, methods=["POST"]), Route("/__interactions/{interaction_id}", self.interaction_state, methods=["PUT", "DELETE"]), - Route("/v1beta/interactions/{interaction_id}:cancel", self.cancel_interaction, methods=["POST"]), + Route("/v1beta/interactions/{interaction_id}/cancel", self.cancel_interaction, methods=["POST"]), Route("/v1beta/interactions/{interaction_id}", self.interaction, methods=["GET", "DELETE"]), Route( - "/{prefix:path}/v1beta/interactions/{interaction_id}:cancel", + "/{prefix:path}/v1beta/interactions/{interaction_id}/cancel", self.cancel_interaction, methods=["POST"], ), diff --git a/tests/integration/providers/_gemini_long_decode_support.py b/tests/integration/providers/_gemini_long_decode_support.py new file mode 100644 index 00000000000..d7bbb502eff --- /dev/null +++ b/tests/integration/providers/_gemini_long_decode_support.py @@ -0,0 +1,594 @@ +from __future__ import annotations + +import base64 +import json +import os +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from functools import reduce +from types import MappingProxyType +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, string_value +from integration._support.database import read_rows +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + +Provider = Literal["gemini", "vertex_ai"] +Endpoint = Literal["chat", "messages", "responses"] +Answer = Callable[[bool], Reply] + +BACKEND: Final = "gemini-4-argon" +API_KEY: Final = "synthetic-gemini-key" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-central1" +MODEL_PATH: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/google/models/{BACKEND}" +CONTINUED_MODEL_PATH: Final = MODEL_PATH.replace("/v1/", "/v1beta1/", 1) +SALT: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") +CAP: Final = 1000 +INPUT_COST: Final = 0.001 +OUTPUT_COST: Final = 0.002 +NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +NOTHING_EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +PROMPT: Final = "Write the long answer." +CHECKPOINT: Final = "checkpoint-token-never-for-clients" +HOP_TOKEN: Final = "hop-token-litellm-sends-back" +CUT_TOKEN: Final = "cut-token-for-the-caller" +REFUSAL: Final = "The scripted upstream refused this continuation request" +PROVIDERS: Final[tuple[Provider, ...]] = ("gemini", "vertex_ai") +ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +SPEND_COLUMNS: Final = ( + 'SELECT request_id, call_type, status, cache_hit, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs"' +) + + +@dataclass(frozen=True, slots=True) +class Hop: + pieces: tuple[str, str] + finish: str + token: str | None + prompt_tokens: int + output_tokens: int + + @property + def text(self) -> str: + return "".join(self.pieces) + + @property + def usage(self) -> tuple[int, int]: + return self.prompt_tokens, self.output_tokens + + +FIRST_HALF: Final = Hop(("The first ", "half, "), "CONTINUATION", HOP_TOKEN, 20, 7) +SECOND_HALF: Final = Hop(("then the ", "second half."), "STOP", None, 27, 5) +CUT: Final = Hop(("Cut at the ", "caller's cap"), "MAX_TOKENS", CUT_TOKEN, 20, 7) +PLAIN: Final = Hop(("A short ", "answer."), "STOP", None, 20, 7) +RESUMED: Final = Hop(("and the ", "rest of it."), "STOP", None, 27, 5) + + +def summed(*hops: Hop) -> tuple[int, int]: + return sum(hop.prompt_tokens for hop in hops), sum(hop.output_tokens for hop in hops) + + +def frame( + response_id: str, text: str, usage: tuple[int, int], finish: str | None = None, token: str | None = None +) -> dict[str, JsonValue]: + prompt_tokens, output_tokens = usage + return { + "responseId": response_id, + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": text}]}, + "index": 0, + **({"finishReason": finish} if finish else {}), + **({"continuationToken": token} if token else {}), + } + ], + "usageMetadata": { + "promptTokenCount": prompt_tokens, + "candidatesTokenCount": output_tokens, + "totalTokenCount": prompt_tokens + output_tokens, + }, + "modelVersion": BACKEND, + } + + +def stream_frames(response_id: str, hop: Hop) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + opening, closing = hop.pieces + return ( + frame(response_id, opening, (hop.prompt_tokens, 1), token=CHECKPOINT), + frame(response_id, closing, hop.usage, hop.finish, hop.token), + ) + + +def sse(frames: Sequence[Mapping[str, JsonValue]], abort_after: int | None = None) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(item)}\n\n".encode() for item in frames), + abort_after=abort_after, + ) + + +def hop_reply(response_id: str, hop: Hop, stream: bool, abort_after: int | None = None) -> Reply: + if stream: + return sse(stream_frames(response_id, hop), abort_after) + return Reply(body=json.dumps(frame(response_id, hop.text, hop.usage, hop.finish, hop.token)).encode()) + + +def answers(response_id: str, hop: Hop) -> Answer: + return lambda stream: hop_reply(response_id, hop, stream) + + +def refuses(status: int) -> Answer: + error: Final = {"error": {"code": status, "message": REFUSAL, "status": "INVALID_ARGUMENT"}} + return lambda _stream: Reply(status=status, body=json.dumps(error).encode()) + + +def drops_before_answering(_stream: bool) -> Reply: + return Reply(drop_connection=True) + + +def drops_mid_body(response_id: str, hop: Hop) -> Answer: + return lambda stream: replace(hop_reply(response_id, hop, stream, abort_after=1), pause_between_chunks=0.5) + + +def streams(request: Request) -> bool: + return ":streamGenerateContent" in request.target + + +def upstream(first: Answer, continued: Answer | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if "continuationToken" not in JSON_OBJECT.validate_json(request.body): + return first(streams(request)) + assert continued is not None, f"Unscripted continuation request: {request.target}" + return continued(streams(request)) + + return respond + + +def register_vertex( + scenario: Scenario, + wire: Wire, + token_url: str, + *, + extra_headers: JsonValue = None, + backend: str = BACKEND, + settings: Mapping[str, JsonValue] = NOTHING_EXTRA, +) -> str: + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=f"{wire.url}{MODEL_PATH.removesuffix(BACKEND)}{backend}", + api_key=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=service_account_json(PROJECT, token_url), + input_cost_per_token=INPUT_COST, + output_cost_per_token=OUTPUT_COST, + extra_headers=extra_headers, + model_info=None, + **settings, + ) + + +def register( + gateway: Gateway, + scenario: Scenario, + provider: Provider, + wire: Wire, + *, + backend: str = BACKEND, + settings: Mapping[str, JsonValue] = NOTHING_EXTRA, +) -> str: + if provider == "gemini": + return scenario.model( + model=f"gemini/{backend}", + api_base=wire.url, + api_key=API_KEY, + input_cost_per_token=INPUT_COST, + output_cost_per_token=OUTPUT_COST, + model_info=None, + **settings, + ) + return register_vertex(scenario, wire, gateway.upstream_url.rstrip("/"), backend=backend, settings=settings) + + +def target(provider: Provider, stream: bool, continued: bool = False) -> str: + vertex_path: Final = CONTINUED_MODEL_PATH if continued else MODEL_PATH + prefix: Final = f"/models/{BACKEND}" if provider == "gemini" else vertex_path + return f"{prefix}:streamGenerateContent?alt=sse" if stream else f"{prefix}:generateContent" + + +def expected_auth(provider: Provider) -> tuple[str, str]: + return ("x-goog-api-key", API_KEY) if provider == "gemini" else ("authorization", "Bearer scripted-token") + + +def received( + wire: Wire, provider: Provider, stream: bool, continued: tuple[bool, ...] +) -> tuple[dict[str, JsonValue], ...]: + requests: Final = wire.drain() + assert [(request.method, request.target) for request in requests] == [ + ("POST", target(provider, stream, resumed)) for resumed in continued + ], requests + header, value = expected_auth(provider) + assert [request.headers[header] for request in requests] == [value] * len(requests), requests + return tuple(JSON_OBJECT.validate_json(request.body) for request in requests) + + +def path_of(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def request_body( + endpoint: Endpoint, + model: str, + stream: bool, + extra: Mapping[str, JsonValue] = NO_CACHE, + *, + prompt: str = PROMPT, + cap: int | None = CAP, +) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + **({"max_tokens": cap} if cap is not None else {}), + **({"stream_options": {"include_usage": True}} if stream else {}), + **extra, + } + case "messages": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + **({"max_tokens": cap} if cap is not None else {}), + **extra, + } + case "responses": + return { + "model": model, + "input": prompt, + "stream": stream, + **({"max_output_tokens": cap} if cap is not None else {}), + **extra, + } + + +def post( + gateway: Gateway, endpoint: Endpoint, body: Mapping[str, JsonValue], *, key: str | None = None +) -> httpx.Response: + return gateway.request("POST", path_of(endpoint), body, key=key, headers={"anthropic-version": "2023-06-01"}) + + +def _step(value: JsonValue, key: str | int) -> JsonValue: + if isinstance(key, str): + return value.get(key) if isinstance(value, dict) else None + return value[key] if isinstance(value, list) and len(value) > key else None + + +def at(value: JsonValue, *keys: str | int) -> JsonValue: + return reduce(_step, keys, value) + + +def text_of(value: JsonValue) -> str: + return value if isinstance(value, str) else "" + + +def items_of(value: JsonValue) -> tuple[JsonValue, ...]: + return tuple(value) if isinstance(value, list) else () + + +def count_of(value: JsonValue) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def sse_payloads(text: str) -> tuple[dict[str, JsonValue], ...]: + lines: Final = tuple(line for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]") + return tuple(JSON_OBJECT.validate_json(line.removeprefix("data: ").encode()) for line in lines) + + +def payloads(text: str, stream: bool) -> tuple[dict[str, JsonValue], ...]: + return sse_payloads(text) if stream else (JSON_OBJECT.validate_json(text.encode()),) + + +@dataclass(frozen=True, slots=True) +class Seen: + ids: frozenset[str] + answer: str + streamed: str + finishes: tuple[str, ...] + usage: tuple[int | None, int | None] | None + initial_prompt_tokens: JsonValue + tokens: tuple[str, ...] + token_carriers: tuple[int, ...] + content_carriers: tuple[int, ...] + incomplete_reason: JsonValue + errors: tuple[dict[str, JsonValue], ...] + raw: str + + +def _named(values: Sequence[JsonValue]) -> tuple[str, ...]: + return tuple(value for value in values if isinstance(value, str)) + + +def _errors(frames: Sequence[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + item + for item in frames + if item.get("error") is not None or item.get("type") in ("error", "response.failed", "response.error") + ) + + +def _chat_unary(payload: dict[str, JsonValue], raw: str) -> Seen: + choices: Final = items_of(payload.get("choices")) + answer: Final = text_of(at(payload, "choices", 0, "message", "content")) + return Seen( + ids=frozenset(_named((payload.get("id"),))), + answer=answer, + streamed=answer, + finishes=_named(tuple(at(choice, "finish_reason") for choice in choices)), + usage=(count_of(at(payload, "usage", "prompt_tokens")), count_of(at(payload, "usage", "completion_tokens"))), + initial_prompt_tokens=at(payload, "usage", "prompt_tokens_details", "initial_prompt_tokens"), + tokens=_named( + tuple(at(choice, "message", "provider_specific_fields", "continuation_token") for choice in choices) + ), + token_carriers=(), + content_carriers=(), + incomplete_reason=None, + errors=_errors((payload,)), + raw=raw, + ) + + +def _chunk_token(chunk: dict[str, JsonValue]) -> JsonValue: + return at(chunk, "choices", 0, "delta", "provider_specific_fields", "continuation_token") + + +def _chat_stream(chunks: tuple[dict[str, JsonValue], ...], raw: str) -> Seen: + contents: Final = tuple(text_of(at(chunk, "choices", 0, "delta", "content")) for chunk in chunks) + metered: Final = tuple(chunk for chunk in chunks if isinstance(chunk.get("usage"), dict)) + answer: Final = "".join(contents) + return Seen( + ids=frozenset(_named(tuple(chunk.get("id") for chunk in chunks))), + answer=answer, + streamed=answer, + finishes=_named(tuple(at(chunk, "choices", 0, "finish_reason") for chunk in chunks)), + usage=( + ( + count_of(at(metered[-1], "usage", "prompt_tokens")), + count_of(at(metered[-1], "usage", "completion_tokens")), + ) + if metered + else None + ), + initial_prompt_tokens=( + at(metered[-1], "usage", "prompt_tokens_details", "initial_prompt_tokens") if metered else None + ), + tokens=_named(tuple(_chunk_token(chunk) for chunk in chunks)), + token_carriers=tuple(index for index, chunk in enumerate(chunks) if _chunk_token(chunk) is not None), + content_carriers=tuple(index for index, content in enumerate(contents) if content), + incomplete_reason=None, + errors=_errors(chunks), + raw=raw, + ) + + +def _messages_unary(payload: dict[str, JsonValue], raw: str) -> Seen: + blocks: Final = items_of(payload.get("content")) + answer: Final = "".join(text_of(at(block, "text")) for block in blocks if at(block, "type") == "text") + return Seen( + ids=frozenset(_named((payload.get("id"),))), + answer=answer, + streamed=answer, + finishes=_named((payload.get("stop_reason"),)), + usage=(count_of(at(payload, "usage", "input_tokens")), count_of(at(payload, "usage", "output_tokens"))), + initial_prompt_tokens=None, + tokens=(), + token_carriers=(), + content_carriers=(), + incomplete_reason=None, + errors=_errors((payload,)), + raw=raw, + ) + + +def _messages_stream(events: tuple[dict[str, JsonValue], ...], raw: str) -> Seen: + closings: Final = tuple(event for event in events if event.get("type") == "message_delta") + answer: Final = "".join( + text_of(at(event, "delta", "text")) for event in events if at(event, "delta", "type") == "text_delta" + ) + return Seen( + ids=frozenset( + _named(tuple(at(event, "message", "id") for event in events if event.get("type") == "message_start")) + ), + answer=answer, + streamed=answer, + finishes=_named(tuple(at(event, "delta", "stop_reason") for event in closings)), + usage=( + ( + count_of(at(closings[-1], "usage", "input_tokens")), + count_of(at(closings[-1], "usage", "output_tokens")), + ) + if closings + else None + ), + initial_prompt_tokens=None, + tokens=(), + token_carriers=(), + content_carriers=(), + incomplete_reason=None, + errors=_errors(events), + raw=raw, + ) + + +def _part_texts(item: JsonValue) -> tuple[str, ...]: + return tuple( + text_of(at(part, "text")) for part in items_of(at(item, "content")) if at(part, "type") == "output_text" + ) + + +def output_text(response: JsonValue) -> str: + messages: Final = tuple(item for item in items_of(at(response, "output")) if at(item, "type") == "message") + return "".join("".join(_part_texts(item)) for item in messages) + + +def _responses(response: JsonValue, streamed: str | None, errors: tuple[dict[str, JsonValue], ...], raw: str) -> Seen: + answer: Final = output_text(response) + return Seen( + ids=frozenset(_named((at(response, "id"),))), + answer=answer, + streamed=answer if streamed is None else streamed, + finishes=_named((at(response, "status"),)), + usage=( + (count_of(at(response, "usage", "input_tokens")), count_of(at(response, "usage", "output_tokens"))) + if response is not None + else None + ), + initial_prompt_tokens=None, + tokens=_named((at(response, "provider_specific_fields", "continuation_token"),)), + token_carriers=(), + content_carriers=(), + incomplete_reason=at(response, "incomplete_details", "reason"), + errors=errors, + raw=raw, + ) + + +def _responses_stream(events: tuple[dict[str, JsonValue], ...], raw: str) -> Seen: + settled: Final = tuple( + event.get("response") for event in events if event.get("type") in ("response.completed", "response.incomplete") + ) + deltas: Final = "".join( + text_of(event.get("delta")) for event in events if event.get("type") == "response.output_text.delta" + ) + return _responses(settled[-1] if settled else None, deltas, _errors(events), raw) + + +def seen(endpoint: Endpoint, stream: bool, frames: tuple[dict[str, JsonValue], ...], raw: str) -> Seen: + match (endpoint, stream): + case ("chat", False): + return _chat_unary(frames[0], raw) + case ("chat", True): + return _chat_stream(frames, raw) + case ("messages", False): + return _messages_unary(frames[0], raw) + case ("messages", True): + return _messages_stream(frames, raw) + case ("responses", False): + return _responses(frames[0], None, _errors(frames), raw) + case _: + return _responses_stream(frames, raw) + + +def answered(endpoint: Endpoint, stream: bool, response: httpx.Response) -> Seen: + assert response.status_code == 200, response.text + return seen(endpoint, stream, payloads(response.text, stream), response.text) + + +def ends_cleanly(endpoint: Endpoint, raw: str) -> bool: + if endpoint == "chat": + return raw.rstrip().endswith("data: [DONE]") + closing: Final = sse_payloads(raw)[-1].get("type") + if endpoint == "messages": + return closing == "message_stop" + return closing in ("response.completed", "response.incomplete") + + +def complete(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "stop" + case "messages": + return "end_turn" + case "responses": + return "completed" + + +def cut_off(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "length" + case "messages": + return "max_tokens" + case "responses": + return "incomplete" + + +def carries_token(endpoint: Endpoint, stream: bool) -> bool: + return endpoint == "chat" or (endpoint == "responses" and stream) + + +def single(values: frozenset[str]) -> str: + assert len(values) == 1, values + return next(iter(values)) + + +def upstream_response_id(endpoint: Endpoint, caller_id: str) -> str: + if endpoint != "responses": + return caller_id + decrypted: Final = decrypt_if_encrypted_with(caller_id.removeprefix("resp_"), SALT) + assert decrypted is not None, caller_id + issued: Final = decrypted.split(";")[0].split("response_id:")[-1] + return base64.b64decode(issued.removeprefix("resp_")).decode().split(";")[-1].removeprefix("response_id:") + + +def assert_one_id(endpoint: Endpoint, stream: bool, observed: Seen, response_id: str) -> None: + caller_id: Final = single(observed.ids) + if endpoint == "messages" and stream: + assert caller_id.startswith("msg_"), caller_id + return + assert upstream_response_id(endpoint, caller_id) == response_id, (caller_id, response_id) + + +def call_type(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "acompletion" + case "messages": + return "anthropic_messages" + case "responses": + return "aresponses" + + +def number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def cost_of(usage: tuple[int, int]) -> float: + prompt_tokens, completion_tokens = usage + return prompt_tokens * INPUT_COST + completion_tokens * OUTPUT_COST + + +def spend_rows(model: str, expected: int = 1) -> Sequence[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE model_group=%s ORDER BY request_id", (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len(rows) == expected, rows + return rows + + +def assert_billed(model: str, endpoint: Endpoint, usage: tuple[int, int], identities: frozenset[str]) -> None: + (row,) = spend_rows(model) + assert (row["call_type"], row["status"], row["prompt_tokens"], row["completion_tokens"]) == ( + call_type(endpoint), + "success", + *usage, + ), row + assert number(row["spend"]) == pytest.approx(cost_of(usage)), row + assert string_value(row["request_id"]) in identities, (row, identities) diff --git a/tests/integration/providers/test_gemini_long_decode_chaos.py b/tests/integration/providers/test_gemini_long_decode_chaos.py new file mode 100644 index 00000000000..7f86a01d69a --- /dev/null +++ b/tests/integration/providers/test_gemini_long_decode_chaos.py @@ -0,0 +1,470 @@ +from __future__ import annotations + +import asyncio +import dataclasses +import re +import threading +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._gemini_long_decode_support import ( + API_KEY, + BACKEND, + CHECKPOINT, + ENDPOINTS, + INPUT_COST, + OUTPUT_COST, + PLAIN, + PROVIDERS, + SPEND_COLUMNS, + Endpoint, + Hop, + Seen, + call_type, + carries_token, + complete, + cost_of, + cut_off, + hop_reply, + number, + path_of, + payloads, + post, + register, + request_body, + seen, + single, + streams, + summed, +) +from pydantic import JsonValue + +Continuation = Callable[[Request, str], Reply] + +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_HOP_TOKEN: Final = "hop-token-" +_PROBE_ID: Final = "probe-" +_MODES: Final = len(ENDPOINTS) * 2 + + +@dataclass(frozen=True, slots=True) +class Call: + endpoint: Endpoint + stream: bool + marker: str + model: str + + +@dataclass(frozen=True, slots=True) +class Served: + call: Call + status: int + text: str + + @property + def observed(self) -> Seen: + assert self.status == 200, self.text + return seen(self.call.endpoint, self.call.stream, payloads(self.text, self.call.stream), self.text) + + +def _opening(marker: str) -> Hop: + return Hop((f"opening marker-{marker} ", "and "), "CONTINUATION", f"{_HOP_TOKEN}{marker}", 20, 7) + + +def _closing(marker: str) -> Hop: + return Hop((f"closing marker-{marker}", "."), "STOP", None, 27, 5) + + +def _joined(marker: str) -> str: + return _opening(marker).text + _closing(marker).text + + +_OPENING_USAGE: Final = _opening("").usage +_CHAIN_USAGE: Final = summed(_opening(""), _closing("")) + + +def _marker_asked(request: Request) -> str | None: + found: Final = _MARKER.search(request.body.decode()) + return found.group(1) if found is not None else None + + +def _token_sent(request: Request) -> str | None: + token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") + return token if isinstance(token, str) else None + + +def _continues(request: Request, token: str) -> Reply: + return hop_reply(f"second-{uuid.uuid4().hex}", _closing(token.removeprefix(_HOP_TOKEN)), streams(request)) + + +def _continues_slowly(request: Request, token: str) -> Reply: + return dataclasses.replace(_continues(request, token), pause_between_chunks=0.3) + + +def _held(release: threading.Event, held: SimpleQueue[str], then: Continuation = _continues) -> Continuation: + def hold(request: Request, token: str) -> Reply: + held.put(token) + assert release.wait(timeout=90), "The continuation requests were never released" + return then(request, token) + + return hold + + +def _down_for(doomed: frozenset[str]) -> Continuation: + def answer(request: Request, token: str) -> Reply: + if token.removeprefix(_HOP_TOKEN) in doomed: + return Reply(drop_connection=True) + return _continues(request, token) + + return answer + + +def _google(continuation: Continuation = _continues) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + marker: Final = _marker_asked(request) + token: Final = _token_sent(request) + if marker is None: + return hop_reply(f"{_PROBE_ID}{uuid.uuid4().hex}", PLAIN, streams(request)) + if token is None: + return hop_reply(f"first-{marker}", _opening(marker), streams(request)) + return continuation(request, token) + + return respond + + +@contextmanager +def _released_on_exit(release: threading.Event) -> Generator[threading.Event]: + try: + yield release + finally: + release.set() + + +def _calls(count: int, models: Sequence[str], stream: Callable[[int], bool]) -> tuple[Call, ...]: + return tuple( + Call( + endpoint=ENDPOINTS[index % len(ENDPOINTS)], + stream=stream(index), + marker=uuid.uuid4().hex, + model=models[(index // _MODES) % len(models)], + ) + for index in range(count) + ) + + +def _alternating(index: int) -> bool: + return index % 2 == 0 + + +async def _send(client: httpx.AsyncClient, key: str, call: Call) -> Served: + question: Final = f"Write the long answer for marker-{call.marker}" + async with client.stream( + "POST", + path_of(call.endpoint), + json=request_body(call.endpoint, call.model, call.stream, prompt=question), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + gateway: Gateway, calls: Sequence[Call], *, tolerate_transport_errors: bool = False +) -> tuple[Served, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, gateway.key, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, Served)) + + +async def _all_mid_chain(held: SimpleQueue[str], burst: asyncio.Task[tuple[Served, ...]], count: int) -> None: + await asyncio.to_thread(eventually, lambda: held.qsize() == count or burst.done(), bool, 60) + assert held.qsize() == count, f"{held.qsize()} of {count} answers reached their continuation request" + + +def _assert_joined(served: Served) -> None: + observed: Final = served.observed + whole: Final = _joined(served.call.marker) + assert (observed.answer, observed.streamed) == (whole, whole), served.text + assert observed.finishes == (complete(served.call.endpoint),), served.text + assert (observed.usage, observed.tokens, observed.errors) == (_CHAIN_USAGE, (), ()), served.text + assert CHECKPOINT not in served.text and _HOP_TOKEN not in served.text, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def _assert_cut_off_with_its_own_token(served: Served) -> None: + observed: Final = served.observed + call: Final = served.call + so_far: Final = _opening(call.marker).text + own: Final = (f"{_HOP_TOKEN}{call.marker}",) if carries_token(call.endpoint, call.stream) else () + assert (observed.answer, observed.streamed) == (so_far, so_far), served.text + assert observed.finishes == (cut_off(call.endpoint),), served.text + assert (observed.usage, observed.tokens, observed.errors) == (_OPENING_USAGE, own, ()), served.text + assert CHECKPOINT not in served.text, served.text + assert set(_MARKER.findall(served.text)) == {call.marker}, served.text + + +def _assert_each_answer_sent_its_own_token_once(received: Sequence[Request], calls: Sequence[Call]) -> None: + sent: Final = sorted( + (marker, _token_sent(request) or "") for request in received if (marker := _marker_asked(request)) is not None + ) + assert sent == sorted( + (*((call.marker, "") for call in calls), *((call.marker, f"{_HOP_TOKEN}{call.marker}") for call in calls)) + ), sent + + +def _spend_rows(model: str) -> Sequence[Mapping[str, JsonValue]]: + return read_rows( + f"{SPEND_COLUMNS} WHERE model_group=%s AND request_id NOT LIKE '{_PROBE_ID}%%' ORDER BY request_id", (model,) + ) + + +def _billed_as(row: Mapping[str, JsonValue], served: Served, usage: tuple[int, int]) -> bool: + identities: Final = (single(served.observed.ids), f"first-{served.call.marker}") + return ( + any(same_response(string_value(row["request_id"]), identity) for identity in identities) + and (row["call_type"], row["status"], row["prompt_tokens"], row["completion_tokens"]) + == (call_type(served.call.endpoint), "success", *usage) + and number(row["spend"]) == pytest.approx(cost_of(usage)) + ) + + +def _billings(rows: Sequence[Mapping[str, JsonValue]], served: Served, usage: tuple[int, int]) -> int: + return sum(1 for row in rows if _billed_as(row, served, usage)) + + +def _assert_landed_once(model: str, joined: Sequence[Served], cut_off_early: Sequence[Served] = ()) -> None: + mine: Final = tuple( + (served, usage) + for served, usage in ( + *((served, _CHAIN_USAGE) for served in joined), + *((served, _OPENING_USAGE) for served in cut_off_early), + ) + if served.call.model == model + ) + rows: Final = eventually(lambda: _spend_rows(model), lambda found: len(found) >= len(mine), seconds=90) + assert len(rows) == len(mine), (len(rows), len(mine), rows) + unbilled: Final = tuple(served.call for served, usage in mine if _billings(rows, served, usage) != 1) + assert unbilled == (), (unbilled, rows) + + +@pytest.mark.timeout(400) +async def test_mixed_burst_of_continued_answers_keeps_every_answer_on_its_own_token_and_bills_it_once( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + with ( + wire_server(_google(_held(release, held))) as wire, + _released_on_exit(release), + gateway.scenario() as scenario, + ): + models: Final = tuple(register(gateway, scenario, provider, wire) for provider in PROVIDERS) + calls: Final = _calls(30, models, _alternating) + burst: Final = asyncio.create_task(_burst(gateway, calls)) + await _all_mid_chain(held, burst, 30) + release.set() + served: Final = await burst + assert len(served) == 30 + for item in served: + _assert_joined(item) + _assert_each_answer_sent_its_own_token_once(wire.drain(), calls) + for model in models: + _assert_landed_once(model, served) + + +@pytest.mark.timeout(400) +async def test_slow_continuation_streams_are_joined_once_each_under_concurrency(gateway: Gateway) -> None: + with wire_server(_google(_continues_slowly)) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + calls: Final = _calls(20, (model,), lambda _: True) + served: Final = await _burst(gateway, calls) + assert len(served) == 20 + for item in served: + _assert_joined(item) + _assert_each_answer_sent_its_own_token_once(wire.drain(), calls) + _assert_landed_once(model, served) + + +@pytest.mark.timeout(400) +async def test_google_outage_on_half_the_continuation_requests_mid_burst_cuts_those_off_and_recovers( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + with gateway.scenario() as scenario: + calls: Final = _calls(30, ("unregistered",), _alternating) + doomed: Final = frozenset(call.marker for call in calls[::2]) + with wire_server(_google(_held(release, held, _down_for(doomed)))) as wire, _released_on_exit(release): + model: Final = register(gateway, scenario, "gemini", wire) + routed: Final = tuple(dataclasses.replace(call, model=model) for call in calls) + burst: Final = asyncio.create_task(_burst(gateway, routed)) + await _all_mid_chain(held, burst, 30) + release.set() + served: Final = await burst + assert len(served) == 30 + lost: Final = tuple(item for item in served if item.call.marker in doomed) + kept: Final = tuple(item for item in served if item.call.marker not in doomed) + for item in lost: + _assert_cut_off_with_its_own_token(item) + for item in kept: + _assert_joined(item) + after: Final = _calls(_MODES, (model,), _alternating) + recovered: Final = await _burst(gateway, after) + assert len(recovered) == _MODES + for item in recovered: + _assert_joined(item) + _assert_each_answer_sent_its_own_token_once(wire.drain(), (*routed, *after)) + _assert_landed_once(model, (*kept, *recovered), lost) + + +def _chaos_config(wire: Wire, directory: Path, model_name: str) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "model_list": [ + { + "model_name": model_name, + "litellm_params": { + "model": f"gemini/{BACKEND}", + "api_base": wire.url, + "api_key": API_KEY, + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + ], + } + path: Final = directory / "gemini-long-decode-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _answers_probes(gateway: Gateway, model: str) -> bool: + try: + return post(gateway, "chat", request_body("chat", model, False)).status_code == 200 + except httpx.TransportError: + return False + + +def _ready(gateway: Gateway, model: str) -> None: + eventually( + lambda: tuple(_answers_probes(gateway, model) for _ in range(2)), all, seconds=graceful_stop_seconds() + 60 + ) + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _worker_pids(log: Path) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(started.group(1)) for started in _STARTED_WORKER.finditer(log.read_text())), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds() + 60, + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +async def test_worker_sigkill_mid_answer_leaves_the_sibling_joining_its_own_answers( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = f"gemini-long-decode-chaos-{uuid.uuid4().hex[:12]}" + calls: Final = _calls(20, (alias,), _alternating) + with wire_server(_google(_held(release, held))) as wire, _released_on_exit(release): + config: Final = _chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = _worker_pids(owned.log) + _ready(candidate, alias) + wire.drain() + burst: Final = asyncio.create_task(_burst(candidate, calls, tolerate_transport_errors=True)) + await _all_mid_chain(held, burst, 20) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_joined(item) + follow_up: Final = _calls(1, (alias,), _alternating) + (answered,) = await _burst(candidate, follow_up) + _assert_joined(answered) + _assert_each_answer_sent_its_own_token_once(wire.drain(), (*calls, *follow_up)) + _assert_landed_once(alias, (*served, answered)) + eventually( + lambda: owned.log.read_text().count("Application startup complete."), + lambda started: started == 3, + seconds=graceful_stop_seconds() + 60, + ) + + +@pytest.mark.timeout(4 * graceful_stop_seconds() + 240) +async def test_proxy_sigterm_mid_answer_finishes_every_answer_and_the_rebooted_proxy_bills_each_once( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = f"gemini-long-decode-chaos-{uuid.uuid4().hex[:12]}" + calls: Final = _calls(20, (alias,), _alternating) + with wire_server(_google(_held(release, held))) as wire, _released_on_exit(release): + config: Final = _chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + _worker_pids(owned.log) + _ready(owned.gateway, alias) + wire.drain() + burst: Final = asyncio.create_task(_burst(owned.gateway, calls)) + await _all_mid_chain(held, burst, 20) + owned.begin_shutdown() + release.set() + served: Final = await burst + assert len(served) == 20 + for item in served: + _assert_joined(item) + await asyncio.to_thread( + eventually, owned.process.poll, lambda code: code is not None, graceful_stop_seconds() + 60 + ) + _assert_each_answer_sent_its_own_token_once(wire.drain(), calls) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as rebooted: + _worker_pids(rebooted.log) + _ready(rebooted.gateway, alias) + wire.drain() + follow_up: Final = _calls(_MODES, (alias,), _alternating) + recovered: Final = await _burst(rebooted.gateway, follow_up) + assert len(recovered) == _MODES + for item in recovered: + _assert_joined(item) + _assert_each_answer_sent_its_own_token_once(wire.drain(), follow_up) + _assert_landed_once(alias, (*served, *recovered)) diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py new file mode 100644 index 00000000000..eccfb8e3486 --- /dev/null +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -0,0 +1,2106 @@ +from __future__ import annotations + +import ast +import contextlib +import itertools +import json +import threading +import uuid +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import anthropic +import httpx +import openai +import pytest +import yaml +from _pytest.mark.structures import ParameterSet +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, string_value +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply, posted_targets +from integration._support.otlp_sink import Span, owned_sinks, recorded_spans +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._gemini_long_decode_support import ( + API_KEY, + BACKEND, + CAP, + CHECKPOINT, + CUT, + CUT_TOKEN, + ENDPOINTS, + FIRST_HALF, + HOP_TOKEN, + NO_CACHE, + NOTHING_EXTRA, + PLAIN, + PROMPT, + PROVIDERS, + REFUSAL, + RESUMED, + SECOND_HALF, + SPEND_COLUMNS, + Answer, + Endpoint, + Hop, + Provider, + Seen, + answered, + answers, + assert_billed, + assert_one_id, + at, + carries_token, + complete, + cost_of, + cut_off, + drops_before_answering, + drops_mid_body, + ends_cleanly, + frame, + hop_reply, + items_of, + number, + output_text, + path_of, + payloads, + post, + received, + refuses, + register, + register_vertex, + request_body, + seen, + single, + spend_rows, + sse, + stream_frames, + streams, + summed, + target, + text_of, + upstream, +) +from pydantic import BaseModel, JsonValue, TypeAdapter + +Client = Literal["sdk_sync", "sdk_async", "httpx"] +Failure = Literal["400", "500", "drop", "notjson", "errorframe"] +Dropped = Literal["first", "continued"] +Unread = Literal["dropped", "empty", "unfinished", "silent"] +DeploymentTimeout = Literal["timeout", "stream_timeout"] +Absent = Literal["missing", "null", "empty"] +Respond: TypeAlias = Callable[[Request], Reply] + +_CLIENTS: Final[tuple[Client, ...]] = ("sdk_sync", "sdk_async", "httpx") +_FAILURES: Final[tuple[Failure, ...]] = ("400", "500", "drop", "notjson", "errorframe") +_UNREAD: Final[tuple[Unread, ...]] = ("dropped", "empty", "unfinished", "silent") +_DEPLOYMENT_TIMEOUTS: Final[tuple[DeploymentTimeout, ...]] = ("timeout", "stream_timeout") +_JOINED: Final = FIRST_HALF.text + SECOND_HALF.text +_CHAIN_USAGE: Final = summed(FIRST_HALF, SECOND_HALF) +_CHAT_MESSAGES: Final = ({"role": "user", "content": PROMPT},) +_USAGE_IN_STREAM: Final = {"include_usage": True} +_NATIVE_BODY: Final[Mapping[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": PROMPT}]}], + "generationConfig": {"maxOutputTokens": CAP}, +} + + +def _mode(stream: bool) -> str: + return "stream" if stream else "sync" + + +def _cells(providers: Sequence[Provider], endpoints: Sequence[Endpoint]) -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") + for provider, endpoint, stream in itertools.product(providers, endpoints, (False, True)) + ) + + +def _provider_modes() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") + for provider, stream in itertools.product(PROVIDERS, (False, True)) + ) + + +def _chain_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, endpoint, stream, client, id=f"{provider}-{endpoint}-{_mode(stream)}-{client}") + for provider, endpoint, stream, client in itertools.product(PROVIDERS, ENDPOINTS, (False, True), _CLIENTS) + ) + + +def _openai(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_openai(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_anthropic(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def _dumped(endpoint: Endpoint, stream: bool, models: Sequence[BaseModel]) -> Seen: + texts: Final = tuple(model.model_dump_json(by_alias=True, exclude_unset=True) for model in models) + return seen(endpoint, stream, tuple(JSON_OBJECT.validate_json(text) for text in texts), "\n".join(texts)) + + +def _chat_sync(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _openai(gateway) + if not stream: + reply: Final = client.chat.completions.create( + model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, extra_body=dict(NO_CACHE) + ) + return _dumped("chat", False, (reply,)) + chunks: Final = tuple( + client.chat.completions.create( + model=model, + messages=_CHAT_MESSAGES, + max_tokens=CAP, + stream=True, + stream_options=_USAGE_IN_STREAM, + extra_body=dict(NO_CACHE), + ) + ) + return _dumped("chat", True, chunks) + + +async def _chat_async(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _async_openai(gateway) + if not stream: + reply: Final = await client.chat.completions.create( + model=model, messages=_CHAT_MESSAGES, max_tokens=CAP, extra_body=dict(NO_CACHE) + ) + return _dumped("chat", False, (reply,)) + chunks: Final = [ + chunk + async for chunk in await client.chat.completions.create( + model=model, + messages=_CHAT_MESSAGES, + max_tokens=CAP, + stream=True, + stream_options=_USAGE_IN_STREAM, + extra_body=dict(NO_CACHE), + ) + ] + return _dumped("chat", True, chunks) + + +def _messages_sync(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _anthropic(gateway) + if not stream: + reply: Final = client.messages.create( + model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, extra_body=dict(NO_CACHE) + ) + return _dumped("messages", False, (reply,)) + events: Final = tuple( + client.messages.create( + model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, stream=True, extra_body=dict(NO_CACHE) + ) + ) + return _dumped("messages", True, events) + + +async def _messages_async(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _async_anthropic(gateway) + if not stream: + reply: Final = await client.messages.create( + model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, extra_body=dict(NO_CACHE) + ) + return _dumped("messages", False, (reply,)) + events: Final = [ + event + async for event in await client.messages.create( + model=model, max_tokens=CAP, messages=_CHAT_MESSAGES, stream=True, extra_body=dict(NO_CACHE) + ) + ] + return _dumped("messages", True, events) + + +def _responses_sync(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _openai(gateway) + if not stream: + reply: Final = client.responses.create( + model=model, input=PROMPT, max_output_tokens=CAP, extra_body=dict(NO_CACHE) + ) + return _dumped("responses", False, (reply,)) + events: Final = tuple( + client.responses.create( + model=model, input=PROMPT, max_output_tokens=CAP, stream=True, extra_body=dict(NO_CACHE) + ) + ) + return _dumped("responses", True, events) + + +async def _responses_async(gateway: Gateway, model: str, stream: bool) -> Seen: + client: Final = _async_openai(gateway) + if not stream: + reply: Final = await client.responses.create( + model=model, input=PROMPT, max_output_tokens=CAP, extra_body=dict(NO_CACHE) + ) + return _dumped("responses", False, (reply,)) + events: Final = [ + event + async for event in await client.responses.create( + model=model, input=PROMPT, max_output_tokens=CAP, stream=True, extra_body=dict(NO_CACHE) + ) + ] + return _dumped("responses", True, events) + + +def _over_httpx(gateway: Gateway, endpoint: Endpoint, model: str, stream: bool) -> Seen: + response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + return observed + + +async def _observe(gateway: Gateway, endpoint: Endpoint, client: Client, model: str, stream: bool) -> Seen: + match (endpoint, client): + case ("chat", "sdk_sync"): + return _chat_sync(gateway, model, stream) + case ("chat", "sdk_async"): + return await _chat_async(gateway, model, stream) + case ("messages", "sdk_sync"): + return _messages_sync(gateway, model, stream) + case ("messages", "sdk_async"): + return await _messages_async(gateway, model, stream) + case ("responses", "sdk_sync"): + return _responses_sync(gateway, model, stream) + case ("responses", "sdk_async"): + return await _responses_async(gateway, model, stream) + case _: + return _over_httpx(gateway, endpoint, model, stream) + + +def _identities(observed: Seen, response_id: str) -> frozenset[str]: + return frozenset({single(observed.ids), response_id}) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "client"), _chain_cells()) +async def test_answer_google_ends_early_is_continued_into_one_joined_answer( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, client: Client +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = await _observe(gateway, endpoint, client, model, stream) + assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed + assert observed.finishes == (complete(endpoint),), observed + assert observed.usage == _CHAIN_USAGE, observed + assert observed.initial_prompt_tokens == (FIRST_HALF.prompt_tokens if endpoint == "chat" else None), observed + assert observed.tokens == () and observed.errors == (), observed + assert CHECKPOINT not in observed.raw and HOP_TOKEN not in observed.raw, observed.raw + assert_one_id(endpoint, stream, observed, first_id) + opening, continuing = received(wire, provider, stream, (False, True)) + assert at(opening, "generationConfig", "max_output_tokens") == CAP, opening + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_answer_cut_at_the_callers_cap_returns_from_one_request_and_chat_carries_its_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + response_id: Final = f"cut-{uuid.uuid4().hex}" + with wire_server(upstream(answers(response_id, CUT))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.streamed) == (CUT.text, CUT.text), observed + assert observed.finishes == (cut_off(endpoint),), observed + assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed + assert (observed.usage, observed.initial_prompt_tokens) == (CUT.usage, None), observed + assert CHECKPOINT not in observed.raw, observed.raw + assert observed.tokens == ((CUT_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert (CUT_TOKEN in observed.raw) == carries_token(endpoint, stream), observed.raw + assert observed.token_carriers == (observed.content_carriers[-1:] if endpoint == "chat" else ()), observed + assert_one_id(endpoint, stream, observed, response_id) + (only,) = received(wire, provider, stream, (False,)) + assert at(only, "generationConfig", "max_output_tokens") == CAP, only + assert_billed(model, endpoint, CUT.usage, _identities(observed, response_id)) + + +_EMPTY_VERTEX_LISTS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_ratings": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + } +) + + +def _vertex_lists(chunk: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: value for key, value in chunk.items() if key.startswith("vertex_ai_")} + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_cut_chat_stream_hands_its_token_on_a_chunk_with_empty_vertex_ai_lists_and_no_reasoning_count( + gateway: Gateway, provider: Provider +) -> None: + with wire_server(upstream(answers(f"cut-{uuid.uuid4().hex}", CUT))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + chunks: Final = payloads(post(gateway, "chat", request_body("chat", model, True)).text, True) + carriers: Final = [chunk for chunk in chunks if "continuation_token" in json.dumps(chunk)] + assert [_vertex_lists(chunk) for chunk in carriers] == [dict(_EMPTY_VERTEX_LISTS)], chunks + assert [_vertex_lists(chunk) for chunk in chunks if chunk not in carriers] == [{}] * (len(chunks) - 1), chunks + metered: Final = [chunk for chunk in chunks if isinstance(chunk.get("usage"), dict)] + assert [at(chunk, "usage", "completion_tokens_details") for chunk in metered] == [ + {"text_tokens": CUT.output_tokens} + ], chunks + received(wire, provider, True, (False,)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_resend_with_a_continuation_token_reaches_google_and_returns_the_rest( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + resumed_id: Final = f"resumed-{uuid.uuid4().hex}" + respond: Final = upstream(answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(resumed_id, RESUMED)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, stream, {**NO_CACHE, "continuation_token": CUT_TOKEN}) + observed: Final = answered(endpoint, stream, post(gateway, endpoint, body)) + assert (observed.answer, observed.streamed) == (RESUMED.text, RESUMED.text), observed + assert observed.finishes == (complete(endpoint),), observed + assert (observed.usage, observed.tokens) == (RESUMED.usage, ()), observed + assert CHECKPOINT not in observed.raw, observed.raw + assert_one_id(endpoint, stream, observed, resumed_id) + (only,) = received(wire, provider, stream, (True,)) + assert only["continuationToken"] == CUT_TOKEN and "continuation_token" not in json.dumps(only), only + assert_billed(model, endpoint, RESUMED.usage, _identities(observed, resumed_id)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_answer_google_finishes_itself_stays_one_request_with_no_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + response_id: Final = f"plain-{uuid.uuid4().hex}" + with wire_server(upstream(answers(response_id, PLAIN))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.streamed) == (PLAIN.text, PLAIN.text), observed + assert observed.finishes == (complete(endpoint),), observed + assert (observed.usage, observed.initial_prompt_tokens) == (PLAIN.usage, None), observed + assert (observed.tokens, observed.token_carriers) == ((), ()), observed + assert CHECKPOINT not in observed.raw, observed.raw + assert_one_id(endpoint, stream, observed, response_id) + (only,) = received(wire, provider, stream, (False,)) + assert "continuationToken" not in only, only + assert_billed(model, endpoint, PLAIN.usage, _identities(observed, response_id)) + + +def _unavailable_mid_stream(stream: bool) -> Reply: + if not stream: + return refuses(503)(stream) + error: Final = {"error": {"code": 503, "message": REFUSAL, "status": "UNAVAILABLE"}} + return Reply(content_type="text/event-stream", body=b"data: " + json.dumps(error).encode() + b"\r\n\r\n") + + +def _failing(failure: Failure) -> Answer: + match failure: + case "400": + return refuses(400) + case "500": + return refuses(500) + case "drop": + return drops_before_answering + case "notjson": + return lambda _stream: Reply(body=b"not json", content_type="text/html") + case "errorframe": + return _unavailable_mid_stream + + +def _failed_cells() -> tuple[ParameterSet, ...]: + reached: Final = ( + *itertools.product(PROVIDERS, ("chat",), (False, True)), + *itertools.product(("gemini",), ("messages", "responses"), (False, True)), + ) + return tuple( + pytest.param(provider, endpoint, stream, failure, id=f"{provider}-{endpoint}-{_mode(stream)}-{failure}") + for (provider, endpoint, stream), failure in itertools.product(reached, _FAILURES) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "failure"), _failed_cells()) +def test_continuation_request_that_fails_to_open_returns_the_answer_so_far_with_its_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, failure: Failure +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with ( + wire_server(upstream(answers(first_id, FIRST_HALF), _failing(failure))) as wire, + gateway.scenario() as scenario, + ): + model: Final = register(gateway, scenario, provider, wire) + response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed + assert observed.finishes == (cut_off(endpoint),), observed + assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed + assert (observed.usage, observed.initial_prompt_tokens) == (FIRST_HALF.usage, None), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert (HOP_TOKEN in observed.raw) == carries_token(endpoint, stream) and observed.errors == (), observed.raw + assert CHECKPOINT not in observed.raw and REFUSAL not in observed.raw, observed.raw + assert len(observed.token_carriers) == (1 if endpoint == "chat" and stream else 0), observed + assert observed.token_carriers[:1] >= observed.content_carriers[-1:], observed + assert_one_id(endpoint, stream, observed, first_id) + opening, continuing = received(wire, provider, stream, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +_SECOND_TOKEN: Final = "token-after-the-second-request" +_MIDDLE: Final = Hop(("then the ", "middle, "), "CONTINUATION", _SECOND_TOKEN, 27, 5) +_LAST: Final = Hop(("and then ", "the end."), "STOP", None, 32, 4) + + +def _in_three_requests(first_id: str, last: Answer) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") + if token is None: + return hop_reply(first_id, FIRST_HALF, streams(request)) + if token == HOP_TOKEN: + return hop_reply(f"middle-{uuid.uuid4().hex}", _MIDDLE, streams(request)) + assert token == _SECOND_TOKEN, request + return last(streams(request)) + + return respond + + +def _assert_each_request_sent_the_token_before_it(wire: Wire, provider: Provider, stream: bool) -> None: + opening, second, third = received(wire, provider, stream, (False, True, True)) + assert (second, third) == ( + {**opening, "continuationToken": HOP_TOKEN}, + {**opening, "continuationToken": _SECOND_TOKEN}, + ), (opening, second, third) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_answer_google_ends_early_twice_is_joined_from_three_requests_each_sent_the_token_before_it( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + joined: Final = FIRST_HALF.text + _MIDDLE.text + _LAST.text + usage: Final = summed(FIRST_HALF, _MIDDLE, _LAST) + respond: Final = _in_three_requests(first_id, answers(f"last-{uuid.uuid4().hex}", _LAST)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.streamed) == (joined, joined), observed + assert (observed.finishes, observed.usage) == ((complete(endpoint),), usage), observed + assert observed.initial_prompt_tokens == (FIRST_HALF.prompt_tokens if endpoint == "chat" else None), observed + assert observed.tokens == () and observed.errors == (), observed + assert not any(token in observed.raw for token in (CHECKPOINT, HOP_TOKEN, _SECOND_TOKEN)), observed.raw + assert_one_id(endpoint, stream, observed, first_id) + _assert_each_request_sent_the_token_before_it(wire, provider, stream) + assert_billed(model, endpoint, usage, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_third_request_that_fails_to_open_returns_both_earlier_parts_with_the_second_requests_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + so_far: Final = FIRST_HALF.text + _MIDDLE.text + usage: Final = summed(FIRST_HALF, _MIDDLE) + with wire_server(_in_three_requests(first_id, refuses(500))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.streamed) == (so_far, so_far), observed + assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), usage), observed + assert observed.tokens == ((_SECOND_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert (_SECOND_TOKEN in observed.raw) == carries_token(endpoint, stream), observed.raw + assert not any(leaked in observed.raw for leaked in (CHECKPOINT, HOP_TOKEN, REFUSAL)), observed.raw + assert observed.errors == (), observed + assert_one_id(endpoint, stream, observed, first_id) + _assert_each_request_sent_the_token_before_it(wire, provider, stream) + assert_billed(model, endpoint, usage, _identities(observed, first_id)) + + +_CUTS: Final = 40 + + +def _part(index: int) -> Hop: + if index == _CUTS: + return Hop((f"part {index} ", "ends it."), "STOP", None, 20, 3) + return Hop((f"part {index} ", "goes on, "), "CONTINUATION", f"token-after-part-{index}", 20, 3) + + +_PARTS: Final = tuple(_part(index) for index in range(_CUTS + 1)) +_PART_AFTER: Final[Mapping[str | None, Hop]] = MappingProxyType( + {None: _PARTS[0], **{earlier.token: later for earlier, later in itertools.pairwise(_PARTS)}} +) + + +def _cut_until_the_last_part(first_id: str) -> Respond: + def respond(request: Request) -> Reply: + token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") + assert token is None or isinstance(token, str), request + response_id: Final = first_id if token is None else f"later-{uuid.uuid4().hex}" + return hop_reply(response_id, _PART_AFTER[token], streams(request)) + + return respond + + +@pytest.mark.parametrize(("provider", "stream"), _provider_modes()) +def test_answer_google_ends_early_forty_times_is_joined_from_forty_one_requests_and_billed_once( + gateway: Gateway, provider: Provider, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + joined: Final = "".join(part.text for part in _PARTS) + usage: Final = summed(*_PARTS) + with wire_server(_cut_until_the_last_part(first_id)) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = answered("chat", stream, post(gateway, "chat", request_body("chat", model, stream))) + assert (observed.answer, observed.finishes, observed.usage) == (joined, ("stop",), usage), observed + assert observed.tokens == () and observed.errors == (), observed + assert_one_id("chat", stream, observed, first_id) + sent: Final = received(wire, provider, stream, (False, *(True,) * _CUTS)) + assert [body.get("continuationToken") for body in sent] == [None, *(part.token for part in _PARTS[:-1])] + assert_billed(model, "chat", usage, _identities(observed, first_id)) + + +_TOKEN_PATH: Final = "/_oauth/token" +_EXPIRED: Final = "The scripted Google refused an access token that had expired" +_REVOKED: Final = "The scripted token endpoint revoked this service account" +_CALLERS_AUTHORIZATION: Final = "Bearer the-callers-own-credential" + + +def _vertex_cells(endpoints: Sequence[Endpoint]) -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(endpoint, stream, id=f"vertex_ai-{endpoint}-{_mode(stream)}") + for endpoint, stream in itertools.product(endpoints, (False, True)) + ) + + +def _token_reply(token: str, lifetime: int) -> Reply: + return Reply(body=json.dumps({"access_token": token, "token_type": "Bearer", "expires_in": lifetime}).encode()) + + +def _issuing_tokens(lifetime: int, google: Respond) -> Respond: + issued: Final = itertools.count() + + def respond(request: Request) -> Reply: + if request.target == _TOKEN_PATH: + return _token_reply(f"issued-{next(issued)}", lifetime) + return google(request) + + return respond + + +def _tokens_that_expire_by_the_next_request(google: Respond) -> Respond: + issued: Final = itertools.count() + unspent: Final[SimpleQueue[str]] = SimpleQueue() + + def respond(request: Request) -> Reply: + if request.target == _TOKEN_PATH: + token: Final = f"issued-{next(issued)}" + unspent.put(token) + return _token_reply(token, 0) + valid: Final = frozenset(f"Bearer {unspent.get()}" for _ in range(unspent.qsize())) + if request.headers.get("authorization") in valid: + return google(request) + error: Final = {"error": {"code": 401, "message": _EXPIRED, "status": "UNAUTHENTICATED"}} + return Reply(status=401, body=json.dumps(error).encode()) + + return respond + + +def _tokens_revoked_once_google_is_asked(google: Respond) -> Respond: + issued: Final = itertools.count() + asked: Final = threading.Event() + + def respond(request: Request) -> Reply: + if request.target != _TOKEN_PATH: + asked.set() + return google(request) + if not asked.is_set(): + return _token_reply(f"issued-{next(issued)}", 0) + refusal: Final = {"error": "invalid_grant", "error_description": _REVOKED} + return Reply(status=400, body=json.dumps(refusal).encode()) + + return respond + + +def _google_calls(requests: Sequence[Request], stream: bool, continued: tuple[bool, ...]) -> tuple[Request, ...]: + calls: Final = tuple(request for request in requests if request.target != _TOKEN_PATH) + assert [(call.method, call.target) for call in calls] == [ + ("POST", target("vertex_ai", stream, resumed)) for resumed in continued + ], requests + return calls + + +def _chain_through(token_endpoint: Callable[[Respond], Respond], first_id: str) -> Respond: + return token_endpoint(upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF))) + + +@pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(ENDPOINTS)) +def test_vertex_continuation_request_sent_after_the_access_token_expired_carries_a_new_one_and_completes( + gateway: Gateway, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with ( + wire_server(_chain_through(_tokens_that_expire_by_the_next_request, first_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = register_vertex(scenario, wire, wire.url) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed + assert (observed.finishes, observed.usage) == ((complete(endpoint),), _CHAIN_USAGE), observed + assert observed.tokens == () and observed.errors == (), observed + assert _EXPIRED not in observed.raw and HOP_TOKEN not in observed.raw, observed.raw + assert_one_id(endpoint, stream, observed, first_id) + _google_calls(wire.drain(), stream, (False, True)) + assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(ENDPOINTS)) +def test_vertex_access_token_that_cannot_be_renewed_before_the_continuation_returns_the_answer_so_far_with_its_token( + gateway: Gateway, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with ( + wire_server(_chain_through(_tokens_revoked_once_google_is_asked, first_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = register_vertex(scenario, wire, wire.url) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed + assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), FIRST_HALF.usage), observed + assert observed.incomplete_reason == ("max_output_tokens" if endpoint == "responses" else None), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert (HOP_TOKEN in observed.raw) == carries_token(endpoint, stream) and observed.errors == (), observed.raw + assert _REVOKED not in observed.raw and CHECKPOINT not in observed.raw, observed.raw + assert_one_id(endpoint, stream, observed, first_id) + _google_calls(wire.drain(), stream, (False,)) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(("chat",))) +def test_vertex_access_token_still_fresh_is_fetched_once_and_sent_on_every_request_of_a_continued_answer( + gateway: Gateway, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with ( + wire_server(_chain_through(partial(_issuing_tokens, 3600), first_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = register_vertex(scenario, wire, wire.url) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.finishes) == (_JOINED, (complete(endpoint),)), observed + requests: Final = wire.drain() + calls: Final = _google_calls(requests, stream, (False, True)) + assert [request.target for request in requests].count(_TOKEN_PATH) == 1, requests + assert len({call.headers["authorization"] for call in calls}) == 1, requests + + +@pytest.mark.parametrize(("endpoint", "stream"), _vertex_cells(("chat",))) +def test_vertex_deployments_own_authorization_header_is_sent_on_every_request_of_a_continued_answer( + gateway: Gateway, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with ( + wire_server(_chain_through(partial(_issuing_tokens, 0), first_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = register_vertex( + scenario, wire, wire.url, extra_headers={"Authorization": _CALLERS_AUTHORIZATION} + ) + observed: Final = _over_httpx(gateway, endpoint, model, stream) + assert (observed.answer, observed.finishes) == (_JOINED, (complete(endpoint),)), observed + calls: Final = _google_calls(wire.drain(), stream, (False, True)) + assert [call.headers["authorization"] for call in calls] == [_CALLERS_AUTHORIZATION] * 2, calls + + +_SILENCE_SECONDS: Final = 5 + + +def _silent(release: threading.Event) -> Answer: + def hold(_stream: bool) -> Reply: + assert release.wait(timeout=90), "The silent continuation request was never released" + return Reply(drop_connection=True) + + return hold + + +@pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) +def test_non_streamed_continuation_request_left_unanswered_returns_the_answer_so_far_at_the_callers_timeout( + gateway: Gateway, provider: Provider, endpoint: Endpoint +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + release: Final = threading.Event() + with ( + wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, + contextlib.ExitStack() as cleanup, + gateway.scenario() as scenario, + ): + cleanup.callback(release.set) + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, False, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) + observed: Final = answered(endpoint, False, post(gateway, endpoint, body)) + assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed + assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, False) else ()), observed + assert_one_id(endpoint, False, observed, first_id) + opening, continuing = received(wire, provider, False, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +_GIVE_UP_SECONDS: Final = 2 + + +def _held(release: threading.Event, answer: Answer) -> Answer: + def hold(stream: bool) -> Reply: + assert release.wait(timeout=90), "The first request was never released" + return answer(stream) + + return hold + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_non_streamed_request_its_caller_gave_up_on_is_still_continued_and_billed_for_the_whole_answer( + gateway: Gateway, provider: Provider +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + release: Final = threading.Event() + respond: Final = upstream( + _held(release, answers(first_id, FIRST_HALF)), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) + ) + with wire_server(respond) as wire, contextlib.ExitStack() as cleanup, gateway.scenario() as scenario: + cleanup.callback(release.set) + model: Final = register(gateway, scenario, provider, wire) + with pytest.raises(httpx.ReadTimeout): + gateway.client.post( + path_of("chat"), + json=request_body("chat", model, False), + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=_GIVE_UP_SECONDS, + ) + release.set() + assert_billed(model, "chat", _CHAIN_USAGE, frozenset({first_id})) + opening, continuing = received(wire, provider, False, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + + +@pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) +def test_streamed_continuation_request_left_unanswered_ends_the_stream_with_the_answer_so_far_at_the_callers_timeout( + gateway: Gateway, provider: Provider, endpoint: Endpoint +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + release: Final = threading.Event() + with ( + wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, + contextlib.ExitStack() as cleanup, + gateway.scenario() as scenario, + ): + cleanup.callback(release.set) + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, True, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) + response: Final = post(gateway, endpoint, body) + observed: Final = answered(endpoint, True, response) + assert ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed + assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed + assert_one_id(endpoint, True, observed, first_id) + opening, continuing = received(wire, provider, True, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +@pytest.mark.parametrize( + ("provider", "endpoint", "setting"), + tuple(itertools.product(PROVIDERS, ENDPOINTS, _DEPLOYMENT_TIMEOUTS)), +) +def test_streamed_continuation_request_left_unanswered_ends_at_the_deployments_timeout_with_the_answer_so_far( + gateway: Gateway, provider: Provider, endpoint: Endpoint, setting: DeploymentTimeout +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + release: Final = threading.Event() + with ( + wire_server(upstream(answers(first_id, FIRST_HALF), _silent(release))) as wire, + contextlib.ExitStack() as cleanup, + gateway.scenario() as scenario, + ): + cleanup.callback(release.set) + model: Final = register(gateway, scenario, provider, wire, settings={setting: _SILENCE_SECONDS}) + response: Final = post(gateway, endpoint, request_body(endpoint, model, True, NO_CACHE)) + observed: Final = answered(endpoint, True, response) + assert ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.finishes) == (FIRST_HALF.text, (cut_off(endpoint),)), observed + assert (observed.usage, observed.errors) == (FIRST_HALF.usage, ()), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed + assert_one_id(endpoint, True, observed, first_id) + opening, continuing = received(wire, provider, True, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_third_request_left_unanswered_returns_both_earlier_parts_and_the_second_token_at_the_callers_timeout( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + so_far: Final = FIRST_HALF.text + _MIDDLE.text + usage: Final = summed(FIRST_HALF, _MIDDLE) + release: Final = threading.Event() + with ( + wire_server(_in_three_requests(first_id, _silent(release))) as wire, + contextlib.ExitStack() as cleanup, + gateway.scenario() as scenario, + ): + cleanup.callback(release.set) + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, stream, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) + response: Final = post(gateway, endpoint, body) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.streamed) == (so_far, so_far), observed + assert (observed.finishes, observed.usage, observed.errors) == ((cut_off(endpoint),), usage, ()), observed + assert observed.tokens == ((_SECOND_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert_one_id(endpoint, stream, observed, first_id) + _assert_each_request_sent_the_token_before_it(wire, provider, stream) + assert_billed(model, endpoint, usage, _identities(observed, first_id)) + + +def _dropping(dropped: Dropped) -> Callable[[Request], Reply]: + if dropped == "first": + return upstream(drops_mid_body(f"dropped-{uuid.uuid4().hex}", PLAIN)) + return upstream( + answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), drops_mid_body(f"dropped-{uuid.uuid4().hex}", SECOND_HALF) + ) + + +def _delivered_before_the_drop(dropped: Dropped) -> str: + if dropped == "first": + return PLAIN.pieces[0] + return FIRST_HALF.text + SECOND_HALF.pieces[0] + + +def _dropped_cells() -> tuple[ParameterSet, ...]: + reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "messages"), ("gemini", "responses")) + return tuple( + pytest.param(provider, endpoint, dropped, id=f"{provider}-{endpoint}-{dropped}") + for (provider, endpoint), dropped in itertools.product(reached, ("first", "continued")) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "dropped"), _dropped_cells()) +def test_stream_that_drops_mid_body_reaches_the_caller_as_an_error_with_no_finish_and_no_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, dropped: Dropped +) -> None: + with wire_server(_dropping(dropped)) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = post(gateway, endpoint, request_body(endpoint, model, True)) + assert response.status_code == 200, response.text + observed: Final = seen(endpoint, True, payloads(response.text, True), response.text) + assert observed.errors != (), response.text + assert observed.streamed == _delivered_before_the_drop(dropped), observed + assert (observed.finishes, observed.tokens, observed.token_carriers) == ((), (), ()), observed + assert CHECKPOINT not in response.text and HOP_TOKEN not in response.text, response.text + bodies: Final = received(wire, provider, True, (False,) if dropped == "first" else (False, True)) + assert [body.get("continuationToken") for body in bodies] == [None, HOP_TOKEN][: len(bodies)], bodies + + +_UNFINISHED_FRAME: Final = b'data: {"candidates": [{"content": {"parts": [{"text": "never finished\n\n' + + +def _ends_on_an_unfinished_frame(response_id: str) -> Answer: + finished: Final = tuple(f"data: {json.dumps(item)}\n\n".encode() for item in stream_frames(response_id, FIRST_HALF)) + return lambda _stream: Reply(content_type="text/event-stream", chunks=(*finished, _UNFINISHED_FRAME)) + + +@pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) +def test_stream_that_ends_on_an_unfinished_frame_is_still_continued_into_one_joined_answer( + gateway: Gateway, provider: Provider, endpoint: Endpoint +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + respond: Final = upstream( + _ends_on_an_unfinished_frame(first_id), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) + ) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = _over_httpx(gateway, endpoint, model, True) + assert (observed.answer, observed.streamed) == (_JOINED, _JOINED), observed + assert (observed.finishes, observed.usage) == ((complete(endpoint),), _CHAIN_USAGE), observed + assert observed.tokens == () and observed.errors == (), observed + assert_one_id(endpoint, True, observed, first_id) + opening, continuing = received(wire, provider, True, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, _CHAIN_USAGE, _identities(observed, first_id)) + + +def _opens_with_no_readable_frame(unread: Unread) -> Answer: + match unread: + case "dropped": + return lambda stream: hop_reply(f"dropped-{uuid.uuid4().hex}", SECOND_HALF, stream, abort_after=0) + case "empty": + return lambda _stream: Reply(content_type="text/event-stream", chunks=()) + case "unfinished": + return lambda _stream: Reply(content_type="text/event-stream", chunks=(_UNFINISHED_FRAME,)) + case "silent": + return lambda _stream: Reply( + content_type="text/event-stream", chunks=(b"\n",), gate_after_first=threading.Event() + ) + + +def _unread_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, endpoint, unread, id=f"{provider}-{endpoint}-{unread}") + for provider, endpoint, unread in itertools.product(PROVIDERS, ENDPOINTS, _UNREAD) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "unread"), _unread_cells()) +def test_continuation_stream_that_opens_but_never_sends_a_readable_frame_ends_with_the_answer_so_far_and_its_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, unread: Unread +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + respond: Final = upstream(answers(first_id, FIRST_HALF), _opens_with_no_readable_frame(unread)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, True, {**NO_CACHE, "timeout": _SILENCE_SECONDS}) + response: Final = post(gateway, endpoint, body) + observed: Final = answered(endpoint, True, response) + assert ends_cleanly(endpoint, response.text), response.text + assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed + assert (observed.finishes, observed.usage) == ((cut_off(endpoint),), FIRST_HALF.usage), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, True) else ()), observed + assert observed.errors == () and CHECKPOINT not in response.text, response.text + assert_one_id(endpoint, True, observed, first_id) + opening, continuing = received(wire, provider, True, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(observed, first_id)) + + +_MALFORMED: Final[tuple[ParameterSet, ...]] = ( + pytest.param(12345, False, id="int-sync"), + pytest.param(12345, True, id="int-stream"), + pytest.param([CUT_TOKEN], False, id="list-sync"), + pytest.param({"token": CUT_TOKEN}, False, id="object-sync"), + pytest.param(True, False, id="bool-sync"), +) +_WRONG_TYPE: Final = "Invalid value at 'continuation_token' (TYPE_STRING)" + + +def _strict_about_the_token(resumed_id: str) -> Callable[[Request], Reply]: + refused: Final = {"error": {"code": 400, "message": _WRONG_TYPE, "status": "INVALID_ARGUMENT"}} + + def respond(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + if "continuationToken" not in body: + return hop_reply(f"restarted-{uuid.uuid4().hex}", PLAIN, streams(request)) + if not isinstance(body["continuationToken"], str): + return Reply(status=400, body=json.dumps(refused).encode()) + return hop_reply(resumed_id, RESUMED, streams(request)) + + return respond + + +@pytest.mark.parametrize(("token", "stream"), _MALFORMED) +def test_continuation_token_of_the_wrong_type_is_refused_and_never_restarts_the_answer( + gateway: Gateway, token: JsonValue, stream: bool +) -> None: + with wire_server(_strict_about_the_token(f"resumed-{uuid.uuid4().hex}")) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + refused: Final = post( + gateway, "chat", request_body("chat", model, stream, {**NO_CACHE, "continuation_token": token}) + ) + assert refused.status_code == 400, refused.text + assert "continuation_token" in refused.text and PLAIN.text not in refused.text, refused.text + sent: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain()) + assert len(sent) <= 1 and all(body.get("continuationToken") == token for body in sent), sent + healthy: Final = answered("chat", stream, post(gateway, "chat", request_body("chat", model, stream))) + assert (healthy.answer, healthy.finishes) == (PLAIN.text, ("stop",)), healthy + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_five_kilobyte_continuation_token_reaches_google_byte_for_byte(gateway: Gateway, provider: Provider) -> None: + token: Final = f"{uuid.uuid4().hex}-" + "t" * 5000 + resumed_id: Final = f"resumed-{uuid.uuid4().hex}" + with wire_server(_strict_about_the_token(resumed_id)) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body("chat", model, False, {**NO_CACHE, "continuation_token": token}) + observed: Final = answered("chat", False, post(gateway, "chat", body)) + assert (observed.answer, observed.finishes, observed.tokens) == (RESUMED.text, ("stop",), ()), observed + (only,) = received(wire, provider, False, (True,)) + assert only["continuationToken"] == token, only + + +def test_continuation_token_sent_twice_in_one_body_reaches_google_once(gateway: Gateway) -> None: + tokens: Final = (f"first-{uuid.uuid4().hex}", f"second-{uuid.uuid4().hex}") + with wire_server(_strict_about_the_token(f"resumed-{uuid.uuid4().hex}")) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + whole: Final = json.dumps(request_body("chat", model, False)) + raw: Final = f'{whole[:-1]}, "continuation_token": "{tokens[0]}", "continuation_token": "{tokens[1]}"}}' + response: Final = gateway.client.request( + "POST", + "/v1/chat/completions", + content=raw.encode(), + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + observed: Final = answered("chat", False, response) + assert observed.answer == RESUMED.text, observed + (only,) = received(wire, "gemini", False, (True,)) + assert only["continuationToken"] in tokens, only + + +def _absent(kind: Absent) -> Mapping[str, JsonValue]: + match kind: + case "missing": + return NO_CACHE + case "null": + return {**NO_CACHE, "continuation_token": None} + case "empty": + return {**NO_CACHE, "continuation_token": ""} + + +def _absent_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(kind, provider, stream, id=f"{kind}-{provider}-{_mode(stream)}") + for kind, provider, stream in itertools.product(("missing", "null", "empty"), PROVIDERS, (False, True)) + ) + + +@pytest.mark.parametrize(("kind", "provider", "stream"), _absent_cells()) +def test_missing_null_and_empty_continuation_token_all_send_a_request_with_no_token( + gateway: Gateway, kind: Absent, provider: Provider, stream: bool +) -> None: + response_id: Final = f"plain-{uuid.uuid4().hex}" + with wire_server(upstream(answers(response_id, PLAIN))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = answered( + "chat", stream, post(gateway, "chat", request_body("chat", model, stream, _absent(kind))) + ) + assert (observed.answer, observed.finishes, observed.tokens) == (PLAIN.text, ("stop",), ()), observed + assert_one_id("chat", stream, observed, response_id) + (only,) = received(wire, provider, stream, (False,)) + assert "continuationToken" not in only and "continuation_token" not in json.dumps(only), only + + +_DRAFTS: Final = (("Draft A", "token-of-draft-a"), ("Draft B", "token-of-draft-b")) + + +def _two_drafts(request: Request) -> Reply: + assert not streams(request), request.target + candidates: Final[list[JsonValue]] = [ + { + "content": {"role": "model", "parts": [{"text": text}]}, + "index": index, + "finishReason": "CONTINUATION", + "continuationToken": token, + } + for index, (text, token) in enumerate(_DRAFTS) + ] + body: Final = { + "responseId": f"drafts-{uuid.uuid4().hex}", + "candidates": candidates, + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 14, "totalTokenCount": 34}, + "modelVersion": BACKEND, + } + return Reply(body=json.dumps(body).encode()) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_request_for_two_candidates_is_never_continued_and_each_choice_keeps_its_own_token( + gateway: Gateway, provider: Provider +) -> None: + with wire_server(_two_drafts) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = post(gateway, "chat", request_body("chat", model, False, {**NO_CACHE, "n": 2})) + assert response.status_code == 200, response.text + choices: Final = items_of(JSON_OBJECT.validate_json(response.content).get("choices")) + assert [ + ( + at(choice, "message", "content"), + at(choice, "finish_reason"), + at(choice, "message", "provider_specific_fields", "continuation_token"), + ) + for choice in choices + ] == [(text, "length", token) for text, token in _DRAFTS], response.text + (only,) = received(wire, provider, False, (False,)) + assert at(only, "generationConfig", "candidate_count") == 2, only + + +def _uncapped_cells() -> tuple[ParameterSet, ...]: + reached: Final = (("gemini", "chat"), ("vertex_ai", "chat"), ("gemini", "responses")) + return tuple( + pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") + for (provider, endpoint), stream in itertools.product(reached, (False, True)) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _uncapped_cells()) +def test_request_with_no_output_cap_is_never_continued_and_returns_googles_cut_with_its_token( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + observed: Final = answered( + endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream, cap=None)) + ) + assert (observed.answer, observed.streamed) == (FIRST_HALF.text, FIRST_HALF.text), observed + assert observed.finishes == (cut_off(endpoint),), observed + assert observed.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), observed + assert observed.token_carriers == (observed.content_carriers[-1:] if endpoint == "chat" else ()), observed + assert CHECKPOINT not in observed.raw, observed.raw + (only,) = received(wire, provider, stream, (False,)) + assert at(only, "generationConfig", "max_output_tokens") is None, only + + +_NOT_A_NUMBER: Final[tuple[tuple[str, Mapping[str, JsonValue], int | None], ...]] = ( + ("n-string", {"n": "1"}, CAP), + ("n-list", {"n": [1]}, CAP), + ("n-empty", {"n": ""}, CAP), + ("n-5kb", {"n": "1" * 5120}, CAP), + ("max_tokens-string", {"max_tokens": "300"}, None), + ("max_tokens-list", {"max_tokens": [300]}, None), + ("max_tokens-empty", {"max_tokens": ""}, None), + ("max_tokens-5kb", {"max_tokens": "3" * 5120}, None), +) +_NOT_A_NUMBER_CELLS: Final[tuple[ParameterSet, ...]] = tuple( + pytest.param(extra, cap, stream, id=f"{name}-{_mode(stream)}") + for (name, extra, cap), stream in itertools.product(_NOT_A_NUMBER, (False, True)) +) + + +def _halves() -> Respond: + return upstream( + answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) + ) + + +@pytest.mark.parametrize(("extra", "cap", "stream"), _NOT_A_NUMBER_CELLS) +def test_count_or_cap_that_is_not_a_number_makes_one_request_and_is_never_continued( + gateway: Gateway, extra: Mapping[str, JsonValue], cap: int | None, stream: bool +) -> None: + with wire_server(_halves()) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + body: Final = request_body("chat", model, stream, {**NO_CACHE, **extra}, cap=cap) + observed: Final = answered("chat", stream, post(gateway, "chat", body)) + assert (observed.answer, observed.errors) == (FIRST_HALF.text, ()), observed + (only,) = received(wire, "gemini", stream, (False,)) + assert "continuationToken" not in only, only + + +def _twice_cells() -> tuple[ParameterSet, ...]: + pairs: Final = ( + ("max_tokens-number-last", f'"max_tokens": "300", "max_tokens": {CAP}', True), + ("max_tokens-string-last", f'"max_tokens": {CAP}, "max_tokens": "300"', False), + ("n-one-last", f'"max_tokens": {CAP}, "n": 2, "n": 1', True), + ("n-string-last", f'"max_tokens": {CAP}, "n": 1, "n": "1"', False), + ) + return tuple( + pytest.param(pair, continued, stream, id=f"{name}-{_mode(stream)}") + for (name, pair, continued), stream in itertools.product(pairs, (False, True)) + ) + + +@pytest.mark.parametrize(("pairs", "continued", "stream"), _twice_cells()) +def test_count_or_cap_sent_twice_in_one_body_is_read_from_its_last_value( + gateway: Gateway, pairs: str, continued: bool, stream: bool +) -> None: + with wire_server(_halves()) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + whole: Final = json.dumps(request_body("chat", model, stream, cap=None)) + response: Final = gateway.client.request( + "POST", + path_of("chat"), + content=f"{whole[:-1]}, {pairs}}}".encode(), + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + observed: Final = answered("chat", stream, response) + assert (observed.answer, observed.errors) == (_JOINED if continued else FIRST_HALF.text, ()), observed + received(wire, "gemini", stream, (False, True) if continued else (False,)) + + +@pytest.mark.parametrize(("provider", "stream"), _provider_modes()) +def test_count_of_true_is_read_as_one_answer_and_continued(gateway: Gateway, provider: Provider, stream: bool) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body("chat", model, stream, {**NO_CACHE, "n": True}) + observed: Final = answered("chat", stream, post(gateway, "chat", body)) + assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed + assert (observed.tokens, observed.errors) == ((), ()), observed + opening, continuing = received(wire, provider, stream, (False, True)) + assert continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + assert_billed(model, "chat", _CHAIN_USAGE, _identities(observed, first_id)) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _cells(PROVIDERS, ENDPOINTS)) +def test_cache_hit_twin_of_a_continued_answer_replays_it_and_bills_the_chain_once( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + replayed: Final = (_JOINED, (complete(endpoint),), _CHAIN_USAGE) + respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") + first: Final = answered(endpoint, stream, post(gateway, endpoint, body)) + assert (first.answer, first.finishes, first.usage) == replayed, first + assert len(wire.drain()) == 2 + assert_billed(model, endpoint, _CHAIN_USAGE, _identities(first, first_id)) + second: Final = answered(endpoint, stream, post(gateway, endpoint, body)) + assert (second.answer, second.finishes, second.usage) == replayed and second.tokens == (), second + assert CHECKPOINT not in second.raw and HOP_TOKEN not in second.raw, second.raw + assert wire.drain() == () + rows: Final = spend_rows(model, expected=2) + (served_from_cache,) = [row for row in rows if str(row["cache_hit"]) == "True"] + (sent_to_google,) = [row for row in rows if str(row["cache_hit"]) != "True"] + assert number(sent_to_google["spend"]) == pytest.approx(cost_of(_CHAIN_USAGE)), rows + assert (served_from_cache["status"], number(served_from_cache["spend"])) == ("success", 0), rows + + +def _rest_for(token: str) -> Hop: + return Hop((f"the rest for {token} ", "is here."), "STOP", None, 27, 5) + + +def _resumes_by_token(request: Request) -> Reply: + token: Final = JSON_OBJECT.validate_json(request.body).get("continuationToken") + hop: Final = _rest_for(token) if isinstance(token, str) else PLAIN + return hop_reply(f"answer-{uuid.uuid4().hex}", hop, streams(request)) + + +def test_request_with_a_token_is_never_served_the_cached_answer_of_one_without_it_or_with_another( + gateway: Gateway, +) -> None: + mine: Final = f"mine-{uuid.uuid4().hex}" + other: Final = f"other-{uuid.uuid4().hex}" + with wire_server(_resumes_by_token) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + body: Final = request_body("chat", model, False, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") + fresh: Final = answered("chat", False, post(gateway, "chat", body)) + spend_rows(model) + cached: Final = answered("chat", False, post(gateway, "chat", body)) + assert (fresh.answer, cached.answer, len(wire.drain())) == (PLAIN.text, PLAIN.text, 1) + resumed: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": mine})) + assert resumed.answer == _rest_for(mine).text, resumed + (sent,) = received(wire, "gemini", False, (True,)) + assert sent["continuationToken"] == mine, sent + spend_rows(model, expected=3) + replayed: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": mine})) + assert (replayed.answer, wire.drain()) == (_rest_for(mine).text, ()), replayed + another: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": other})) + assert another.answer == _rest_for(other).text, another + (crossed,) = received(wire, "gemini", False, (True,)) + assert crossed["continuationToken"] == other, crossed + + +@pytest.mark.parametrize("absent", ("", None), ids=("empty", "null")) +def test_empty_or_null_continuation_token_is_served_the_cached_answer_of_the_request_without_one( + gateway: Gateway, absent: str | None +) -> None: + with wire_server(_resumes_by_token) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, "gemini", wire) + body: Final = request_body("chat", model, False, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") + fresh: Final = answered("chat", False, post(gateway, "chat", body)) + spend_rows(model) + assert (fresh.answer, len(wire.drain())) == (PLAIN.text, 1), fresh + twin: Final = answered("chat", False, post(gateway, "chat", {**body, "continuation_token": absent})) + assert (twin.answer, wire.drain()) == (PLAIN.text, ()), twin + rows: Final = spend_rows(model, expected=2) + assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, True], rows + + +def _refused_continuation_cells() -> tuple[ParameterSet, ...]: + reached: Final = ( + *itertools.product(PROVIDERS, ("chat",), (False, True)), + *itertools.product(("gemini",), ("messages", "responses"), (False, True)), + ) + return tuple( + pytest.param(provider, endpoint, stream, id=f"{provider}-{endpoint}-{_mode(stream)}") + for provider, endpoint, stream in reached + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream"), _refused_continuation_cells()) +def test_answer_so_far_after_a_refused_continuation_is_cached_and_replayed_without_asking_google_again( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + so_far: Final = (FIRST_HALF.text, (cut_off(endpoint),), FIRST_HALF.usage) + with wire_server(upstream(answers(first_id, FIRST_HALF), refuses(500))) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + body: Final = request_body(endpoint, model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") + first: Final = answered(endpoint, stream, post(gateway, endpoint, body)) + assert (first.answer, first.finishes, first.usage) == so_far, first + assert first.tokens == ((HOP_TOKEN,) if carries_token(endpoint, stream) else ()), first + assert len(wire.drain()) == 2 + assert_billed(model, endpoint, FIRST_HALF.usage, _identities(first, first_id)) + replayed: Final = answered(endpoint, stream, post(gateway, endpoint, body)) + assert (replayed.answer, replayed.finishes, replayed.usage) == so_far, replayed + assert wire.drain() == () + rows: Final = spend_rows(model, expected=2) + assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, True], rows + + +_OPENAI_ANSWER: Final = "An answer from another provider." + + +def _openai_answering(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + return chat_reply( + f"chatcmpl-{uuid.uuid4().hex}", string_value(body["model"]), _OPENAI_ANSWER, stream=body.get("stream") is True + ) + + +def _register_openai(scenario: Scenario, wire: Wire) -> str: + return scenario.model(model="openai/gpt-5.6", api_base=f"{wire.url}/v1") + + +@pytest.mark.parametrize("stream", (False, True), ids=_mode) +def test_continuation_token_on_another_providers_request_misses_the_cached_answer_of_the_request_without_it( + gateway: Gateway, stream: bool +) -> None: + with wire_server(answering_model_discovery(_openai_answering)) as wire, gateway.scenario() as scenario: + model: Final = _register_openai(scenario, wire) + body: Final = request_body("chat", model, stream, NOTHING_EXTRA, prompt=f"{PROMPT} {uuid.uuid4().hex}") + tokened: Final = {**body, "continuation_token": CUT_TOKEN} + assert answered("chat", stream, post(gateway, "chat", body)).answer == _OPENAI_ANSWER + spend_rows(model) + assert answered("chat", stream, post(gateway, "chat", tokened)).answer == _OPENAI_ANSWER + spend_rows(model, expected=2) + assert answered("chat", stream, post(gateway, "chat", tokened)).answer == _OPENAI_ANSWER + rows: Final = spend_rows(model, expected=3) + assert posted_targets(wire) == ("/v1/chat/completions",) * 2 + assert sorted(str(row["cache_hit"]) == "True" for row in rows) == [False, False, True], rows + + +def _openai_stream_reporting_cached_prompt_tokens(request: Request) -> Reply: + chunk: Final[dict[str, JsonValue]] = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion.chunk", + "created": 1, + "model": string_value(JSON_OBJECT.validate_json(request.body)["model"]), + } + usage: Final[dict[str, JsonValue]] = { + "prompt_tokens": 5, + "completion_tokens": 3, + "total_tokens": 8, + "prompt_tokens_details": {"cached_tokens": 2}, + } + frames: Final[tuple[dict[str, JsonValue], ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": _OPENAI_ANSWER}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(f"data: {json.dumps(item)}\n\n".encode() for item in frames), b"data: [DONE]\n\n"), + ) + + +def _spend_metadata(model: str) -> dict[str, JsonValue]: + spend_rows(model) + (row,) = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)) + metadata: Final = row["metadata"] + return JSON_OBJECT.validate_json(metadata) if isinstance(metadata, str) else JSON_OBJECT.validate_python(metadata) + + +def test_streamed_answer_from_another_provider_logs_an_empty_initial_prompt_count_the_caller_never_sees( + gateway: Gateway, +) -> None: + with ( + wire_server(answering_model_discovery(_openai_stream_reporting_cached_prompt_tokens)) as wire, + gateway.scenario() as scenario, + ): + model: Final = _register_openai(scenario, wire) + response: Final = post(gateway, "chat", request_body("chat", model, True)) + assert answered("chat", True, response).answer == _OPENAI_ANSWER + assert "initial_prompt_tokens" not in response.text, response.text + logged: Final = at(_spend_metadata(model), "usage_object", "prompt_tokens_details") + assert isinstance(logged, dict) and logged.get("initial_prompt_tokens", "absent") is None, logged + + +def _native_path(model: str, stream: bool) -> str: + return f"/v1beta/models/{model}:{'streamGenerateContent' if stream else 'generateContent'}" + + +def _native(gateway: Gateway, model: str, stream: bool, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", _native_path(model, stream), dict(body), params={"alt": "sse"} if stream else None) + + +def _native_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") + for provider, stream in itertools.product(PROVIDERS, (False, True)) + ) + + +@pytest.mark.parametrize(("provider", "stream"), _native_cells()) +def test_native_route_hands_googles_continuation_fields_to_the_caller_from_one_request( + gateway: Gateway, provider: Provider, stream: bool +) -> None: + respond: Final = upstream( + answers(f"first-{uuid.uuid4().hex}", FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF) + ) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = _native(gateway, model, stream, _NATIVE_BODY) + assert response.status_code == 200, response.text + candidates: Final = tuple(at(item, "candidates", 0) for item in payloads(response.text, stream)) + assert "".join(text_of(at(item, "content", "parts", 0, "text")) for item in candidates) == FIRST_HALF.text + assert tuple(at(item, "continuationToken") for item in candidates) == ( + (CHECKPOINT, HOP_TOKEN) if stream else (HOP_TOKEN,) + ), response.text + assert at(candidates[-1], "finishReason") == "CONTINUATION", response.text + (only,) = received(wire, provider, stream, (False,)) + assert "continuationToken" not in only, only + (row,) = spend_rows(model) + assert (row["status"], row["prompt_tokens"], row["completion_tokens"]) == ("success", *FIRST_HALF.usage), row + assert number(row["spend"]) == pytest.approx(cost_of(FIRST_HALF.usage)), row + + +@pytest.mark.parametrize(("provider", "stream"), _native_cells()) +def test_native_route_passes_the_callers_continuation_token_to_google( + gateway: Gateway, provider: Provider, stream: bool +) -> None: + respond: Final = upstream( + answers(f"restarted-{uuid.uuid4().hex}", PLAIN), answers(f"resumed-{uuid.uuid4().hex}", RESUMED) + ) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + response: Final = _native(gateway, model, stream, {**_NATIVE_BODY, "continuationToken": CUT_TOKEN}) + assert response.status_code == 200, response.text + candidates: Final = tuple(at(item, "candidates", 0) for item in payloads(response.text, stream)) + assert "".join(text_of(at(item, "content", "parts", 0, "text")) for item in candidates) == RESUMED.text + (sent,) = wire.drain() + assert sent.target in (target(provider, stream), target(provider, stream, continued=True)), sent.target + assert JSON_OBJECT.validate_json(sent.body)["continuationToken"] == CUT_TOKEN, sent.body + + +Searched = Literal["first", "each"] + +_SEARCHES: Final = ("scripted search one", "scripted search two") + + +def _grounded(reply: dict[str, JsonValue], searches: Sequence[str]) -> dict[str, JsonValue]: + candidate: Final = at(reply, "candidates", 0) + assert isinstance(candidate, dict), reply + grounding: Final[dict[str, JsonValue]] = { + "webSearchQueries": [*searches], + "groundingChunks": [{"web": {"uri": "https://example.com/source", "title": "example.com"}}], + } + return {**reply, "candidates": [{**candidate, "groundingMetadata": grounding}]} + + +def _searching(response_id: str, hop: Hop, searches: Sequence[str]) -> Answer: + def reply(stream: bool) -> Reply: + if not searches: + return hop_reply(response_id, hop, stream) + if stream: + opening, closing = stream_frames(response_id, hop) + return sse((opening, _grounded(closing, searches))) + whole: Final = frame(response_id, hop.text, hop.usage, hop.finish, hop.token) + return Reply(body=json.dumps(_grounded(whole, searches)).encode()) + + return reply + + +def _searches_counted(response: httpx.Response, stream: bool) -> JsonValue: + metered: Final = tuple(item for item in payloads(response.text, stream) if isinstance(item.get("usage"), dict)) + assert metered, response.text + return at(metered[-1], "usage", "prompt_tokens_details", "web_search_requests") + + +def _search_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, stream, searched, id=f"{provider}-{_mode(stream)}-{searched}") + for provider, stream, searched in itertools.product(PROVIDERS, (False, True), ("first", "each")) + ) + + +@pytest.mark.parametrize(("provider", "stream", "searched"), _search_cells()) +def test_google_searches_across_a_continued_answer_are_all_counted_and_billed_as_one_request_would_be( + gateway: Gateway, provider: Provider, stream: bool, searched: Searched +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + opening_searches: Final = _SEARCHES if searched == "first" else _SEARCHES[:1] + chained: Final = upstream( + _searching(first_id, FIRST_HALF, opening_searches), + _searching(f"second-{uuid.uuid4().hex}", SECOND_HALF, _SEARCHES[len(opening_searches) :]), + ) + in_one_request: Final = upstream(_searching(f"plain-{uuid.uuid4().hex}", PLAIN, _SEARCHES)) + with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + twin: Final = register(gateway, scenario, provider, twin_wire) + searching: Final[Mapping[str, JsonValue]] = {**NO_CACHE, "tools": [{"googleSearch": {}}]} + response: Final = post(gateway, "chat", request_body("chat", model, stream, searching)) + observed: Final = answered("chat", stream, response) + assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed + assert _searches_counted(response, stream) == len(_SEARCHES), response.text + opening, continuing = received(wire, provider, stream, (False, True)) + assert "tools" in opening and continuing == {**opening, "continuationToken": HOP_TOKEN}, (opening, continuing) + alone: Final = post(gateway, "chat", request_body("chat", twin, stream, searching)) + assert answered("chat", stream, alone).answer == PLAIN.text, alone.text + assert _searches_counted(alone, stream) == len(_SEARCHES), alone.text + (chain_row,) = spend_rows(model) + (twin_row,) = spend_rows(twin) + search_fee: Final = number(twin_row["spend"]) - cost_of(PLAIN.usage) + assert search_fee > 0, twin_row + assert number(chain_row["spend"]) == pytest.approx(cost_of(_CHAIN_USAGE) + search_fee), (chain_row, twin_row) + + +Spelling = Literal["citations", "citationSources"] +Earlier = Literal["multibyte", "thought"] + +_SPELLINGS: Final[tuple[Spelling, ...]] = ("citations", "citationSources") +_EARLIER: Final[tuple[Earlier, ...]] = ("multibyte", "thought") +_CITED_OPENING: Final = Hop(("Grüße aus 東京 🌸, ", "wo der Fuji steht, "), "CONTINUATION", HOP_TOKEN, 20, 7) +_CITED_CLOSING: Final = Hop(("dann der ", "zweite Teil über Zürich."), "STOP", None, 27, 5) +_CITED_WHOLE: Final = Hop((_CITED_OPENING.text, _CITED_CLOSING.text), "STOP", None, 20, 12) +_THOUGHT: Final = "Ich plane die Antwort über 東京 zuerst." + + +@dataclass(frozen=True, slots=True) +class _Source: + phrase: str + uri: str + + +@dataclass(frozen=True, slots=True) +class _Carried: + answer: str + annotations: tuple[str, ...] + citations: tuple[str, ...] + retrieved: tuple[str, ...] + + +_OPENING_SOURCE: Final = _Source("der Fuji", "https://example.com/fuji") +_CLOSING_SOURCE: Final = _Source("über Zürich", "https://example.com/zurich") + + +def _byte_span(text: str, phrase: str) -> dict[str, JsonValue]: + # Google measures citation and grounding segment offsets in UTF-8 bytes (ai.google.dev/api/generate-content, + # Segment and CitationSource, read 2026-10-14) + start: Final = text.encode().index(phrase.encode()) + return {"startIndex": start, "endIndex": start + len(phrase.encode())} + + +def _cited_metadata(text: str, sources: Sequence[_Source], spelling: Spelling) -> dict[str, JsonValue]: + return { + "citationMetadata": {spelling: [{**_byte_span(text, source.phrase), "uri": source.uri} for source in sources]}, + "urlContextMetadata": { + "urlMetadata": [ + {"retrievedUrl": source.uri, "urlRetrievalStatus": "URL_RETRIEVAL_STATUS_SUCCESS"} for source in sources + ] + }, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": source.uri, "title": source.phrase}} for source in sources], + "groundingSupports": [ + { + "segment": {**_byte_span(text, source.phrase), "text": source.phrase}, + "groundingChunkIndices": [index], + } + for index, source in enumerate(sources) + ], + }, + } + + +def _opening_metadata(spelling: Spelling) -> dict[str, JsonValue]: + return _cited_metadata(_CITED_OPENING.text, (_OPENING_SOURCE,), spelling) + + +def _decorated( + reply: dict[str, JsonValue], thought: str | None, metadata: Mapping[str, JsonValue] +) -> dict[str, JsonValue]: + candidate: Final = at(reply, "candidates", 0) + assert isinstance(candidate, dict), reply + thoughts: Final[tuple[JsonValue, ...]] = ({"text": thought, "thought": True},) if thought else () + content: Final[dict[str, JsonValue]] = { + "role": "model", + "parts": [*thoughts, *items_of(at(candidate, "content", "parts"))], + } + return {**reply, "candidates": [{**candidate, "content": content, **metadata}]} + + +def _citing(response_id: str, hop: Hop, metadata: Mapping[str, JsonValue], thought: str | None) -> Answer: + def reply(stream: bool) -> Reply: + if stream: + opening, closing = stream_frames(response_id, hop) + return sse((_decorated(opening, thought, NOTHING_EXTRA), _decorated(closing, None, metadata))) + whole: Final = frame(response_id, hop.text, hop.usage, hop.finish, hop.token) + return Reply(body=json.dumps(_decorated(whole, thought, metadata)).encode()) + + return reply + + +def _canonical(values: Iterable[JsonValue]) -> tuple[str, ...]: + return tuple(sorted(json.dumps(value, sort_keys=True) for value in values)) + + +def _entries(listed: JsonValue, *keys: str) -> Iterator[JsonValue]: + return itertools.chain.from_iterable( + items_of(at(entry, key)) for entry, key in itertools.product(items_of(listed), keys) + ) + + +def _chat_carried(payloads_seen: Sequence[dict[str, JsonValue]], message: str) -> _Carried: + return _Carried( + answer="".join(text_of(at(payload, "choices", 0, message, "content")) for payload in payloads_seen), + annotations=_canonical( + itertools.chain.from_iterable( + items_of(at(payload, "choices", 0, message, "annotations")) for payload in payloads_seen + ) + ), + citations=_canonical( + itertools.chain.from_iterable( + _entries(payload.get("vertex_ai_citation_metadata"), "citations", "citationSources") + for payload in payloads_seen + ) + ), + retrieved=_canonical( + itertools.chain.from_iterable( + _entries(payload.get("vertex_ai_url_context_metadata"), "urlMetadata") for payload in payloads_seen + ) + ), + ) + + +def _output_parts(response: JsonValue) -> Iterator[JsonValue]: + messages: Final = (item for item in items_of(at(response, "output")) if at(item, "type") == "message") + return itertools.chain.from_iterable(items_of(at(item, "content")) for item in messages) + + +def _responses_carried(response: JsonValue) -> _Carried: + return _Carried( + answer=output_text(response), + annotations=_canonical( + itertools.chain.from_iterable(items_of(at(part, "annotations")) for part in _output_parts(response)) + ), + citations=(), + retrieved=(), + ) + + +def _completed(events: Sequence[dict[str, JsonValue]]) -> JsonValue: + (completed,) = (event.get("response") for event in events if event.get("type") == "response.completed") + return completed + + +def _carried(endpoint: Endpoint, stream: bool, response: httpx.Response) -> _Carried: + assert response.status_code == 200, response.text + seen_payloads: Final = payloads(response.text, stream) + match (endpoint, stream): + case ("chat", False): + return _chat_carried(seen_payloads, "message") + case ("chat", True): + return _chat_carried(seen_payloads, "delta") + case ("responses", False): + return _responses_carried(seen_payloads[0]) + case ("responses", True): + return _responses_carried(_completed(seen_payloads)) + case _: + raise AssertionError(f"No citations are read from {endpoint}") + + +def _citation_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param( + provider, + endpoint, + stream, + spelling, + earlier, + id=f"{provider}-{endpoint}-{_mode(stream)}-{spelling}-{earlier}", + ) + for provider, endpoint, stream, spelling, earlier in itertools.product( + PROVIDERS, ("chat", "responses"), (False, True), _SPELLINGS, _EARLIER + ) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling", "earlier"), _citation_cells()) +def test_citations_and_url_context_of_a_continued_answer_match_what_one_request_for_the_whole_answer_returns( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling, earlier: Earlier +) -> None: + thought: Final = _THOUGHT if earlier == "thought" else None + chained: Final = upstream( + _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, _opening_metadata(spelling), thought), + _citing( + f"second-{uuid.uuid4().hex}", + _CITED_CLOSING, + _cited_metadata(_CITED_CLOSING.text, (_CLOSING_SOURCE,), spelling), + None, + ), + ) + in_one_request: Final = upstream( + _citing( + f"plain-{uuid.uuid4().hex}", + _CITED_WHOLE, + _cited_metadata(_CITED_WHOLE.text, (_OPENING_SOURCE, _CLOSING_SOURCE), spelling), + thought, + ) + ) + with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + twin: Final = register(gateway, scenario, provider, twin_wire) + continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) + alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) + assert alone.answer == _CITED_WHOLE.text and len(alone.annotations) == 2, alone + assert endpoint != "chat" or (len(alone.citations), len(alone.retrieved)) == (2, 2), alone + assert continued == alone, (continued, alone) + assert len(wire.drain()) == 2 + + +def _unplaced_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param(provider, endpoint, stream, spelling, id=f"{provider}-{endpoint}-{_mode(stream)}-{spelling}") + for provider, endpoint, stream, spelling in itertools.product( + PROVIDERS, ("chat", "responses"), (False, True), _SPELLINGS + ) + ) + + +def _cited_uris(carried: _Carried) -> frozenset[JsonValue]: + return frozenset(JSON_OBJECT.validate_json(cited).get("uri") for cited in carried.citations) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling"), _unplaced_cells()) +def test_later_request_citing_a_source_with_no_offsets_returns_what_one_request_for_the_whole_answer_returns( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling +) -> None: + # Google marks every CitationSource field optional, endIndex included (ai.google.dev/api/generate-content, + # CitationSource, read 2026-10-14) + unplaced: Final[JsonValue] = {"uri": _CLOSING_SOURCE.uri} + placed: Final[JsonValue] = {**_byte_span(_CITED_OPENING.text, _OPENING_SOURCE.phrase), "uri": _OPENING_SOURCE.uri} + chained: Final = upstream( + _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, {"citationMetadata": {spelling: [placed]}}, None), + _citing(f"second-{uuid.uuid4().hex}", _CITED_CLOSING, {"citationMetadata": {spelling: [unplaced]}}, None), + ) + in_one_request: Final = upstream( + _citing(f"plain-{uuid.uuid4().hex}", _CITED_WHOLE, {"citationMetadata": {spelling: [placed, unplaced]}}, None) + ) + with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + twin: Final = register(gateway, scenario, provider, twin_wire) + continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) + alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) + assert alone.answer == continued.answer == _CITED_WHOLE.text, (continued, alone) + assert stream or endpoint != "chat" or _OPENING_SOURCE.uri in _cited_uris(alone), alone + assert _cited_uris(continued) == _cited_uris(alone), (continued, alone) + assert len(wire.drain()) == 2 + + +_FIRST_BYTE_SOURCE: Final = _Source("dann der", "https://example.com/dann") + + +def _cited_from_first_byte(text: str, source: _Source, spelling: Spelling) -> dict[str, JsonValue]: + assert text.startswith(source.phrase), text + end: Final = len(source.phrase.encode()) + return { + "citationMetadata": {spelling: [{"endIndex": end, "uri": source.uri}]}, + "urlContextMetadata": { + "urlMetadata": [{"retrievedUrl": source.uri, "urlRetrievalStatus": "URL_RETRIEVAL_STATUS_SUCCESS"}] + }, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": source.uri, "title": source.phrase}}], + "groundingSupports": [{"segment": {"endIndex": end, "text": source.phrase}, "groundingChunkIndices": [0]}], + }, + } + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "spelling"), _unplaced_cells()) +def test_later_request_span_from_its_first_byte_with_no_start_index_points_where_one_request_for_the_whole_answer_does( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, spelling: Spelling +) -> None: + chained: Final = upstream( + _citing(f"first-{uuid.uuid4().hex}", _CITED_OPENING, _opening_metadata(spelling), None), + _citing( + f"second-{uuid.uuid4().hex}", + _CITED_CLOSING, + _cited_from_first_byte(_CITED_CLOSING.text, _FIRST_BYTE_SOURCE, spelling), + None, + ), + ) + in_one_request: Final = upstream( + _citing( + f"plain-{uuid.uuid4().hex}", + _CITED_WHOLE, + _cited_metadata(_CITED_WHOLE.text, (_OPENING_SOURCE, _FIRST_BYTE_SOURCE), spelling), + None, + ) + ) + with wire_server(chained) as wire, wire_server(in_one_request) as twin_wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire) + twin: Final = register(gateway, scenario, provider, twin_wire) + continued: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, model, stream))) + alone: Final = _carried(endpoint, stream, post(gateway, endpoint, request_body(endpoint, twin, stream))) + assert alone.answer == _CITED_WHOLE.text and len(alone.annotations) == 2, alone + assert endpoint != "chat" or (len(alone.citations), len(alone.retrieved)) == (2, 2), alone + assert continued == alone, (continued, alone) + assert len(wire.drain()) == 2 + + +OddMetadata = Literal["citations_null", "citations_empty", "url_context_null", "url_context_empty", "both_missing"] + +_ODD_METADATA: Final[Mapping[OddMetadata, Mapping[str, JsonValue]]] = MappingProxyType( + { + "citations_null": {"citationMetadata": None}, + "citations_empty": {"citationMetadata": {}}, + "url_context_null": {"urlContextMetadata": None}, + "url_context_empty": {"urlContextMetadata": {}}, + "both_missing": NOTHING_EXTRA, + } +) +_OLDER_BACKEND: Final = "gemini-3.8-flash" +_ODD_ANSWERS: Final[tuple[tuple[str, tuple[Hop, ...]], ...]] = ( + (BACKEND, (PLAIN,)), + (BACKEND, (FIRST_HALF, SECOND_HALF)), + (_OLDER_BACKEND, (PLAIN,)), + (_OLDER_BACKEND, (FIRST_HALF, SECOND_HALF)), +) + + +def _odd(response_id: str, hop: Hop, backend: str, metadata: Mapping[str, JsonValue]) -> Answer: + def reply(stream: bool) -> Reply: + frames: Final = ( + stream_frames(response_id, hop) + if stream + else (frame(response_id, hop.text, hop.usage, hop.finish, hop.token),) + ) + odd: Final = tuple({**_decorated(item, None, metadata), "modelVersion": backend} for item in frames) + return sse(odd) if stream else Reply(body=json.dumps(odd[0]).encode()) + + return reply + + +def _odd_cells() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param( + provider, + endpoint, + stream, + backend, + hops, + odd, + id=f"{provider}-{endpoint}-{_mode(stream)}-{backend}-{len(hops)}-requests-{odd}", + ) + for provider, endpoint, stream, (backend, hops), odd in itertools.product( + PROVIDERS, ENDPOINTS, (False, True), _ODD_ANSWERS, _ODD_METADATA + ) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "backend", "hops", "odd"), _odd_cells()) +def test_null_empty_or_missing_citation_and_url_context_metadata_still_returns_the_whole_answer( + gateway: Gateway, + provider: Provider, + endpoint: Endpoint, + stream: bool, + backend: str, + hops: tuple[Hop, ...], + odd: OddMetadata, +) -> None: + replies: Final = tuple(_odd(f"odd-{uuid.uuid4().hex}", hop, backend, _ODD_METADATA[odd]) for hop in hops) + with wire_server(upstream(*replies)) as wire, gateway.scenario() as scenario: + model: Final = register(gateway, scenario, provider, wire, backend=backend) + response: Final = post(gateway, endpoint, request_body(endpoint, model, stream)) + observed: Final = answered(endpoint, stream, response) + assert not stream or ends_cleanly(endpoint, response.text), response.text + assert observed.answer == "".join(hop.text for hop in hops) and observed.errors == (), observed + assert observed.finishes == (complete(endpoint),), observed + assert len(wire.drain()) == len(hops) + + +_OTHER_BODY: Final[Mapping[str, JsonValue]] = { + **_NATIVE_BODY, + "contents": [{"role": "user", "parts": [{"text": "Write the short answer."}]}], +} + + +def _pass_through_config(directory: Path, wire_url: str) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "gemini-pass-through.yaml" + path.write_text( + yaml.safe_dump({**base, "environment_variables": {"GEMINI_API_BASE": wire_url, "GEMINI_API_KEY": API_KEY}}) + ) + return path + + +def _scripted(script: Sequence[tuple[Mapping[str, JsonValue], Reply]]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + sent: Final = JSON_OBJECT.validate_json(request.body) + matching: Final = tuple(reply for body, reply in script if body == sent) + assert len(matching) == 1, request + return matching[0] + + return respond + + +def _sent_by_google(reply: Reply, stream: bool) -> tuple[dict[str, JsonValue], ...]: + return payloads(b"".join(reply.chunks or (reply.body,)).decode(), stream) + + +def _pass_through(gateway: Gateway, stream: bool, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + f"/gemini{_native_path(BACKEND, stream)}", + dict(body), + params={"alt": "sse"} if stream else None, + headers={"x-goog-api-key": gateway.key}, + ) + + +def _pass_through_row(response: httpx.Response) -> dict[str, JsonValue]: + (row,) = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE request_id=%s", (response.headers["x-litellm-call-id"],)), + lambda found: len(found) >= 1, + seconds=70, + ) + return row + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("stream", (False, True), ids=_mode) +def test_gemini_pass_through_route_relays_googles_continuation_fields_as_sent_and_bills_each_request_on_its_own( + gateway: Gateway, tmp_path: Path, stream: bool +) -> None: + script: Final = ( + (_NATIVE_BODY, hop_reply("pass-through-opening", FIRST_HALF, stream)), + ({**_NATIVE_BODY, "continuationToken": HOP_TOKEN}, hop_reply("pass-through-rest", SECOND_HALF, stream)), + (_OTHER_BODY, hop_reply("pass-through-other", PLAIN, stream)), + ) + with wire_server(_scripted(script)) as wire: + with owned_proxy_process(gateway, tmp_path, {}, config=_pass_through_config(tmp_path, wire.url)) as owned: + responses: Final = tuple(_pass_through(owned.gateway, stream, body) for body, _ in script) + assert [response.status_code for response in responses] == [200] * len(script), [ + response.text for response in responses + ] + assert [payloads(response.text, stream) for response in responses] == [ + _sent_by_google(reply, stream) for _, reply in script + ] + requests: Final = wire.drain() + assert [(request.method, request.target.split("?", 1)[0]) for request in requests] == [ + ("POST", _native_path(BACKEND, stream)) + ] * len(script), requests + assert [JSON_OBJECT.validate_json(request.body) for request in requests] == [body for body, _ in script] + rows: Final = tuple(_pass_through_row(response) for response in responses) + assert [ + (row["call_type"], row["status"], row["prompt_tokens"], row["completion_tokens"]) for row in rows + ] == [("pass_through_endpoint", "success", *hop.usage) for hop in (FIRST_HALF, SECOND_HALF, PLAIN)], rows + opening_spend, rest_spend, other_spend = (number(row["spend"]) for row in rows) + assert opening_spend == pytest.approx(other_spend), rows + assert min(opening_spend, rest_spend) > 0, rows + + +_LOGGED_CANDIDATES: Final = TypeAdapter(list[JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Traced: + gateway: Gateway + sink: str + + +def _traced_config(directory: Path, sink: str) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = JSON_OBJECT.validate_python(base.get("litellm_settings", {})) + path: Final = directory / "gemini-long-decode-otel.yaml" + path.write_text( + yaml.safe_dump( + { + **base, + "litellm_settings": {**settings, "callbacks": ["otel"]}, + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": sink, "use_simple_processor": True} + }, + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def traced(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Traced]: + directory: Final = tmp_path_factory.mktemp("long-decode-otel") + with owned_sinks(directory) as sinks: + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} + config: Final = _traced_config(directory, sinks.operator) + with owned_proxy_process(session_gateway, directory, overrides, config=config) as owned: + yield _Traced(owned.gateway, sinks.operator) + + +def _raw_request_spans(sink: str, provider: Provider, marker: str) -> tuple[Span, ...]: + return tuple( + span + for span in recorded_spans(sink)[1] + if span["name"] == "raw_gen_ai_request" and marker in str(span["attributes"].get(f"llm.{provider}.contents")) + ) + + +def _logged(span: Span, provider: Provider, field: str) -> str: + return string_value(span["attributes"][f"llm.{provider}.{field}"]) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("last", (SECOND_HALF, CUT), ids=("finished", "cut")) +@pytest.mark.parametrize(("provider", "endpoint"), tuple(itertools.product(PROVIDERS, ENDPOINTS))) +def test_continued_answers_logged_raw_response_is_one_google_response_holding_the_whole_answer( + traced: _Traced, provider: Provider, endpoint: Endpoint, last: Hop +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + marker: Final = f"{PROMPT} {uuid.uuid4().hex}" + respond: Final = upstream(answers(first_id, FIRST_HALF), answers(f"last-{uuid.uuid4().hex}", last)) + with wire_server(respond) as wire, traced.gateway.scenario() as scenario: + model: Final = register(traced.gateway, scenario, provider, wire) + response: Final = post(traced.gateway, endpoint, request_body(endpoint, model, False, prompt=marker)) + assert response.status_code == 200, response.text + (span,) = eventually( + lambda: _raw_request_spans(traced.sink, provider, marker), + lambda found: len(found) >= 1, + seconds=60, + return_last_on_timeout=True, + ) + assert span["attributes"].get(f"llm.{provider}.responseId") == first_id, span + (candidate,) = _LOGGED_CANDIDATES.validate_python(ast.literal_eval(_logged(span, provider, "candidates"))) + parts: Final = items_of(at(candidate, "content", "parts")) + assert "".join(text_of(at(part, "text")) for part in parts) == FIRST_HALF.text + last.text, candidate + ending: Final = (at(candidate, "finishReason"), at(candidate, "continuationToken")) + assert ending == (last.finish, last.token), candidate + usage: Final = JSON_OBJECT.validate_python(ast.literal_eval(_logged(span, provider, "usageMetadata"))) + assert (at(usage, "promptTokenCount"), at(usage, "candidatesTokenCount")) == summed(FIRST_HALF, last), usage diff --git a/tests/integration/sdk/test_gemini_long_decode_sync_sdk.py b/tests/integration/sdk/test_gemini_long_decode_sync_sdk.py new file mode 100644 index 00000000000..146283c7766 --- /dev/null +++ b/tests/integration/sdk/test_gemini_long_decode_sync_sdk.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import itertools +import json +import threading +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from queue import SimpleQueue +from typing import Final + +import pytest +from _pytest.mark.structures import ParameterSet +from integration._support.client import JSON_OBJECT +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._gemini_long_decode_support import ( + API_KEY, + BACKEND, + CAP, + CHECKPOINT, + FIRST_HALF, + HOP_TOKEN, + LOCATION, + MODEL_PATH, + PROJECT, + PROMPT, + PROVIDERS, + SECOND_HALF, + Provider, + Seen, + answers, + seen, + summed, + target, + upstream, +) +from pydantic import JsonValue + +import litellm +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import ModelResponse + +Respond = Callable[[Request], Reply] + +_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" +_JOINED: Final = FIRST_HALF.text + SECOND_HALF.text +_CHAIN_USAGE: Final = summed(FIRST_HALF, SECOND_HALF) + + +@dataclass(frozen=True, slots=True) +class _Route: + api_base: str + api_key: str | None + vertex_project: str | None + vertex_location: str | None + vertex_credentials: str | None + + +def _mode(stream: bool) -> str: + return "stream" if stream else "unary" + + +_CELLS: Final[tuple[ParameterSet, ...]] = tuple( + pytest.param(provider, stream, id=f"{provider}-{_mode(stream)}") + for provider, stream in itertools.product(PROVIDERS, (False, True)) +) +_VERTEX_CELLS: Final[tuple[ParameterSet, ...]] = tuple( + pytest.param(stream, id=f"vertex_ai-{_mode(stream)}") for stream in (False, True) +) + + +def _route(provider: Provider, wire: Wire) -> _Route: + if provider == "gemini": + return _Route(wire.url, API_KEY, None, None, None) + return _Route(f"{wire.url}{MODEL_PATH}", None, PROJECT, LOCATION, service_account_json(PROJECT, wire.url)) + + +def _complete(provider: Provider, wire: Wire, stream: bool) -> Seen: + route: Final = _route(provider, wire) + response: Final = litellm.completion( + model=f"{provider}/{BACKEND}", + messages=[{"role": "user", "content": PROMPT}], + max_tokens=CAP, + stream=stream, + stream_options={"include_usage": True} if stream else None, + api_base=route.api_base, + api_key=route.api_key, + vertex_project=route.vertex_project, + vertex_location=route.vertex_location, + vertex_credentials=route.vertex_credentials, + num_retries=0, + ) + if isinstance(response, CustomStreamWrapper): + chunks: Final = tuple(JSON_OBJECT.validate_json(chunk.model_dump_json()) for chunk in response) + return seen("chat", True, chunks, "\n".join(json.dumps(chunk) for chunk in chunks)) + assert isinstance(response, ModelResponse), type(response) + payload: Final = JSON_OBJECT.validate_json(response.model_dump_json()) + return seen("chat", False, (payload,), json.dumps(payload)) + + +def _token_reply(token: str, lifetime: int) -> Reply: + return Reply(body=json.dumps({"access_token": token, "token_type": "Bearer", "expires_in": lifetime}).encode()) + + +def _tokens_lasting_an_hour(google: Respond) -> Respond: + issued: Final = itertools.count() + + def respond(request: Request) -> Reply: + if request.target == _TOKEN_PATH: + return _token_reply(f"issued-{next(issued)}", 3600) + 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 _chain(first_id: str) -> Respond: + return upstream(answers(first_id, FIRST_HALF), answers(f"second-{uuid.uuid4().hex}", SECOND_HALF)) + + +def _google_calls(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if request.target != _TOKEN_PATH) + + +def _sent_tokens(calls: tuple[Request, ...]) -> list[JsonValue]: + return [JSON_OBJECT.validate_json(call.body).get("continuationToken") for call in calls] + + +@pytest.mark.parametrize(("provider", "stream"), _CELLS) +def test_sync_completion_joins_an_answer_google_cut_with_a_continuation_token(provider: Provider, stream: bool) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with wire_server(_tokens_lasting_an_hour(_chain(first_id))) as wire: + observed: Final = _complete(provider, wire, stream) + calls: Final = _google_calls(wire.drain()) + assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed + assert observed.tokens == () and observed.ids == frozenset({first_id}), observed + assert CHECKPOINT not in observed.raw and HOP_TOKEN not in observed.raw, observed.raw + assert [call.target for call in calls] == [target(provider, stream), target(provider, stream, True)], calls + assert _sent_tokens(calls) == [None, HOP_TOKEN], calls + + +@pytest.mark.parametrize("stream", _VERTEX_CELLS) +def test_sync_vertex_continuation_sent_after_the_access_token_expired_carries_a_new_one_and_completes( + stream: bool, +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with wire_server(_tokens_that_expire_by_the_next_request(_chain(first_id))) as wire: + observed: Final = _complete("vertex_ai", wire, stream) + calls: Final = _google_calls(wire.drain()) + assert (observed.answer, observed.finishes, observed.usage) == (_JOINED, ("stop",), _CHAIN_USAGE), observed + assert observed.tokens == () and _EXPIRED not in observed.raw, observed.raw + assert [call.target for call in calls] == [target("vertex_ai", stream), target("vertex_ai", stream, True)] + assert len({call.headers.get("authorization") for call in calls}) == 2, calls + assert _sent_tokens(calls) == [None, HOP_TOKEN], calls + + +@pytest.mark.parametrize("stream", _VERTEX_CELLS) +def test_sync_vertex_access_token_that_cannot_be_renewed_returns_the_answer_so_far_with_its_token( + stream: bool, +) -> None: + first_id: Final = f"first-{uuid.uuid4().hex}" + with wire_server(_tokens_revoked_once_google_is_asked(_chain(first_id))) as wire: + observed: Final = _complete("vertex_ai", wire, stream) + calls: Final = _google_calls(wire.drain()) + assert (observed.answer, observed.finishes, observed.usage) == (FIRST_HALF.text, ("length",), FIRST_HALF.usage) + assert observed.tokens == (HOP_TOKEN,) and observed.ids == frozenset({first_id}), observed + assert _REVOKED not in observed.raw and CHECKPOINT not in observed.raw, observed.raw + assert [call.target for call in calls] == [target("vertex_ai", stream)], calls diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 75cb2912076..16eb0a3434e 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -499,8 +499,16 @@ def test_cancel_then_delete_on_another_replica_releases_without_a_spend_row(rig: key: Final = scenario.key() created: Final = _create(rig.creator, model, key) _state(rig, created, InteractionState(status="in_progress")) + _drain(rig) cancelled: Final = rig.settler.request("POST", f"/v1beta/interactions/{created}/cancel", {}, key=key) assert cancelled.status_code == 200, cancelled.text + assert JSON_OBJECT.validate_json(cancelled.content)["status"] == "cancelled", cancelled.text + posted: Final = tuple( + string_value(object_value(entry)["path"]) + for entry in _drain(rig) + if string_value(object_value(entry)["method"]) == "POST" + ) + assert len(posted) == 1 and posted[0].endswith(f"/v1beta/interactions/{created}/cancel"), posted before_delete: Final = _settlement(created) assert before_delete is not None and before_delete["claimed"] is False, before_delete deleted: Final = _delete(rig.settler, created, key) diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index c28455f3aca..b84529de0d6 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -414,6 +414,24 @@ def test_provider_specific_cache_key_ignores_litellm_owned_kwargs(monkeypatch: p assert cache.get_cache_key(**{**request, "top_k": 6}) != base_key +def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = { + "model": "gemini/gemini-4-argon", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16000, + } + + first_segment_key: Final = cache.get_cache_key(**request) + resumed_key: Final = cache.get_cache_key(**request, continuation_token="dG9rLTE=") + + assert resumed_key != first_segment_key + assert cache.get_cache_key(**request, continuation_token="dG9rLTI=") not in (first_segment_key, resumed_key) + assert cache.get_cache_key(**request, continuation_token="dG9rLTE=") == resumed_key + assert cache.get_cache_key(**request, continuation_token=None) == first_segment_key + assert cache.get_cache_key(**request, continuation_token="") == first_segment_key + + class PhaseRecordingCache(InMemoryCache): """Records the target and the active span each read / write ran under, as a Redis span would.""" diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 39ca62d3935..21f42ca2dff 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -63,6 +63,7 @@ NATIVE_TOP_LEVEL_FIELD_CASES = [ ("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}), ("cachedContent", "cachedContents/abc123"), ("labels", {"team": "search"}), + ("continuationToken", "dG9rLTE="), ] diff --git a/tests/unit/interactions/test_gemini_interactions_transformation.py b/tests/unit/interactions/test_gemini_interactions_transformation.py index 5e394e9f218..ebc99e8e88b 100644 --- a/tests/unit/interactions/test_gemini_interactions_transformation.py +++ b/tests/unit/interactions/test_gemini_interactions_transformation.py @@ -422,7 +422,7 @@ class TestInteractionOperationUrls: ( "transform_cancel_interaction_request", "interaction-789", - "interaction-789:cancel", + "interaction-789/cancel", ), ], ) @@ -450,7 +450,7 @@ class TestInteractionOperationUrls: assert ( url - == "https://generativelanguage.googleapis.com/v1beta/interactions/..%2F..%2Finteractions%2Fother%3Fx%3D1%23frag:cancel" + == "https://generativelanguage.googleapis.com/v1beta/interactions/..%2F..%2Finteractions%2Fother%3Fx%3D1%23frag/cancel" ) assert params == {} diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index badc2291654..123f44ca779 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -404,6 +404,32 @@ class TestInteractionsAcreateOffline: {"event_type": "step.delta", "object": "interaction", "delta": {"type": "text", "text": "Hello"}}, ] + @pytest.mark.usefixtures("fake_provider_credentials") + @pytest.mark.asyncio + async def test_acreate_gemini_keeps_continuation_token_out_of_the_request(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "interaction-offline", + "object": "interaction", + "model": "gemini-3.8-flash", + "status": "completed", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hi"}]}], + "usage": {"input_tokens": 6, "output_tokens": 3}, + }, + ) + ) + response: Final = await interactions.acreate( + model="gemini/gemini-3.8-flash", + input="Prove it", + continuation_token="dG9rLTE=", + api_key="gemini-offline", + ) + # The Interactions API answered 400 "Unknown parameter 'continuation_token'." on 2026-10-13 + assert json.loads(route.calls.last.request.content) == {"model": "gemini-3.8-flash", "input": "Prove it"} + assert response.status == "completed" + @pytest.mark.usefixtures("fake_provider_credentials") @pytest.mark.asyncio async def test_acreate_simple_litellm_responses_bridge(self, respx_mock: MockRouter) -> None: diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 1f716b41acd..9a7f96000e6 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -246,6 +246,7 @@ class TestMapFinishReasonGemini: ("ESCALATION", "content_filter"), ("UNEXPECTED_TOOL_CALL", "stop"), ("MISSING_THOUGHT_SIGNATURE", "stop"), + ("CONTINUATION", "length"), ], ) def test_gemini_finish_reasons(self, gemini_reason, expected): diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index a646135efc3..84c139e867b 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1932,3 +1932,20 @@ def test_messages_reach_token_count_honours_disable_token_counter(monkeypatch: p messages = _threshold_test_messages(turns=3) assert messages_reach_token_count(model="gpt-4o", messages=messages, threshold=0) is True assert messages_reach_token_count(model="gpt-4o", messages=messages, threshold=1) is False + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_get_modified_max_tokens_leaves_a_gemini_4_long_decode_budget_alone(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "token_counter", lambda **_: 1_000) + + assert ( + get_modified_max_tokens( + model="gemini/gemini-4-argon", + base_model="gemini/gemini-4-argon", + messages=[{"role": "user", "content": "Write a complete proof of the prime number theorem."}], + user_max_tokens=200_000, + buffer_perc=None, + buffer_num=None, + ) + == 200_000 + ) diff --git a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py new file mode 100644 index 00000000000..4465ad7f3df --- /dev/null +++ b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py @@ -0,0 +1,1310 @@ +import json +import sys +import threading +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from itertools import chain, count, repeat +from typing import Final + +import httpx +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.gemini.long_decode import ( + collect_hops_async, + collect_hops_sync, + initial_prompt_tokens_of, + with_initial_prompt_tokens, +) +from litellm.types.llms.vertex_ai import RequestBody +from litellm.types.utils import ChatCompletionDeltaToolCall, ChatCompletionMessageToolCall, Usage + +MODEL: Final = "gemini/gemini-3.8-flash" +MESSAGES: Final = ({"role": "user", "content": "Prove it"},) + +HOP_1: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "hello"}]}, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rLTE=", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 30, "totalTokenCount": 40}, +} +HOP_2: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": " world"}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 10, + "cachedContentTokenCount": 9, + "candidatesTokenCount": 20, + "totalTokenCount": 30, + }, +} +HOP_2_TRUNCATED: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": " world"}]}, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rLTI=", + } + ], + "usageMetadata": HOP_2["usageMetadata"], +} + +TWO_CHOICES_CUT: Final = { + "candidates": [{**HOP_1["candidates"][0], "index": 0}, {**HOP_2_TRUNCATED["candidates"][0], "index": 1}], + "usageMetadata": HOP_1["usageMetadata"], +} +GOOGLE_REJECTION: Final = {"error": {"code": 400, "status": "INVALID_ARGUMENT"}} + +VERTEX_PUBLISHER_BASE: Final = "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/google" + +HOP_1_SEARCHES: Final = ("argon launch date", "argon output limit", "argon pricing") +HOP_2_SEARCHES: Final = ("argon pricing", "argon context window") +STREAM_HOP_1_END: Final = { + "candidates": [{**HOP_1["candidates"][0], "content": {"role": "model", "parts": [{"text": "lo"}]}}], + "usageMetadata": HOP_1["usageMetadata"], +} + + +def _answered_as(body: Mapping[str, object], response_id: str) -> dict[str, object]: + return {**body, "responseId": response_id} + + +def _ending_with(body: Mapping[str, object], finish_reason: str) -> dict[str, object]: + candidates: Final = body["candidates"] + assert isinstance(candidates, list) + return {**body, "candidates": [{**candidates[0], "finishReason": finish_reason}, *candidates[1:]]} + + +def _searching(body: Mapping[str, object], queries: Sequence[str]) -> Mapping[str, object]: + if not queries: + return body + candidates: Final = body["candidates"] + assert isinstance(candidates, list) + grounded: Final = {**candidates[0], "groundingMetadata": {"webSearchQueries": list(queries)}} + return {**body, "candidates": [grounded, *candidates[1:]]} + + +def _replay(hops: Sequence[httpx.Response], sent: list[dict[str, object]]) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)) + return hops[len(sent) - 1] + + return handler + + +def _recording_paths(hops: Sequence[httpx.Response], paths: list[str]) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + return hops[len(paths) - 1] + + return handler + + +def _logging_raw_responses_to(logged: list[object]) -> Callable[[Mapping[str, object]], None]: + def log_raw_response(details: Mapping[str, object]) -> None: + if details.get("log_event_type") == "post_api_call": + logged.append(details["original_response"]) + + return log_raw_response + + +def _sse(*chunks: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + content="".join(f"data: {json.dumps(chunk)}\r\n\r\n" for chunk in chunks).encode(), + headers={"content-type": "text/event-stream"}, + ) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_long_decode_resends_the_original_request_with_the_token_and_bills_every_hop(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + kwargs: Final = {"model": MODEL, "messages": list(MESSAGES), "max_tokens": 100_000, "api_key": "fake-gemini-key"} + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert sent[1] == {**sent[0], "continuationToken": "dG9rLTE="} + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens == 20 + assert response.usage.completion_tokens == 50 + assert response.usage.prompt_tokens_details.cached_tokens == 9 + assert response.usage.prompt_tokens_details.initial_prompt_tokens == 10 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_long_decode_answers_under_the_first_requests_response_id(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + httpx.Response(200, json=_answered_as(HOP_1, "resp-1")), + httpx.Response(200, json=_answered_as(HOP_2, "resp-2")), + ) + kwargs: Final = {"model": MODEL, "messages": list(MESSAGES), "max_tokens": 100_000, "api_key": "fake-gemini-key"} + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert len(sent) == 2 + assert response.id == "resp-1" + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("sent_as_a_string", [{"n": "1"}, {"max_tokens": "100000"}]) +@pytest.mark.asyncio +async def test_request_whose_n_or_max_tokens_is_a_string_is_sent_once_and_never_continued( + stream: bool, sent_as_a_string: Mapping[str, str] +): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END) if stream else httpx.Response(200, json=HOP_1), + httpx.Response(400, json=GOOGLE_REJECTION), + ) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "stream": stream, + "api_key": "fake-gemini-key", + **sent_as_a_string, + } + + response: Final = await litellm.acompletion( + **kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent))) + ) + text: Final = ( + "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + if stream + else response.choices[0].message.content + ) + + assert len(sent) == 1 + assert text == "hello" + + +@pytest.mark.parametrize("hop_2_searches", [(), HOP_2_SEARCHES]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_long_decode_bills_the_searches_of_every_hop(is_async: bool, hop_2_searches: Sequence[str]): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + httpx.Response(200, json=_searching(HOP_1, HOP_1_SEARCHES)), + httpx.Response(200, json=_searching(HOP_2, hop_2_searches)), + ) + kwargs: Final = {"model": MODEL, "messages": list(MESSAGES), "max_tokens": 100_000, "api_key": "fake-gemini-key"} + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert len(sent) == 2 + assert response.choices[0].message.content == "hello world" + assert response.usage.prompt_tokens_details.web_search_requests == len({*HOP_1_SEARCHES, *hop_2_searches}) + assert [item["webSearchQueries"] for item in response.vertex_ai_grounding_metadata] == [ + list(searches) for searches in (HOP_1_SEARCHES, hop_2_searches) if searches + ] + + +HOP_1_GROUNDED: Final = "Kórnik Arboretum reached Level IV." +HOP_1_CITED: Final = " Python is easy to learn." +HOP_2_GROUNDED: Final = " The garden opened in 1826." +URL_CONTEXT: Final = { + "urlMetadata": [ + {"retrievedUrl": "https://www.python.org/about/", "urlRetrievalStatus": "URL_RETRIEVAL_STATUS_SUCCESS"} + ] +} + + +def _byte_span(text: str, part: str) -> tuple[int, int]: + start: Final = text.encode().index(part.encode()) + return start, start + len(part.encode()) + + +def _support(text: str, part: str) -> dict[str, object]: + start, end = _byte_span(text, part) + segment: Final = {"startIndex": start, "endIndex": end, "text": part} if start else {"endIndex": end, "text": part} + return {"segment": segment, "groundingChunkIndices": [0]} + + +def _citation(text: str, part: str) -> dict[str, object]: + start, end = _byte_span(text, part) + return {"startIndex": start, "endIndex": end, "uri": f"https://example.com/{len(part)}"} + + +def _citing( + body: Mapping[str, object], parts: Sequence[Mapping[str, object]], metadata: Mapping[str, object] +) -> dict[str, object]: + candidates: Final = body["candidates"] + assert isinstance(candidates, list) + return {**body, "candidates": [{**candidates[0], "content": {"role": "model", "parts": list(parts)}, **metadata}]} + + +def _hops_citing_sources(citation_key: str, stream: bool = False) -> tuple[httpx.Response, httpx.Response]: + thought: Final = {"text": "Planning the report", "thought": True} + hop_1_text: Final = HOP_1_GROUNDED + HOP_1_CITED + hop_1_metadata: Final = { + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com/arboretum", "title": "arboretum"}}], + "groundingSupports": [_support(hop_1_text, HOP_1_GROUNDED)], + }, + "citationMetadata": {citation_key: [_citation(hop_1_text, HOP_1_CITED.strip())]}, + "urlContextMetadata": URL_CONTEXT, + } + hop_2: Final = _citing( + HOP_2, + [{"text": HOP_2_GROUNDED}], + { + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com/garden", "title": "garden"}}], + "groundingSupports": [_support(HOP_2_GROUNDED, HOP_2_GROUNDED)], + }, + "citationMetadata": {citation_key: [_citation(HOP_2_GROUNDED, "1826")]}, + }, + ) + if stream: + opening: Final = {"candidates": [{"content": {"role": "model", "parts": [thought, {"text": HOP_1_GROUNDED}]}}]} + return _sse(opening, _citing(HOP_1, [{"text": HOP_1_CITED}], hop_1_metadata)), _sse(hop_2) + hop_1: Final = _citing(HOP_1, [thought, {"text": hop_1_text}], hop_1_metadata) + return httpx.Response(200, json=hop_1), httpx.Response(200, json=hop_2) + + +def _supported_spans(grounding: Sequence[Mapping[str, object]]) -> list[Mapping[str, object]]: + return [support["segment"] for support in chain.from_iterable(item["groundingSupports"] for item in grounding)] + + +def _cited_spans(citation_metadata: Sequence[Mapping[str, object]], citation_key: str) -> list[Mapping[str, object]]: + return list(chain.from_iterable(item[citation_key] for item in citation_metadata)) + + +def _bytes_at(answer: str, span: Mapping[str, object]) -> str: + return answer.encode()[span.get("startIndex", 0) : span["endIndex"]].decode() + + +@pytest.mark.parametrize("citation_key", ["citations", "citationSources"]) +@pytest.mark.asyncio +async def test_unary_long_decode_keeps_every_hops_citations_and_points_them_into_the_whole_answer(citation_key: str): + sent: Final[list[dict[str, object]]] = [] + + response: Final = await litellm.acompletion( + model=MODEL, + messages=list(MESSAGES), + max_tokens=100_000, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(_hops_citing_sources(citation_key), sent))), + ) + answer: Final = response.choices[0].message.content + supported: Final = _supported_spans(response.vertex_ai_grounding_metadata) + cited: Final = _cited_spans(response.vertex_ai_citation_metadata, citation_key) + + assert len(sent) == 2 + assert answer == HOP_1_GROUNDED + HOP_1_CITED + HOP_2_GROUNDED + assert [_bytes_at(answer, span) for span in supported] == [span["text"] for span in supported] + assert [span["text"] for span in supported] == [HOP_1_GROUNDED, HOP_2_GROUNDED] + assert [_bytes_at(answer, span) for span in cited] == [HOP_1_CITED.strip(), "1826"] + assert response.vertex_ai_url_context_metadata == [URL_CONTEXT] + + +NO_CITATIONS_OR_URL_CONTEXT: Final = ( + pytest.param({}, id="missing"), + pytest.param({"citationMetadata": None, "urlContextMetadata": None}, id="null"), + pytest.param({"citationMetadata": {}, "urlContextMetadata": {}}, id="empty"), +) + + +def _carrying(body: Mapping[str, object], fields: Mapping[str, object]) -> dict[str, object]: + candidates: Final = body["candidates"] + assert isinstance(candidates, list) + return {**body, "candidates": [{**candidates[0], **fields}, *candidates[1:]]} + + +async def _answer(hops: Sequence[httpx.Response], sent: list[dict[str, object]]): + return await litellm.acompletion( + model=MODEL, + messages=list(MESSAGES), + max_tokens=100_000, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent))), + ) + + +@pytest.mark.parametrize("fields", NO_CITATIONS_OR_URL_CONTEXT) +@pytest.mark.parametrize( + "bodies, answer", [((HOP_2,), " world"), ((HOP_1, HOP_2), "hello world")], ids=["one_request", "continued"] +) +@pytest.mark.asyncio +async def test_unary_answer_whose_candidates_carry_no_citations_or_url_context_is_returned_whole( + bodies: Sequence[Mapping[str, object]], answer: str, fields: Mapping[str, object] +): + sent: Final[list[dict[str, object]]] = [] + + response: Final = await _answer(tuple(httpx.Response(200, json=_carrying(body, fields)) for body in bodies), sent) + + assert len(sent) == len(bodies) + assert response.choices[0].message.content == answer + assert response.choices[0].finish_reason == "stop" + + +@pytest.mark.parametrize("fields", NO_CITATIONS_OR_URL_CONTEXT) +@pytest.mark.asyncio +async def test_unary_long_decode_keeps_earlier_citations_and_url_context_when_a_follow_up_carries_none( + fields: Mapping[str, object], +): + sent: Final[list[dict[str, object]]] = [] + citation: Final = _citation("hello", "hello") + hop_1: Final = _carrying(HOP_1, {"citationMetadata": {"citations": [citation]}, "urlContextMetadata": URL_CONTEXT}) + + response: Final = await _answer( + (httpx.Response(200, json=hop_1), httpx.Response(200, json=_carrying(HOP_2, fields))), sent + ) + + assert response.choices[0].message.content == "hello world" + assert [item for item in response.vertex_ai_url_context_metadata if item != {}] == [URL_CONTEXT] + assert _cited_spans([item for item in response.vertex_ai_citation_metadata if item != {}], "citations") == [ + citation + ] + + +def _hop_2_spanning(span: Mapping[str, object]) -> dict[str, object]: + return _carrying( + HOP_2, + { + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com/world", "title": "world"}}], + "groundingSupports": [{"segment": {**span, "text": " wo"}, "groundingChunkIndices": [0]}], + }, + "citationMetadata": {"citations": [{**span, "uri": "https://example.com/world"}]}, + }, + ) + + +def _chunk_metadata(chunks: Sequence[object], field: str) -> list[Mapping[str, object]]: + return list(chain.from_iterable(getattr(chunk, field, None) or () for chunk in chunks)) + + +async def _answer_with_spans( + stream: bool, hop_2: Mapping[str, object], sent: list[dict[str, object]] +) -> tuple[str, list[Mapping[str, object]], list[Mapping[str, object]]]: + if stream: + chunks: Final = await _collect_stream(True, (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(hop_2)), sent) + return ( + _stream_text(chunks), + _supported_spans(_chunk_metadata(chunks, "vertex_ai_grounding_metadata")), + _cited_spans(_chunk_metadata(chunks, "vertex_ai_citation_metadata"), "citations"), + ) + response: Final = await _answer((httpx.Response(200, json=HOP_1), httpx.Response(200, json=hop_2)), sent) + return ( + response.choices[0].message.content, + _supported_spans(response.vertex_ai_grounding_metadata), + _cited_spans(response.vertex_ai_citation_metadata, "citations"), + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_long_decode_counts_a_follow_ups_null_start_index_from_the_start_of_its_text(stream: bool): + sent: Final[list[dict[str, object]]] = [] + + answer, supported, cited = await _answer_with_spans( + stream, _hop_2_spanning({"startIndex": None, "endIndex": len(" wo")}), sent + ) + + assert answer == "hello world" + assert [_bytes_at(answer, s) for s in supported] == [" wo"] + assert [_bytes_at(answer, s) for s in cited] == [" wo"] + + +@pytest.mark.parametrize( + "span", + [ + pytest.param({}, id="no_offsets"), + pytest.param({"startIndex": 1}, id="missing_end"), + pytest.param({"startIndex": 1, "endIndex": None}, id="null_end"), + ], +) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_long_decode_returns_a_follow_ups_span_with_no_end_index_as_google_sent_it( + stream: bool, span: Mapping[str, object] +): + sent: Final[list[dict[str, object]]] = [] + + answer, supported, cited = await _answer_with_spans(stream, _hop_2_spanning(span), sent) + + assert answer == "hello world" + assert supported == [{**span, "text": " wo"}] + assert cited == [{**span, "uri": "https://example.com/world"}] + + +@pytest.mark.parametrize( + "max_tokens, expected_requests, expected_token", + [(None, 1, "dG9rLTE="), (30, 1, "dG9rLTE="), (50, 2, "dG9rLTI=")], +) +@pytest.mark.asyncio +async def test_unary_long_decode_stops_at_the_callers_cap_and_returns_the_token_to_resume( + max_tokens: int | None, expected_requests: int, expected_token: str +): + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2_TRUNCATED)) + + response: Final = await litellm.acompletion( + model=MODEL, + messages=list(MESSAGES), + max_tokens=max_tokens, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent))), + ) + + assert len(sent) == expected_requests + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == expected_token + assert response.usage.prompt_tokens_details.initial_prompt_tokens == (10 if expected_requests > 1 else None) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_answer_cut_at_the_callers_cap_is_returned_with_its_token_in_one_request(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + # The Gemini API answered 400 "max_output_tokens (16000) has already been exhausted" to a resend + # after MAX_TOKENS on 2026-10-13 + hops: Final = ( + httpx.Response(200, json=_ending_with(HOP_1, "MAX_TOKENS")), + httpx.Response(400, json=GOOGLE_REJECTION), + ) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "continuation_token": "cHJldg==", + "api_key": "fake-gemini-key", + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert len(sent) == 1 + assert response.choices[0].message.content == "hello" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTE=" + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_request_for_more_than_one_choice_is_never_continued(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + # Vertex AI answered 400 "Candidate count greater than 1 is not supported when continuation_token is set" + # on 2026-10-13 + hops: Final = (httpx.Response(200, json=TWO_CHOICES_CUT), httpx.Response(400, json=GOOGLE_REJECTION)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "n": 2, + "api_key": "fake-gemini-key", + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert len(sent) == 1 + assert [choice.message.content for choice in response.choices] == ["hello", " world"] + assert [choice.finish_reason for choice in response.choices] == ["length", "length"] + assert [choice.message.provider_specific_fields["continuation_token"] for choice in response.choices] == [ + "dG9rLTE=", + "dG9rLTI=", + ] + + +@pytest.mark.asyncio +async def test_unary_failed_continuation_hop_returns_the_answer_so_far_with_its_token(): + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(429, json={"error": {"code": 429}})) + + response: Final = await litellm.acompletion( + model=MODEL, + messages=list(MESSAGES), + max_tokens=100_000, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent))), + ) + + assert len(sent) == 2 + assert response.choices[0].message.content == "hello" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTE=" + assert response.usage.completion_tokens == 30 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_long_decode_logs_one_raw_response_object_with_the_whole_answer_and_only_its_last_token( + is_async: bool, +): + sent: Final[list[dict[str, object]]] = [] + logged: Final[list[object]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2_TRUNCATED)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 50, + "api_key": "fake-gemini-key", + "logger_fn": _logging_raw_responses_to(logged), + } + + if is_async: + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + else: + litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + + assert len(sent) == 2 + assert len(logged) == 1 + raw: Final = logged[0] + assert isinstance(raw, str) + body: Final = json.loads(raw) + assert isinstance(body, dict) + candidate: Final = body["candidates"][0] + assert "".join(part["text"] for part in candidate["content"]["parts"]) == "hello world" + assert candidate["continuationToken"] == "dG9rLTI=" + assert "dG9rLTE=" not in raw + assert body["usageMetadata"]["candidatesTokenCount"] == 50 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_answer_in_one_request_logs_googles_body_unchanged(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + logged: Final[list[object]] = [] + googles_body: Final = json.dumps(HOP_2, indent=2) + hops: Final = (httpx.Response(200, content=googles_body.encode(), headers={"content-type": "application/json"}),) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 50, + "api_key": "fake-gemini-key", + "logger_fn": _logging_raw_responses_to(logged), + } + + if is_async: + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + else: + litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + + assert len(sent) == 1 + assert logged == [googles_body] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_long_decode_follows_more_hops_than_the_interpreter_allows_nested_calls(is_async: bool): + hop_count: Final = sys.getrecursionlimit() + 50 + one_token_hop: Final = {**HOP_1, "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1}} + responses: Final = chain( + repeat(httpx.Response(200, json=one_token_hop), hop_count - 2), + (httpx.Response(200, json=HOP_2),), + ) + request_body: Final = RequestBody(contents=[], generationConfig={"max_output_tokens": hop_count * 100}) + + async def post_async(_body: RequestBody) -> httpx.Response: + return next(responses) + + hops: Final = ( + await collect_hops_async(one_token_hop, request_body, post_async) + if is_async + else collect_hops_sync(one_token_hop, request_body, lambda _body: next(responses)) + ) + + assert len(hops) == hop_count + assert hops[-1] == HOP_2 + + +@pytest.mark.asyncio +async def test_continuation_token_param_is_sent_at_the_top_level_of_generate_content(): + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=MODEL, + messages=list(MESSAGES), + continuation_token="dG9rLTE=", + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert sent[0]["continuationToken"] == "dG9rLTE=" + assert "continuation_token" not in json.dumps(sent[0]) + + +STREAM_FIRST_CHUNK: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}}]} + + +def _stream_hops() -> tuple[httpx.Response, httpx.Response]: + return (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2)) + + +async def _stream_chunks( + is_async: bool, + handler: Callable[[httpx.Request], httpx.Response], + api_base: str | None = None, + continuation_token: str | None = None, + n: int | None = None, + timeout: float | None = None, +) -> list: + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "stream": True, + "stream_options": {"include_usage": True}, + "api_key": "fake-gemini-key", + "api_base": api_base, + "continuation_token": continuation_token, + "n": n, + "timeout": timeout, + } + if is_async: + stream: Final = await litellm.acompletion( + **kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + ) + return [chunk async for chunk in stream] + return list( + litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))) + ) + + +async def _collect_stream(is_async: bool, hops: Sequence[httpx.Response], sent: list[dict[str, object]]) -> list: + return await _stream_chunks(is_async, _replay(hops, sent)) + + +def _stream_text(chunks: list) -> str: + return "".join(c.choices[0].delta.content or "" for c in chunks if c.choices) + + +def _stream_finish_reasons(chunks: list) -> list[str]: + return [c.choices[0].finish_reason for c in chunks if c.choices and c.choices[0].finish_reason] + + +def _delta_fields(chunks: list) -> list[Mapping[str, object]]: + return [ + choice.delta.provider_specific_fields + for chunk in chunks + for choice in chunk.choices # comprehension-ok: flattening stream chunks into their choices + if choice.delta.provider_specific_fields + ] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_long_decode_chains_hops_into_one_stream_with_one_finish_and_summed_usage(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + + chunks: Final = await _collect_stream(is_async, _stream_hops(), sent) + + assert sent[1] == {**sent[0], "continuationToken": "dG9rLTE="} + assert "".join(c.choices[0].delta.content or "" for c in chunks if c.choices) == "hello world" + assert [c.choices[0].finish_reason for c in chunks if c.choices and c.choices[0].finish_reason] == ["stop"] + assert not any("continuation_token" in fields for fields in _delta_fields(chunks)) + assert chunks[-1].usage.prompt_tokens == 20 + assert chunks[-1].usage.completion_tokens == 50 + assert chunks[-1].usage.prompt_tokens_details.initial_prompt_tokens == 10 + + +TOOL_CALL_PART: Final = { + "functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, + "thoughtSignature": "c2lnLTI=", +} +HOP_2_CALLING_A_TOOL: Final = { + "candidates": [{"content": {"role": "model", "parts": [TOOL_CALL_PART]}, "finishReason": "STOP"}], + "usageMetadata": HOP_2["usageMetadata"], +} + + +@dataclass(frozen=True, slots=True) +class _ToolAnswer: + requests: int + text: str | None + tool_calls: tuple[Mapping[str, object], ...] + finish_reasons: tuple[str | None, ...] + + +def _without_id(tool_call: ChatCompletionMessageToolCall | ChatCompletionDeltaToolCall) -> Mapping[str, object]: + return tool_call.model_dump(exclude={"id"}) + + +async def _tool_answer(stream: bool, hops: Sequence[httpx.Response]) -> _ToolAnswer: + sent: Final[list[dict[str, object]]] = [] + if not stream: + choice: Final = (await _answer(hops, sent)).choices[0] + return _ToolAnswer( + len(sent), + choice.message.content, + tuple(map(_without_id, choice.message.tool_calls or ())), + (choice.finish_reason,), + ) + chunks: Final = await _collect_stream(True, hops, sent) + deltas: Final = chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks if chunk.choices) + return _ToolAnswer( + len(sent), _stream_text(chunks), tuple(map(_without_id, deltas)), tuple(_stream_finish_reasons(chunks)) + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_long_decode_returns_a_follow_ups_tool_call_with_its_signature_as_one_request_would(stream: bool): + one_request_parts: Final = {"role": "model", "parts": [{"text": "lo" if stream else "hello"}, TOOL_CALL_PART]} + one_request_end: Final = { + **HOP_2_CALLING_A_TOOL, + "candidates": [{**HOP_2_CALLING_A_TOOL["candidates"][0], "content": one_request_parts}], + } + one_request: Final = ( + (_sse(STREAM_FIRST_CHUNK, one_request_end),) if stream else (httpx.Response(200, json=one_request_end),) + ) + continued: Final = ( + (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2_CALLING_A_TOOL)) + if stream + else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2_CALLING_A_TOOL)) + ) + + answer: Final = await _tool_answer(stream, continued) + + assert answer.requests == 2 + assert [call["provider_specific_fields"] for call in answer.tool_calls] == [{"thought_signature": "c2lnLTI="}] + assert (answer.text, answer.tool_calls, answer.finish_reasons) == ( + "hello", + (await _tool_answer(stream, one_request)).tool_calls, + ("tool_calls",), + ) + + +@pytest.mark.parametrize("fields", NO_CITATIONS_OR_URL_CONTEXT) +@pytest.mark.parametrize( + "hop_chunks, answer", + [(((HOP_2,),), " world"), (((STREAM_FIRST_CHUNK, STREAM_HOP_1_END), (HOP_2,)), "hello world")], + ids=["one_request", "continued"], +) +@pytest.mark.asyncio +async def test_streamed_answer_whose_chunks_carry_no_citations_or_url_context_is_returned_whole( + hop_chunks: Sequence[Sequence[Mapping[str, object]]], answer: str, fields: Mapping[str, object] +): + sent: Final[list[dict[str, object]]] = [] + hops: Final = tuple(_sse(*(_carrying(chunk, fields) for chunk in hop)) for hop in hop_chunks) + + chunks: Final = await _collect_stream(True, hops, sent) + + assert len(sent) == len(hop_chunks) + assert _stream_text(chunks) == answer + assert _stream_finish_reasons(chunks) == ["stop"] + + +@pytest.mark.parametrize("citation_key", ["citations", "citationSources"]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_long_decode_points_every_hops_citations_into_the_whole_answer( + is_async: bool, citation_key: str +): + sent: Final[list[dict[str, object]]] = [] + + chunks: Final = await _collect_stream(is_async, _hops_citing_sources(citation_key, stream=True), sent) + answer: Final = _stream_text(chunks) + supported: Final = _supported_spans(_chunk_metadata(chunks, "vertex_ai_grounding_metadata")) + cited: Final = _cited_spans(_chunk_metadata(chunks, "vertex_ai_citation_metadata"), citation_key) + + assert len(sent) == 2 + assert answer == HOP_1_GROUNDED + HOP_1_CITED + HOP_2_GROUNDED + assert [span["text"] for span in supported] == [HOP_1_GROUNDED, HOP_2_GROUNDED] + assert [_bytes_at(answer, span) for span in supported] == [HOP_1_GROUNDED, HOP_2_GROUNDED] + assert [_bytes_at(answer, span) for span in cited] == [HOP_1_CITED.strip(), "1826"] + + +@pytest.mark.parametrize("next_request_fails", [False, True]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_long_decode_keeps_the_first_requests_response_id_on_every_chunk( + is_async: bool, next_request_fails: bool +): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(_answered_as(STREAM_FIRST_CHUNK, "resp-1"), _answered_as(STREAM_HOP_1_END, "resp-1")), + httpx.Response(500, json={"error": {"code": 500}}) + if next_request_fails + else _sse(_answered_as(HOP_2, "resp-2")), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 2 + assert _stream_finish_reasons(chunks) == ["length" if next_request_fails else "stop"] + assert {chunk.id for chunk in chunks} == {"resp-1"} + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_answer_cut_at_the_callers_cap_ends_with_length_and_its_token_in_one_request(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(STREAM_FIRST_CHUNK, _ending_with(STREAM_HOP_1_END, "MAX_TOKENS")), + httpx.Response(400, json=GOOGLE_REJECTION), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 1 + assert _stream_text(chunks) == "hello" + assert _stream_finish_reasons(chunks) == ["length"] + assert {"continuation_token": "dG9rLTE="} in _delta_fields(chunks) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_request_for_more_than_one_choice_is_never_continued(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = (_sse(TWO_CHOICES_CUT), httpx.Response(400, json=GOOGLE_REJECTION)) + + chunks: Final = await _stream_chunks(is_async, _replay(hops, sent), n=2) + + assert len(sent) == 1 + assert _stream_finish_reasons(chunks) == ["length"] + assert {"continuation_token": "dG9rLTE="} in _delta_fields(chunks) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_long_decode_reads_every_chunk_of_each_continuation_stream(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _stream_hops()[0], + _sse( + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": " wor"}]}, "continuationToken": "Y2twdC0x"} + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + }, + { + **HOP_2_TRUNCATED, + "candidates": [{**HOP_2_TRUNCATED["candidates"][0], "content": {"parts": [{"text": "ld"}]}}], + }, + ), + _sse( + {"candidates": [{"content": {"role": "model", "parts": [{"text": " ag"}]}}]}, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "ain"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 7, "totalTokenCount": 17}, + }, + ), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert [body.get("continuationToken") for body in sent] == [None, "dG9rLTE=", "dG9rLTI="] + assert _stream_text(chunks) == "hello world again" + assert _stream_finish_reasons(chunks) == ["stop"] + assert chunks[-1].usage.prompt_tokens == 30 + assert chunks[-1].usage.completion_tokens == 57 + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_continuation_stream_is_read_whole_after_a_stream_that_ended_on_an_unfinished_frame( + is_async: bool, +): + sent: Final[list[dict[str, object]]] = [] + first_stream, next_stream = _stream_hops() + cut_off_first_stream: Final = httpx.Response( + 200, + content=first_stream.content + b'data: {"candidates": [{"content": {"parts": [{"text": "lo\r\n\r\n', + headers=first_stream.headers, + ) + + chunks: Final = await _collect_stream(is_async, (cut_off_first_stream, next_stream), sent) + + assert len(sent) == 2 + assert _stream_text(chunks) == "hello world" + assert _stream_finish_reasons(chunks) == ["stop"] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_initial_prompt_tokens_is_the_prompt_count_the_first_stream_ends_with(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + first_stream: Final = _sse( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}}], + "usageMetadata": {"promptTokenCount": 7, "candidatesTokenCount": 3, "totalTokenCount": 10}, + }, + STREAM_HOP_1_END, + ) + + chunks: Final = await _collect_stream(is_async, (first_stream, _sse(HOP_2)), sent) + + assert chunks[-1].usage.prompt_tokens_details.initial_prompt_tokens == 10 + assert chunks[-1].usage.prompt_tokens == 20 + + +def _api_versions(paths: Sequence[str]) -> list[str]: + return [path.split("/projects/")[0] for path in paths] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_continuation_request_to_a_vertex_url_goes_to_the_api_version_that_accepts_the_token( + is_async: bool, +): + paths: Final[list[str]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + transport: Final = httpx.MockTransport(_recording_paths(hops, paths)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "api_key": "fake-gemini-key", + "api_base": VERTEX_PUBLISHER_BASE, + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + + assert response.choices[0].message.content == "hello world" + # Vertex AI answered 400 for continuationToken on v1 and 200 on v1beta1, probed at location global on 2026-10-13 + assert _api_versions(paths) == ["/v1", "/v1beta1"] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_continuation_request_to_a_vertex_url_goes_to_the_api_version_that_accepts_the_token( + is_async: bool, +): + paths: Final[list[str]] = [] + + chunks: Final = await _stream_chunks(is_async, _recording_paths(_stream_hops(), paths), VERTEX_PUBLISHER_BASE) + + assert _stream_text(chunks) == "hello world" + assert _api_versions(paths) == ["/v1", "/v1beta1"] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_request_the_client_continues_on_a_vertex_url_goes_to_the_api_version_that_accepts_the_token( + is_async: bool, +): + paths: Final[list[str]] = [] + + chunks: Final = await _stream_chunks( + is_async, _recording_paths((_sse(HOP_2),), paths), VERTEX_PUBLISHER_BASE, continuation_token="dG9rLTE=" + ) + + assert _stream_text(chunks) == " world" + assert _api_versions(paths) == ["/v1beta1"] + + +@pytest.mark.parametrize("hop_2_searches", [(), HOP_2_SEARCHES]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_long_decode_bills_the_searches_of_every_hop(is_async: bool, hop_2_searches: Sequence[str]): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(_searching(STREAM_HOP_1_END, HOP_1_SEARCHES)), + _sse(_searching(HOP_2, hop_2_searches)), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 2 + assert chunks[-1].usage.prompt_tokens_details.web_search_requests == len({*HOP_1_SEARCHES, *hop_2_searches}) + + +def _checkpoint(text: str, token: str) -> dict[str, object]: + return {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}, "continuationToken": token}]} + + +def _client_tokens(chunks: list) -> list[object]: + return [fields["continuation_token"] for fields in _delta_fields(chunks) if "continuation_token" in fields] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_mid_stream_checkpoint_tokens_never_reach_the_client(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(_checkpoint("hel", "Y2twdC0x"), STREAM_HOP_1_END), + _sse(_checkpoint(" big", "Y2twdC0y"), _checkpoint(" wide", "Y2twdC0z"), HOP_2), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 2 + assert _stream_text(chunks) == "hello big wide world" + assert _stream_finish_reasons(chunks) == ["stop"] + assert _client_tokens(chunks) == [] + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_cut_off_stream_gives_the_client_only_its_final_token(is_async: bool): + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + _sse(_checkpoint("hel", "Y2twdC0x"), _ending_with(STREAM_HOP_1_END, "MAX_TOKENS")), + httpx.Response(400, json=GOOGLE_REJECTION), + ) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 1 + assert _stream_text(chunks) == "hello" + assert _stream_finish_reasons(chunks) == ["length"] + assert _client_tokens(chunks) == ["dG9rLTE="] + + +class _BodyThatDropsBeforeItsFirstFrame(httpx.SyncByteStream, httpx.AsyncByteStream): + def __iter__(self) -> Iterator[bytes]: + raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") + + def __aiter__(self) -> AsyncIterator[bytes]: + raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") + + +@pytest.mark.parametrize( + "failed_hop", + [ + lambda: httpx.Response(500, json={"error": {"code": 500}}), + lambda: httpx.Response(200, content=b"not json"), + lambda: _sse({"error": {"code": 503, "message": "The model is overloaded.", "status": "UNAVAILABLE"}}, HOP_2), + lambda: httpx.Response(200, stream=_BodyThatDropsBeforeItsFirstFrame()), + ], + ids=["refused", "not_json", "error_frame", "dropped"], +) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_failed_continuation_hop_ends_the_stream_with_length_and_the_token( + is_async: bool, failed_hop: Callable[[], httpx.Response] +): + sent: Final[list[dict[str, object]]] = [] + hops: Final = (_stream_hops()[0], failed_hop()) + + chunks: Final = await _collect_stream(is_async, hops, sent) + + assert len(sent) == 2 + assert "".join(c.choices[0].delta.content or "" for c in chunks if c.choices) == "hello" + assert [c.choices[0].finish_reason for c in chunks if c.choices and c.choices[0].finish_reason] == ["length"] + assert {"continuation_token": "dG9rLTE="} in _delta_fields(chunks) + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_streaming_continuation_request_is_sent_with_the_callers_timeout(is_async: bool): + read_timeouts: Final[list[object]] = [] + hops: Final = _stream_hops() + + def handler(request: httpx.Request) -> httpx.Response: + read_timeouts.append(request.extensions["timeout"]["read"]) + return hops[len(read_timeouts) - 1] + + chunks: Final = await _stream_chunks(is_async, handler, timeout=7.5) + + assert _stream_text(chunks) == "hello world" + assert read_timeouts == [7.5, 7.5] + + +@dataclass(frozen=True, slots=True) +class _TokenEndpoint: + credentials: str + down: threading.Event + + +def _service_account_json(token_uri: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": "p", + "private_key_id": "k", + "private_key": private_key, + "client_email": "sa@p.iam.gserviceaccount.com", + "token_uri": token_uri, + } + ) + + +@pytest.fixture +def token_endpoint() -> Iterator[_TokenEndpoint]: + down: Final = threading.Event() + minted: Final = count(1) + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers["Content-Length"])) + refused: Final = down.is_set() + payload: Final = json.dumps( + {"error": "invalid_grant"} if refused else {"access_token": f"token-{next(minted)}", "expires_in": 0} + ).encode() + self.send_response(400 if refused else 200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *_args: object) -> None: + return None + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + yield _TokenEndpoint(_service_account_json(f"http://127.0.0.1:{server.server_port}/token"), down) + server.shutdown() + server.server_close() + + +def _google_taking_each_access_token_once( + hops: Sequence[httpx.Response], used: list[str], on_answer: Callable[[], object] +) -> Callable[[httpx.Request], httpx.Response]: + def handler(request: httpx.Request) -> httpx.Response: + token: Final = request.headers["authorization"] + if token in used: + return httpx.Response(401, json={"error": {"code": 401, "status": "UNAUTHENTICATED"}}) + used.append(token) + on_answer() + return hops[len(used) - 1] + + return handler + + +async def _vertex_answer( + is_async: bool, + stream: bool, + handler: Callable[[httpx.Request], httpx.Response], + credentials: str, + extra_headers: Mapping[str, str] | None = None, +) -> tuple[str, list[str], list[object]]: + kwargs: Final = { + "model": "vertex_ai/gemini-3.8-flash", + "extra_headers": extra_headers, + "messages": list(MESSAGES), + "max_tokens": 100_000, + "stream": stream, + "vertex_project": "p", + "vertex_location": "global", + "vertex_credentials": credentials, + } + transport: Final = httpx.MockTransport(handler) + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + if stream: + chunks: Final = [chunk async for chunk in response] if is_async else list(response) + return _stream_text(chunks), _stream_finish_reasons(chunks), _client_tokens(chunks) + choice: Final = response.choices[0] + fields: Final = choice.message.provider_specific_fields or {} + return choice.message.content, [choice.finish_reason], [fields["continuation_token"]] if fields else [] + + +def _two_hops(stream: bool) -> tuple[httpx.Response, httpx.Response]: + return _stream_hops() if stream else (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_finishes_when_each_continuation_outlives_the_access_token_before_it( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + used: Final[list[str]] = [] + google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, lambda: None) + + answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) + + assert answer == ("hello world", ["stop"], []) + assert len(used) == 2 + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_returns_the_answer_so_far_with_its_token_when_no_new_access_token_can_be_had( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + used: Final[list[str]] = [] + google: Final = _google_taking_each_access_token_once(_two_hops(stream), used, token_endpoint.down.set) + + answer: Final = await _vertex_answer(is_async, stream, google, token_endpoint.credentials) + + assert answer == ("hello", ["length"], ["dG9rLTE="]) + assert len(used) == 1 + + +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_every_continuation( + is_async: bool, stream: bool, token_endpoint: _TokenEndpoint +): + hops: Final = _two_hops(stream) + seen: Final[list[str]] = [] + + def google(request: httpx.Request) -> httpx.Response: + seen.append(request.headers["authorization"]) + return hops[len(seen) - 1] + + answer: Final = await _vertex_answer( + is_async, stream, google, token_endpoint.credentials, {"Authorization": "Bearer mine"} + ) + + assert answer == ("hello world", ["stop"], []) + assert seen == ["Bearer mine", "Bearer mine"] + + +def test_initial_prompt_tokens_is_absent_on_a_single_hop_and_the_first_hops_count_otherwise(): + assert initial_prompt_tokens_of((HOP_1,)) is None + assert initial_prompt_tokens_of((HOP_1, HOP_2)) == 10 + + +def test_with_initial_prompt_tokens_adds_the_details_block_and_keeps_the_counts(): + usage: Final = Usage(prompt_tokens=20, completion_tokens=50, total_tokens=70) + + marked: Final = with_initial_prompt_tokens(usage, 10) + + assert marked.prompt_tokens_details.initial_prompt_tokens == 10 + assert (marked.prompt_tokens, marked.completion_tokens, marked.total_tokens) == (20, 50, 70) + assert usage.prompt_tokens_details is None diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3f49869b19f..40ede74b98b 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5777,28 +5777,20 @@ def test_sync_streaming_uses_custom_client(): def test_transform_response_does_not_leak_body_on_parse_failure(): - leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]} - raw_response = MagicMock() - raw_response.json.return_value = leaky_body - raw_response.text = json.dumps(leaky_body) - raw_response.headers = {} + raw_response = httpx.Response(200, json=[{"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]}]) - with patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody", - side_effect=KeyError("missing required field"), - ): - with pytest.raises(VertexAIError) as exc_info: - VertexGeminiConfig().transform_response( - model="gemini-pro", - raw_response=raw_response, - model_response=ModelResponse(), - logging_obj=MagicMock(), - request_data={}, - messages=[], - optional_params={}, - litellm_params={}, - encoding=None, - ) + with pytest.raises(VertexAIError) as exc_info: + VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) msg = str(exc_info.value) assert "secret content" not in msg @@ -9459,6 +9451,8 @@ def test_vertex_ai_streaming_response_id(): model="gemini-pro", messages=[], logging_obj=logging_obj, + request_body={}, + hop_headers=lambda: {}, ) iterator = iter(iterator) first_chunk = next(iterator) diff --git a/tests/unit/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py index 3d6c5b538e9..34502d5cebe 100644 --- a/tests/unit/llms/vertex_ai/test_http_status_201.py +++ b/tests/unit/llms/vertex_ai/test_http_status_201.py @@ -28,6 +28,9 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): self.mock_logging_obj = MagicMock() self.mock_logging_obj.post_call = MagicMock() + async def _hop_headers(self): + return self.mock_headers + @patch( "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.get_async_httpx_client" ) @@ -51,6 +54,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): api_base="https://mock-vertex-ai-api.com", headers=self.mock_headers, data=self.mock_data, + request_body=json.loads(self.mock_data), + hop_headers=self._hop_headers, model=self.mock_model, messages=self.messages, logging_obj=self.mock_logging_obj, @@ -82,6 +87,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): api_base="https://mock-vertex-ai-api.com", headers=self.mock_headers, data=self.mock_data, + request_body=json.loads(self.mock_data), + hop_headers=lambda: self.mock_headers, model=self.mock_model, messages=self.messages, logging_obj=self.mock_logging_obj, @@ -111,6 +118,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): api_base="https://mock-vertex-ai-api.com", headers=self.mock_headers, data=self.mock_data, + request_body=json.loads(self.mock_data), + hop_headers=lambda: self.mock_headers, model=self.mock_model, messages=self.messages, logging_obj=self.mock_logging_obj, diff --git a/tests/unit/test_gemini_4_argon_model_metadata.py b/tests/unit/test_gemini_4_argon_model_metadata.py new file mode 100644 index 00000000000..d4c3a69ca20 --- /dev/null +++ b/tests/unit/test_gemini_4_argon_model_metadata.py @@ -0,0 +1,42 @@ +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm import get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + +REPO_ROOT: Final = Path(__file__).parents[2] +MODELS: Final = ( + ("gemini/gemini-4-argon", "gemini-4-argon", "gemini"), + ("vertex_ai/gemini-4-argon", "gemini-4-argon", "vertex_ai"), +) +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) + + +def _cost_map_entry(path: Path, model: str) -> dict[str, object]: + return COST_MAP_ADAPTER.validate_json(path.read_bytes())[model] + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize(("model", "expected_model", "expected_provider"), MODELS) +def test_gemini_4_argon_routes_to_its_provider_on_the_thinking_level_path( + model: str, expected_model: str, expected_provider: str +) -> None: + routed_model, provider, _, _ = get_llm_provider(model=model) + assert (routed_model, provider) == (expected_model, expected_provider) + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["litellm_provider"] == expected_provider + assert info["max_output_tokens"] == info["max_input_tokens"] + assert VertexGeminiConfig.is_gemini_3_or_newer(routed_model) is True + + +@pytest.mark.parametrize("model", [row[0] for row in MODELS]) +def test_gemini_4_argon_entry_source_and_backup_match(model: str) -> None: + main_entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json", model) + backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json", model) + + assert backup_entry == main_entry diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 321f54466d9..6d94fc976fe 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7992,7 +7992,7 @@ export interface paths { * Cancel Interaction * @description Cancel an interaction by ID. * - * Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel + * Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}/cancel */ post: operations["cancel_interaction_interactions__interaction_id__cancel_post"]; delete?: never; @@ -23422,7 +23422,7 @@ export interface paths { * Cancel Interaction * @description Cancel an interaction by ID. * - * Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel + * Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}/cancel */ post: operations["cancel_interaction_v1beta_interactions__interaction_id__cancel_post"]; delete?: never;