fix(proxy/gemini): keep streaming-parity baseline so eligible_pct can't exceed 100 (#2824)
## Description The non-streaming Gemini `generateContent` finalizer builds its `RequestOutcome` with `optimized_tokens` set to Gemini's own `promptTokenCount` (the provider's tokenizer scale, which correctly feeds billing and the dashboard), while `original_tokens` stays a local estimator count. Those two are on different rulers. Every delta the beacon derives from the pair is a same-ruler difference: `tokens_saved`, `tokens_inflated`, `attempted_input_tokens`, and the beacon's `eligible_pct` / `yield_pct`. When Gemini counts the forwarded prompt higher than our local estimator does, `attempted_input_tokens` (which is `optimized_tokens + tokens_saved`) exceeds the local `original_tokens`, and the request ships a structurally-impossible `eligible_pct > 100` plus a phantom `tokens_inflated`. This is the exact class of bug #2756 removed, on a path #2756 did not touch: it fixed the non-streaming OpenAI handler, and the streaming finalizer (`_finalize_stream_response`) already guards against it by lifting the baseline onto the provider scale. The non-streaming Gemini path had neither treatment. The fix mirrors the streaming finalizer's already-tested handling: when a provider count is present, lift the baseline to `max(original_tokens, promptTokenCount + tokens_saved)` so `attempted_input_tokens <= original_tokens` holds and `tokens_inflated` collapses to 0. It is guarded on a present count, so a null or absent `promptTokenCount` leaves the local baseline untouched and the existing zero-usage preservation test still holds. `optimized_tokens` still carries the provider count, so billing and the dashboard are unchanged. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/handlers/gemini.py` (`handle_gemini_request`, non-streaming `generateContent` branch): compute `effective_original_tokens = max(original_tokens, total_input_tokens + tokens_saved)` when `total_input_tokens > 0` (else keep `original_tokens`), and pass it as the outcome's `original_tokens`. Mirrors the streaming finalizer's provider-usage handling. - `tests/test_proxy/test_gemini_savings_profile.py`: added `test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible`, which drives a request where Gemini's `promptTokenCount` (150) exceeds the local post-compression count (80), and asserts `attempted_input_tokens <= original_tokens`, `tokens_inflated == 0`, the provider count is still carried in `optimized_tokens`, and the baseline is lifted to 170. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text # Fail-before (source fix stashed, new test kept): tests/test_proxy/test_gemini_savings_profile.py::test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible FAILED assert outcome.attempted_input_tokens <= outcome.original_tokens AssertionError: assert 170 <= 100 # Pass-after (fix applied): tests/test_proxy/test_gemini_savings_profile.py::test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible PASSED # Full file + related outcome suites: tests/test_proxy/test_gemini_savings_profile.py tests/test_proxy_gemini_native_integration.py tests/test_request_outcome.py tests/test_outcome_token_scale.py 47 passed, 18 skipped # uvx ruff@0.15.17 check -> All checks passed! # uvx mypy@1.20.2 headroom/proxy/handlers/gemini.py -> Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv (litellm installed), pytest 9.1.1 with pytest-asyncio 1.4.0 (asyncio_mode=auto per pyproject), ruff 0.15.17 and mypy 1.20.2 via uvx. - Exact command / steps: confirmed the streaming sibling already lifts the baseline (`_finalize_stream_response` in `headroom/proxy/handlers/streaming.py` sets `effective_original_tokens = max(original_tokens, provider_input_tokens + tokens_saved)` for openai/gemini), then fail-before with `git stash push headroom/proxy/handlers/gemini.py` and `python -m pytest tests/test_proxy/test_gemini_savings_profile.py -k inflate_eligible` (the assertion fails with `170 <= 100`, i.e. eligible_pct 170%), then pass-after with `git stash pop` and rerunning (passes), then the full file plus the outcome suites (47 passed, 18 skipped). - Observed result: with Gemini reporting `promptTokenCount=150` against a local post-compression count of 80 (saved 20), the outcome now reports `original_tokens=170`, `attempted_input_tokens=170` (so `eligible_pct <= 100`) and `tokens_inflated=0`, while `optimized_tokens` stays 150 so billing and the dashboard are unchanged. Before the fix the same request reported `original_tokens=100`, `attempted_input_tokens=170` (eligible_pct 170%) and `tokens_inflated=50`. - Not tested: a live streamed call to real Gemini/Vertex (no provider credentials in this environment). The provider-count-above-local case is reproduced with a mock response mirroring Gemini's `usageMetadata` shape, and the baseline-lift it mirrors is existing, tested code on the streaming path. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md`: it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Additional Notes Docs and manual testing are N/A: this aligns the non-streaming Gemini finalizer with the already-correct streaming finalizer, no API surface change. The baseline lift is guarded on a present provider count, so the existing zero-usage preservation test (`test_gemini_zero_usage_prompt_count_is_preserved`) is unaffected: a null or zero `promptTokenCount` keeps the local baseline and leaves `optimized_tokens` at 0.
This commit is contained in:
Abhay Singh
@@ -738,6 +738,24 @@ class GeminiHandlerMixin:
|
||||
|
||||
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)
|
||||
|
||||
# optimized_tokens carries Gemini's own promptTokenCount, which is
|
||||
# on the provider's tokenizer scale (it feeds billing/dashboard),
|
||||
# while original_tokens is a LOCAL estimator count. When Gemini
|
||||
# counts the forwarded prompt higher than our estimator does,
|
||||
# attempted_input_tokens (optimized + saved) exceeded the local
|
||||
# original_tokens and shipped a structurally-impossible
|
||||
# eligible_pct > 100 plus a phantom tokens_inflated. Lift the
|
||||
# baseline onto the provider scale when a provider count is
|
||||
# present, mirroring the streaming finalizer's tested handling in
|
||||
# _finalize_stream_response so the two Gemini paths agree. Guarded
|
||||
# on a present count so a null/absent promptTokenCount leaves the
|
||||
# local baseline untouched.
|
||||
effective_original_tokens = (
|
||||
max(original_tokens, total_input_tokens + tokens_saved)
|
||||
if total_input_tokens > 0
|
||||
else original_tokens
|
||||
)
|
||||
|
||||
# Eligible-tracking is TODO for Gemini; pass the full
|
||||
# pre-compression request size as the fallback denominator.
|
||||
# This makes Gemini's contribution to the aggregate
|
||||
@@ -757,7 +775,7 @@ class GeminiHandlerMixin:
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
status_code=response.status_code,
|
||||
original_tokens=original_tokens,
|
||||
original_tokens=effective_original_tokens,
|
||||
optimized_tokens=total_input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
|
||||
@@ -199,3 +199,76 @@ def test_gemini_zero_usage_prompt_count_is_preserved():
|
||||
outcome = captured["outcome"]
|
||||
assert outcome.optimized_tokens == 0
|
||||
assert outcome.uncached_input_tokens == 0
|
||||
|
||||
|
||||
def test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible():
|
||||
"""When Gemini's promptTokenCount exceeds our local estimate, the outcome must
|
||||
not ship attempted_input_tokens > original_tokens (a structurally impossible
|
||||
eligible_pct > 100) or a phantom tokens_inflated. The local baseline is lifted
|
||||
onto the provider scale, matching the streaming finalizer's tested handling."""
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
|
||||
# Local pipeline count: 100 tokens before compression, 80 after (saved 20).
|
||||
# Return genuinely-changed messages so the handler adopts the pipeline's
|
||||
# tokens_before/after (the override only fires when messages actually change).
|
||||
def passthrough_apply(**kwargs):
|
||||
sent = kwargs["messages"]
|
||||
compressed = [dict(m) for m in sent]
|
||||
if compressed:
|
||||
compressed[0] = {**compressed[0], "content": "compressed"}
|
||||
return SimpleNamespace(
|
||||
messages=compressed,
|
||||
transforms_applied=["gemini_compress"],
|
||||
timing={},
|
||||
tokens_before=100,
|
||||
tokens_after=80,
|
||||
waste_signals=None,
|
||||
)
|
||||
|
||||
# Gemini counts the forwarded prompt at 150 -- higher than our local 80, so
|
||||
# attempted = 150 + 20 = 170 would exceed a local original of 100.
|
||||
resp = MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.headers = {"content-type": "application/json"}
|
||||
resp.content = (
|
||||
b'{"candidates":[{"content":{"parts":[{"text":"ok"}]}}],'
|
||||
b'"usageMetadata":{"promptTokenCount":150,"candidatesTokenCount":2}}'
|
||||
)
|
||||
resp.json.return_value = {
|
||||
"candidates": [{"content": {"parts": [{"text": "ok"}]}}],
|
||||
"usageMetadata": {"promptTokenCount": 150, "candidatesTokenCount": 2},
|
||||
}
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def recording_outcome(outcome): # noqa: ANN001
|
||||
captured["outcome"] = outcome
|
||||
|
||||
big = "word " * 4000
|
||||
app = create_app(config)
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
proxy.openai_pipeline.apply = MagicMock(side_effect=passthrough_apply)
|
||||
proxy._retry_request = AsyncMock(return_value=resp)
|
||||
proxy._record_request_outcome = AsyncMock(side_effect=recording_outcome)
|
||||
|
||||
r = client.post(
|
||||
"/v1beta/models/gemini-2.0-flash:generateContent?key=test-key",
|
||||
json={"contents": [{"parts": [{"text": big}]}]},
|
||||
)
|
||||
|
||||
assert r.status_code == 200, r.text
|
||||
outcome = captured["outcome"]
|
||||
# The provider's own count is still carried for billing/dashboard.
|
||||
assert outcome.optimized_tokens == 150
|
||||
# The eligible ratio cannot exceed 100%: attempted must not exceed original.
|
||||
assert outcome.attempted_input_tokens <= outcome.original_tokens
|
||||
# No phantom growth (optimized - original clamped to >= 0 was 50 before).
|
||||
assert outcome.tokens_inflated == 0
|
||||
# Baseline lifted onto the provider scale: max(local 100, provider 150 + saved 20).
|
||||
assert outcome.original_tokens == 170
|
||||
|
||||
Reference in New Issue
Block a user