From e458aa1230a1e23a06aa6ebff88561f5fa510dda Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 15 Jul 2026 08:12:49 -0700
Subject: [PATCH] fix(interactions): bill google_search grounding queries per
 query

---
 .../usage_object_transformation.py            | 12 ++++++-
 .../test_usage_object_transformation.py       | 29 +++++++++++++++
 tests/test_litellm/test_cost_calculator.py    | 36 +++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py
index 3f591180b77..23fafc4132d 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py
@@ -62,6 +62,14 @@ def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, i
     }
 
 
+def _google_search_query_count(usage_object: Mapping[str, Any]) -> int:
+    return sum(
+        _token_count(entry.get("count"))
+        for entry in tuple(usage_object.get("grounding_tool_count") or ())
+        if isinstance(entry, Mapping) and entry.get("type") == "google_search"
+    )
+
+
 def _subtract_cached_from_input(
     input_sums: Mapping[str, int],
     cached_sums: Mapping[str, int],
@@ -116,12 +124,14 @@ class InteractionsUsageObjectTransformation:
         completion_tokens = _token_count(usage_object.get("total_output_tokens")) + reasoning_tokens
         total_tokens = _token_count(usage_object.get("total_tokens")) or (prompt_tokens + completion_tokens)
 
+        web_search_requests = _google_search_query_count(usage_object)
         prompt_tokens_details = (
             PromptTokensDetailsWrapper(
                 cached_tokens=total_cached_tokens or None,
+                web_search_requests=web_search_requests or None,
                 **input_sums,
             )
-            if input_sums or total_cached_tokens
+            if input_sums or total_cached_tokens or web_search_requests
             else None
         )
         completion_tokens_details = (
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
index 01241cf260d..2d8092959ce 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py
@@ -106,6 +106,35 @@ def test_tool_use_tokens_billed_as_input():
     assert usage.prompt_tokens_details.text_tokens == 140
 
 
+def test_google_search_grounding_count_maps_to_web_search_requests():
+    usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
+        {
+            "total_input_tokens": 103,
+            "input_tokens_by_modality": [{"modality": "text", "tokens": 103}],
+            "total_output_tokens": 226,
+            "total_thought_tokens": 351,
+            "grounding_tool_count": [
+                {"type": "google_search", "count": 3},
+                {"type": "url_context", "count": 2},
+            ],
+        }
+    )
+    assert usage.prompt_tokens_details is not None
+    assert usage.prompt_tokens_details.web_search_requests == 3
+
+
+def test_no_grounding_leaves_web_search_requests_unset():
+    usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
+        {
+            "total_input_tokens": 10,
+            "input_tokens_by_modality": [{"modality": "text", "tokens": 10}],
+            "total_output_tokens": 5,
+        }
+    )
+    assert usage.prompt_tokens_details is not None
+    assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None
+
+
 def test_document_modality_folds_into_text():
     usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
         {
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 42fef833944..b144c614474 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -3514,6 +3514,42 @@ def test_completion_cost_bills_interactions_api_response():
     assert cost > 0
 
 
+def test_completion_cost_bills_interactions_google_search_per_query():
+    from litellm.types.interactions import InteractionsAPIResponse
+
+    model_info = litellm.get_model_info(model="gemini-3-flash-preview", custom_llm_provider="gemini")
+    response = InteractionsAPIResponse(
+        id="interactions/search123",
+        model="gemini-3-flash-preview",
+        status="completed",
+        steps=[],
+        usage={
+            "total_tokens": 680,
+            "total_input_tokens": 103,
+            "input_tokens_by_modality": [{"modality": "text", "tokens": 103}],
+            "total_cached_tokens": 0,
+            "total_output_tokens": 226,
+            "total_tool_use_tokens": 0,
+            "total_thought_tokens": 351,
+            "grounding_tool_count": [{"type": "google_search", "count": 3}],
+        },
+    )
+
+    cost = completion_cost(completion_response=response, custom_llm_provider="gemini")
+
+    per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"]
+    reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"]
+    expected = (
+        103 * model_info["input_cost_per_token"]
+        + 226 * model_info["output_cost_per_token"]
+        + 351 * reasoning_rate
+        + 3 * per_query_cost
+    )
+    assert model_info.get("web_search_billing_unit") == "per_query"
+    assert cost == pytest.approx(expected)
+    assert cost > 3 * per_query_cost
+
+
 def test_completion_cost_bills_interactions_video_output_at_video_rate():
     from litellm.types.interactions import InteractionsAPIResponse