Commit graph

7 commits

Author SHA1 Message Date
alexcastgonz
6762d65eb8
fix(vertex_ai): price streamed pass-through calls with the regional uplift (#44795)
* fix(vertex_ai): price streamed pass-through calls with the regional uplift

The streamed Vertex pass-through handler builds the response with
stream_chunk_builder, which stamps a cost into _hidden_params without the
vertex_location. The handler then prices the response again with the
location from the URL, but only into kwargs, and the logger prefers the
stamped cost. Streamed calls on a regional endpoint were billed at the
global rate.

Stamp the location-aware cost on the response in the shared
generateContent payload builder so the logger records it.

Fixes #44755

* style(vertex_ai): keep the response_cost assignment within the 120-char limit

* refactor(vertex_ai): set the passthrough response cost with set_hidden_param

Writing _hidden_params directly mutated a parameter and needed a pyright ignore, which the type discipline gate counts as LIT004 and LIT011. set_hidden_param is the shared helper for this.

* test(integration): check Vertex pass-through spend rows keep the regional uplift when streamed

* fix(proxy): price streamed Vertex pass-through spend like the non-streamed call

Drop the location-blind cost stream_chunk_builder stores instead of stamping a cost, so the logger prices streamed generateContent and rawPredict calls with the URL's location and the deployment's custom pricing, the same way it prices non-streamed calls

* fix(proxy): price custom-priced Vertex pass-through deployments at their own rates

The stamped response_cost from the generateContent logging helper is returned before the
logger checks custom pricing, so a deployment with custom input/output rates was billed at
the cost-map price. The helper now passes custom_pricing and the deployment id into
completion_cost, so the stamp holds the deployment rate for streamed and non-streamed calls,
and deployments without custom pricing keep the regional uplift.

* fix(proxy): bill Vertex /v1beta and interrupted pass-through streams by the deployment location and price

* fix(proxy): give the new pyright ignore a reason for the LIT004 gate

* test(integration): bill a custom-priced Vertex pass-through stream at the deployment price, cut-off streams included

* fix(proxy): let the logger price non-streamed /gemini pass-through calls so both modes bill the same

* fix(proxy): keep a client base_model out of pass-through pricing and price gateway Gemini streams as Gemini

A top-level base_model in a pass-through body now reaches the logging
params only for a key or team with allow_client_pricing_override, so a
plain key cannot pick a cheaper model's price once the logger prices
these calls. The test that pinned every litellm-owned body name into the
logging params now leaves base_model to its own key-gated test.

The Vertex handler now classifies a non-Google upstream host as Vertex
only for Vertex resource paths, so a /gemini stream through a
GEMINI_API_BASE gateway is priced as Gemini without the regional uplift.

Adds the audit's integration spend test for Claude on Vertex, /v1beta,
the named custom-priced /gemini deployment, and a cut OpenAI stream.

* test(integration): price the gateway /gemini stream case as Gemini and cover client base_model and gateway host pricing

The cost harness sends /gemini through GEMINI_API_BASE, which the pass-through now logs and bills as gemini, so the streamed case moves from the Vertex cases to the Gemini cases under the Gemini price. New spend tests cover a body base_model repricing only for keys allowed client pricing, a GEMINI_API_BASE gateway billing as Gemini in both modes, and a configured non-Google gateway picking the provider from the path.

* fix(proxy): price pass-through Gemini by the route and ignore a non-string base_model

The Vertex pass-through handler picked Gemini pricing from the URL path,
so a configured pass-through to a Vertex gateway with a short path (no
/projects/ segment) was billed as Gemini, and its Claude streams at $0.
The provider now comes from the route: a call through /gemini or to
generativelanguage.googleapis.com prices as Gemini, and everything else
as Vertex, the same rule main used. Matching the GEMINI_API_BASE host
was rejected because one gateway can serve both APIs, and the path rule
was rejected because short gateway paths look like Gemini paths.

A non-string base_model (a number, list or object) made the cost
calculation fail and bill $0 once the stamped stream costs were gone.
get_base_model_from_metadata now ignores it and falls back to the
deployment's base_model, so the call bills the list price. A 400 was the
other option, but the logger runs after the upstream answered.

* test(integration): bill a non-string body base_model at the list price on every pass-through route

* fix(proxy): bill a non-string model_info base_model at the list price and drop a non-object client model_info

A key allowed client pricing can send metadata.model_info in a pass-through
body. A non-string model_info.base_model made the cost lookup fail, and
once the stamped stream and Gemini costs were gone those calls billed $0
on /gemini and on streamed /vertex_ai and /v1beta. The lookup now ignores
a base_model that is not a string, the same rule the root base_model got,
so the call bills the list price.

A model_info that is not an object (a string or a list) broke several
logging steps that read it as a mapping, so a non-streamed /gemini call
fired no success callback and no pass-through call wrote a spend row. The
pass-through route now leaves a non-object model_info out of the logging
metadata, so the call bills the list price. A 400 was the other option,
but it would turn calls that answer 200 today into errors, and guarding
every reader of model_info would spread the same check over many files.

* fix: reject a non-object model_info before the provider call again

The previous commit made the base_model lookup return None for a
model_info that is not an object. That lookup also runs during call
setup on every router route and SDK entrypoint, where such a value used
to fail before any provider request. With the early return the call went
upstream, the provider billed it, and the request then failed in a later
reader that still expects an object. On /v1/messages and /v1/responses
the key was also charged for a response the client never got.

The lookup again reads base_model from any truthy model_info, so a
non-object value fails during setup the way it did before this PR. A
non-string base_model inside an object still bills the list price, and
the pass-through route still drops a non-object client model_info at
intake, so pass-through calls keep billing the list price.

* fix: match the Gemini API host exactly when picking the pass-through cost provider

The provider check accepted any host ending in
generativelanguage.googleapis.com, so a Vertex host such as
vertex-generativelanguage.googleapis.com was priced as Gemini and lost
the regional uplift. It now accepts that host or a subdomain of it.

* test: pin that a subdomain of the Gemini API host keeps the Gemini price

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-19 08:10:52 +00:00
Mateo Wang
7090588bdf
feat(gemini)!: auto-continue gemini-4-argon long decodes with continuationToken (#46240)
* feat(gemini): continue MAX_TOKENS decodes with continuationToken across chat, streaming, and interactions

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* test(gemini): pin that a one-request answer logs Google's body unchanged
2026-10-15 22:54:23 +00:00
devin-ai-integration[bot]
17f621e116
fix(vertex_ai): route Model Garden models through the generateContent adapter (#41836)
* fix(vertex_ai): route Model Garden models through the generateContent adapter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(vertex_ai): annotate generateContent route test parameters

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(vertex_ai): type litellm_params in get_vertex_ai_model_route

* fix(vertex_ai): bill and log Model Garden generateContent calls like chat completions

Streamed Model Garden calls through the generateContent adapter now request
stream usage where the deployment supports it, keep cached-token details, and
look up the provider-relative cost key, so they bill like the non-stream call.
The adapter forwards router metadata, passes upstream errors through with their
status, countTokens falls back to local counting for non-Gemini routes, and the
native Gemini config is limited to the Gemini routes.

* fix(google_genai): keep adapter stream usage, admin opt-out, single failure count, and base_model routing

Streamed generateContent through the chat adapter now ends with the upstream usage, honors
always_include_stream_usage false and the caller's stream_options, counts one failure per
failed call, and routes Vertex deployments by their base_model like chat does

* fix(google_genai): ask for stream usage on every adapter stream so Claude streams end with usage

* test(google_genai): expect the trailing usage chunk on bridged OpenAI streams

The OpenAI bridge test now sends the usage chunk OpenAI returns for include_usage and expects it translated into usageMetadata. The response_format support helper goes back to its single-purpose form since its only other caller is gone.

* fix(google_genai): end adapter streams with one STOP event carrying usage, and keep upstream usage after the finish chunk

* fix(google_genai): send the held STOP event when no usage arrives, keep a caller's stream_options, and route countTokens like generateContent

* test(integration): cover the moved Vertex routes, adapter stream usage, the shared iterator and usage settings over the wire

* test(integration): add cache-hit twins for the stream usage setting and non-stream cells for every iterator member

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-15 08:06:49 -07:00
devin-ai-integration[bot]
e2971e0af4
refactor(llms): expose public names for private provider helpers (#45037)
* refactor(litellm): migrate private usage in llms

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(litellm): preserve Bedrock batch signature marker

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(litellm): migrate llms private usage symbols

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(litellm): preserve UUID Watsonx project IDs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(litellm): retarget llms mocks to public names

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(litellm): limit llms changes to renames and forwarders

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(litellm): keep GCS mock client patching private Vertex auth methods

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(litellm): assert forwarder arguments and type forwarder signatures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-07 08:49:25 -07:00
devin-ai-integration[bot]
c3a23fe499
refactor: expose core private helpers under public names (#44871)
* refactor: expose core private symbols with compatibility aliases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* types: narrow core migration diagnostics

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve runtime behavior in core symbol migration

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve core private usage migration behavior

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve private value rebinding compatibility

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: restore optional imports and cover public helpers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: exempt router property from call coverage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: align recursive detector ignore names

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-07 00:05:27 +00:00
hsm207
319b08b4b1
fix(google_genai): preserve proxy_server_request in completion adapter (#43536) 2026-09-28 21:29:55 -07:00
yuneng-jiang
f6882246d4
test: move tests/test_litellm root and small trees into tests/unit (#43186)
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: rename fork-flag to unit-flag now that it applies on every event

* test: move tests/test_litellm root and small trees into tests/unit

Pure renames, no content changes. Follow-up commits in this PR fix
references, merge the three files that already existed in tests/unit,
keep live-provider tests in tests/test_litellm and wire CI.

* test: carry tests/test_litellm conftest isolation into tests/unit

Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS,
proxy-URL and keychain env, and session-end client cleanup now reset for
unit tests too. The environment isolation owns its MonkeyPatch so a test's
own monkeypatch is undone before the model-cost teardown runs.

* test: merge, split and prune the moved root and small-tree tests

Merge batches/test_batch_utils.py and the chat_completions and messages
dispatch tests into the files that already existed in tests/unit. Keep
the live Gemini interactions tests, the async image-fetch format test and
the OpenAI embedding scorer test in tests/test_litellm since they need
real network or keys. Put test_router.py under tests/unit/test_router so
the existing package no longer shadows it. Delete eight tests the audit
found superseded by stronger ones kept in this move.

* ci: run the moved root and small-tree tests under their legacy flags

Add the misc and responses-caching-types flags to unit_selection.sh and
CircleCI, extend enterprise-routing and mcp-integration, and point the
legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest
and change classifier at the new paths.

* test: make the new tests/unit directories packages

tests/unit/test_package_layout.py requires every directory to carry an
__init__.py, and without one the moved and retained
test_litellm_responses_bridge.py modules collide on import.

* test: scope the unit socket block to tests/unit in shared sessions

The GHA shards collect the legacy test-path and the unit selection in one
pytest session. The unit conftest's loopback-only block leaked into legacy
modules that reach the network at import. The legacy conftest now lifts the
restriction at collect and setup time, and the unit conftest re-applies it
when collecting its own modules.

* test: give the shard-script tests their own GITHUB_OUTPUT

They only passed where the runner set it. The CircleCI unit job's env
allowlist drops it, so the script's redirect failed there.

* test: point the router and module-deletion checks at tests/unit

router_code_coverage and code_qa_check_tests only searched tests/test_litellm,
so the moved router tests no longer counted. The two silent-experiment tests
the audit deleted were the only direct callers of those methods; they are
replaced with tests that assert the forwarded shadow request and the
recursion guard.

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-25 11:30:43 -07:00