mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
feat(gemini)!: auto-continue gemini-4-argon long decodes with continuationToken (#46240)
* feat(gemini): continue MAX_TOKENS decodes with continuationToken across chat, streaming, and interactions * fix(gemini): continue long decode on CONTINUATION and cancel interactions at /cancel * feat(gemini): add public gemini-4-argon cost-map rows for AI Studio and Vertex AI Add gemini/gemini-4-argon and vertex_ai/gemini-4-argon alongside the existing early-access row, carrying the same limits and capabilities. The AI Studio row mirrors the early-access entry; the Vertex row differs only in provider and pricing source. Prices follow once Google publishes them. Extend the mapped metadata test to cover all three ids, that each routes to its provider, that its backup copy matches, and that the public rows carry the same limits as the early-access row. * feat(gemini): raise gemini-4-argon to the 1M long decode cap and report initial_prompt_tokens across hops * chore(prices): build on the merged gemini-4-argon rows and only raise the long decode cap * fix(gemini): follow continuation hops in a loop instead of recursion The unary hop collectors and the stream iterator called themselves once per continuation request, which the recursion check in CI rejects and which would hit the interpreter's call depth on a very long output. They now loop, and a test follows more hops than the interpreter allows nested calls * fix(gemini): bill Google Search grounding from every continuation request Google sends search grounding once, on the request that ran the search, and does not repeat it on a continuation request. The merged response kept only the last request's grounding, so usage.web_search_requests and the per-search fee dropped every earlier search. Unary merges and the stream iterator now carry the grounding of every request in the chain * fix(gemini): continue long gemini-4-argon answers on streams and Vertex AI A continued stream stopped after the first chunk of its second request. The stream wrapper asks the iterator for a fresh start before every chunk, and that reset it to the first request's finished stream. The iterator now reads from the newest request's stream On Vertex AI the continuation request went to the v1 API version, which rejects continuationToken with a 400. A request that carries the token now goes to v1beta1, which accepts it, including a token that v1 issued On a stream, initial_prompt_tokens came from a usage chunk in the middle of the stream, which reports a smaller prompt count than the last chunk. It is now the count the first request's stream ends with, the same number a unary call reports continuation_token is no longer forwarded on the Interactions API, which answers 400 for the field. It is dropped there as it was before this branch * chore(gemini): name the type errors the dual-mode stream iterator trips The stream iterator holds a sync or an async line iterator in one untyped attribute. Pointing it at a continuation request's stream gave the type checker a concrete union for it, so six existing lines that only ever run in one of the two modes started to report errors This adds a rule-named suppression with a reason on each of those lines and changes no code. Splitting the class into a sync and an async iterator would remove them and is left for a follow-up * fix(gemini): continue only when Google ends a request with CONTINUATION, and stop streaming mid-answer tokens A long decode now sends the next request only when Google's finish reason is CONTINUATION, so an answer cut at the caller's own max_tokens returns in one request with its continuation_token instead of drawing a rejected extra request A request for more than one choice is never continued, since Google rejects a continuation with more than one candidate Chat-format streams no longer carry the tokens Google sends mid-answer. Only a stream that ends cut off carries a continuation_token, on its last content chunk The response cache key now includes continuation_token, so a resend with a token is not served the earlier segment * fix(gemini): keep one response id across a continued answer A continued chat stream changed its chunk id each time the next request to Google started, and a unary answer came back under the last request's id. Every chunk and the unary answer now carry the first request's id, including the chunk that ends a cut-off stream A request whose n or max_tokens arrives as a string is sent once and never continued, where it used to fail with a 500 on the comparison. An empty continuation_token no longer changes the cache key * test(gemini): add integration cells for long decode continuation Covers continued answers on chat, /v1/messages and /v1/responses, streamed and not, on the Gemini API and Vertex AI against a scripted Google, plus concurrency, outage and worker restart runs Also moves the Google stand-in cancel route to /cancel, the URL this PR now calls for Interactions cancel, and checks the cancel test posts there exactly once * fix(vertex_ai): fetch the access token again before each continuation request A long Gemini 4 answer on Vertex AI can run for over half an hour of follow-up requests, and every one of them reused the bearer token fetched before the first request. Once that token expired Google answered 401 and the caller got a cut-off answer with finish_reason length Each Google request in a chain now asks for the token again, which is a cache hit until the token nears expiry, on the streamed and unary paths, sync and async. A refresh that fails counts as a failed follow-up, so the caller still gets the answer so far with its continuation token. A caller's own Authorization in extra_headers still wins on every request, and the Gemini API key path is unchanged * fix(vertex_ai): start each continuation stream with an empty parser buffer When a Google stream ended on an unfinished frame, its leftover text stayed in the stream parser and every frame of the next continuation stream was appended to it. None of them could be decoded, so the whole continuation was billed and dropped without an error Each continuation stream now starts with a fresh buffer, and a test covers a stream that ends on an unfinished frame, sync and async * fix(vertex_ai): treat a continuation stream with no readable frame as a failed follow-up A continuation stream that answered 200 and then sent Google's error frame first, sent a body that is not JSON, or dropped before its first frame ended the caller's stream with `stop` or an error and no token, so a cut answer looked complete or was lost. It now ends the way a refused follow-up does, with the answer so far, `finish_reason` length and that request's token Every continuation stream request now also gets the caller's per-request timeout, which #46349 added to the first one, so a follow-up that never answers ends at that timeout the same way * fix(vertex_ai): keep every follow-up's citations and URL context in a continued answer A non-streamed answer that took several Google requests kept citationMetadata and urlContextMetadata from the last request only, and left every later request's offsets counting from that request's own text. So a URL read on the first request was missing from the answer, citations from earlier requests were dropped, and a later request's citations and search supports pointed at the wrong part of the joined answer The joined answer now keeps every request's citations (Vertex AI's `citations` and AI Studio's `citationSources`) and URL context, and shifts each later request's citation and search support offsets by the UTF-8 bytes of the answer text sent before it, thought parts excluded, since Google counts these offsets in bytes * fix(vertex_ai): answer when a candidate sends null citationMetadata or urlContextMetadata A Gemini candidate whose citationMetadata or urlContextMetadata was JSON null made the call fail on every Gemini model, non-streamed with a 422 or 500 and streamed with an error event, because the per-request metadata helper handed the null back as is and the caller tried to extend a list with it. A continued answer also failed when a follow-up's citation or search support came with a null startIndex A null field is now kept as one entry on a single request, as it was before this branch, a follow-up's null URL context adds nothing to a continued answer, and a null startIndex counts from the start of that request's text. Tests cover missing, null and empty fields on one request and on a continued answer, streamed and not * fix(vertex_ai): pass through spans with no end index and point a continued stream's citations into the whole answer A later request's citation or search support with no endIndex made the whole non-streamed answer fail with a 500, so a span Google sends without an end index now comes back as Google sent it, the way one request for the whole answer returns it. A span with an end index but no start index still counts from 0 A continued stream handed callers each later request's citation and search support offsets counted from that request's own text. The stream now adds the answer bytes of every earlier request, so the offsets point into the whole answer like the non-streamed path and one request for the whole answer Adds the audit's wire tests for both, which fail 24 cells before this change * test(vertex_ai): cover a later request whose cited span starts at its first byte with no start index Google leaves out a zero startIndex, so the second request of a continued answer can cite its own opening with only an endIndex. These wire cells check that span lands where one request for the whole answer puts it, on Gemini and Vertex AI, chat and Responses, unary and streamed * fix(vertex_ai): log a continued answer's raw response as one object A non-streamed answer joined from several Google requests logged its raw response as a JSON list of every request's body. The OpenTelemetry logger reads that raw response as one Gemini body to set its `llm.gemini.*` attributes, so on a continued answer it hit a validation error and set none of them It now logs the merged body, which carries only the last request's token, and a single request still logs Google's body unchanged. The new unit test fails with the change reverted Also adds integration cells for a cut chat stream's token chunk, an answer joined from 41 requests, a caller that gave up on a non-streamed call, counts and caps that are not numbers or are sent twice, `n` sent as true, a partial answer cached after a refused continuation, the token in another provider's cache key, and the sync SDK path with its Vertex token refetch * test(integration): check a continued answer logs one merged raw response A non-streamed Gemini answer that took several Google requests now hands logging one merged response, so the OpenTelemetry span carries the llm.gemini attributes. These 12 cells fail on the merge base and on the commit before that fix, and pass on the branch * test(vertex_ai): cover a follow-up tool call and its thought signature in long decode * test(gemini): pin that a one-request answer logs Google's body unchanged
This commit is contained in:
Mateo Wang
GitHub
parent
a02f1b4939
commit
7090588bdf
29 changed files with 5654 additions and 166 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
{},
|
||||
)
|
||||
|
||||
|
|
|
|||
354
litellm/llms/vertex_ai/gemini/long_decode.py
Normal file
354
litellm/llms/vertex_ai/gemini/long_decode.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
),
|
||||
|
|
|
|||
594
tests/integration/providers/_gemini_long_decode_support.py
Normal file
594
tests/integration/providers/_gemini_long_decode_support.py
Normal file
|
|
@ -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)
|
||||
470
tests/integration/providers/test_gemini_long_decode_chaos.py
Normal file
470
tests/integration/providers/test_gemini_long_decode_chaos.py
Normal file
|
|
@ -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))
|
||||
2106
tests/integration/providers/test_gemini_long_decode_wire.py
Normal file
2106
tests/integration/providers/test_gemini_long_decode_wire.py
Normal file
File diff suppressed because it is too large
Load diff
204
tests/integration/sdk/test_gemini_long_decode_sync_sdk.py
Normal file
204
tests/integration/sdk/test_gemini_long_decode_sync_sdk.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ NATIVE_TOP_LEVEL_FIELD_CASES = [
|
|||
("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}),
|
||||
("cachedContent", "cachedContents/abc123"),
|
||||
("labels", {"team": "search"}),
|
||||
("continuationToken", "dG9rLTE="),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 == {}
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
1310
tests/unit/llms/vertex_ai/gemini/test_long_decode.py
Normal file
1310
tests/unit/llms/vertex_ai/gemini/test_long_decode.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
42
tests/unit/test_gemini_4_argon_model_metadata.py
Normal file
42
tests/unit/test_gemini_4_argon_model_metadata.py
Normal file
|
|
@ -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
|
||||
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue