From fb75a132e5c903916209b43ed4061ad472712ef5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 28 Oct 2026 22:17:42 -0700
Subject: [PATCH] fix(proxy): count snake_case candidate_count in the Gemini
 budget hold

---
 litellm/proxy/spend_tracking/budget_reservation.py    | 11 ++++++++++-
 .../proxy/spend_tracking/test_budget_reservation.py   |  3 ++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py
index 5cada93bd33..4efcc06fdce 100644
--- a/litellm/proxy/spend_tracking/budget_reservation.py
+++ b/litellm/proxy/spend_tracking/budget_reservation.py
@@ -7,6 +7,7 @@ import time
 from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
+from itertools import product
 from types import MappingProxyType
 from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast
 
@@ -1651,6 +1652,7 @@ def _estimate_output_tokens(
 
 
 _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
+_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config")
 
 
 def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object:
@@ -1658,6 +1660,13 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field:
     return config.get(field) if config is not None else None
 
 
+def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]:
+    return tuple(
+        _nested_field(request_body, config_field, field)
+        for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count"))
+    )
+
+
 def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
     candidates: Final = (
         *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
@@ -1671,7 +1680,7 @@ def _get_output_multiplier(request_body: Mapping[str, object]) -> int:
     candidates: Final = (
         request_body.get("n"),
         request_body.get("best_of"),
-        _nested_field(request_body, "generationConfig", "candidateCount"),
+        *_gemini_candidate_counts(request_body),
     )
     return max((1, *(count for count in map(_to_int, candidates) if count is not None)))
 
diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py
index 97f65a4402d..01acfbf9a61 100644
--- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py
+++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py
@@ -193,7 +193,8 @@ GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling")
     (
         ({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}),
         ({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}),
-        ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {}),
+        ({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"n": 3}),
+        ({"generationConfig": {"candidate_count": 3}}, {"n": 3}),
         ({}, {}),
     ),
 )