Merge 217ca5b379 into 0eb6816081

litellm/llms/vertex_ai/gemini/grounding_requests.py

@@ -34,16 +34,60 @@ def _attributes_queries_to_maps(item: Mapping[str, object]) -> bool:
     return _is_maps_item(item) and "web" not in _chunk_kinds(item)
 
 
-def calculate_grounding_requests(grounding_metadata: Sequence[Mapping[str, object]]) -> GroundingRequests:
+_SERVER_SIDE_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"GOOGLE_SEARCH_WEB", "GOOGLE_SEARCH"})
+
+
+def _tool_call_search_queries(part: Mapping[str, object]) -> frozenset[str]:
+    tool_call: Final = part.get("toolCall")
+    if not isinstance(tool_call, Mapping) or tool_call.get("toolType") not in _SERVER_SIDE_WEB_SEARCH_TOOL_TYPES:
+        return frozenset()
+    args: Final = tool_call.get("args")
+    queries: Final = args.get("queries") if isinstance(args, Mapping) else None
+    if not isinstance(queries, list):
+        return frozenset()
+    return frozenset(query for query in queries if isinstance(query, str) and query)
+
+
+def server_side_search_queries(candidates: Sequence[Mapping[str, object]] | None) -> frozenset[str]:
+    """Distinct Google Search queries carried by server-side tool invocations.
+
+    When ``includeServerSideToolInvocations`` is enabled, Gemini reports completed
+    searches as ``toolCall`` parts (``toolType`` GOOGLE_SEARCH_WEB) holding the queries
+    it ran, instead of, or in addition to, ``webSearchQueries`` in groundingMetadata.
+    These are billed the same way, so they have to be counted too (#49012).
+    """
+    if not candidates:
+        return frozenset()
+    queries: set[str] = set()
+    for candidate in candidates:
+        if not isinstance(candidate, Mapping):
+            continue
+        content: Final = candidate.get("content")
+        parts: Final = content.get("parts") if isinstance(content, Mapping) else None
+        if not isinstance(parts, list):
+            continue
+        for part in parts:
+            if isinstance(part, Mapping):
+                queries |= _tool_call_search_queries(part)
+    return frozenset(queries)
+
+
+def calculate_grounding_requests(
+    grounding_metadata: Sequence[Mapping[str, object]],
+    extra_web_queries: frozenset[str] = frozenset(),
+) -> GroundingRequests:
     """Billable grounding requests across candidates, counting each distinct query once.
 
     Duplicate queries within and across grounding metadata items collapse to the
-    distinct-query count (#36377), and empty strings are ignored. Maps grounding is
-    floored at one request whenever a candidate carries maps chunks or a widget token,
-    since per-prompt billing charges the prompt even when no query is reported.
+    distinct-query count (#36377), and empty strings are ignored. ``extra_web_queries``
+    carries web-search queries reported outside groundingMetadata (server-side tool
+    invocations, #49012); they join the distinct web-query set, so a query reported in
+    both places is still counted once. Maps grounding is floored at one request whenever
+    a candidate carries maps chunks or a widget token, since per-prompt billing charges
+    the prompt even when no query is reported.
     """
     items: Final = tuple(item for item in grounding_metadata if isinstance(item, Mapping))
-    web_queries: Final = frozenset(
+    web_queries: Final = extra_web_queries | frozenset(
         query for item in items if not _attributes_queries_to_maps(item) for query in _queries(item)
     )
     maps_queries: Final = frozenset(

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -115,7 +115,7 @@ from ..context_caching.rejected_cache_retry import (
     post_with_fresh_cache_sync,
 )
 from ..vertex_llm_base import VertexBase
-from .grounding_requests import calculate_grounding_requests
+from .grounding_requests import calculate_grounding_requests, server_side_search_queries
 from .long_decode import (
     HOP_FAILURES,
     AsyncStreamResume,
@@ -2772,8 +2772,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         return cls._calculate_web_search_requests(grounding_metadata)
 
     @staticmethod
-    def _set_grounding_usage_counters(usage: Usage, grounding_metadata: Sequence[Mapping[str, object]]) -> None:
-        grounding_requests: Final = calculate_grounding_requests(grounding_metadata)
+    def _set_grounding_usage_counters(
+        usage: Usage,
+        grounding_metadata: Sequence[Mapping[str, object]],
+        extra_web_search_queries: frozenset[str] = frozenset(),
+    ) -> None:
+        grounding_requests: Final = calculate_grounding_requests(grounding_metadata, extra_web_search_queries)
         details: Final = cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details)
         if grounding_requests.web_search_requests is not None:
             details.web_search_requests = grounding_requests.web_search_requests
@@ -2785,8 +2789,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         cls,
         usage: Usage,
         grounding_metadata: Sequence[Mapping[str, object]],
+        extra_web_search_queries: frozenset[str] = frozenset(),
     ) -> None:
-        return cls._set_grounding_usage_counters(usage, grounding_metadata)
+        return cls._set_grounding_usage_counters(usage, grounding_metadata, extra_web_search_queries)
 
     @staticmethod
     def _create_streaming_choice(
@@ -3295,7 +3300,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 cache_creation_tokens=_cache_creation_tokens(logging_obj),
             )
 
-            VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)
+            VertexGeminiConfig._set_grounding_usage_counters(
+                usage, grounding_metadata, server_side_search_queries(_candidates)
+            )
 
             setattr(model_response, "usage", usage)
 
@@ -4149,6 +4156,7 @@ class ModelResponseIterator:
         self._initial_prompt_tokens: int | None = None
         self._grounding_offset: tuple[Mapping[str, object], ...] = ()
         self._hop_grounding: tuple[Mapping[str, object], ...] = ()
+        self._search_queries: frozenset[str] = frozenset()
         self._answer_offset: int = 0
         self._hop_answer_bytes: int = 0
 
@@ -4291,7 +4299,9 @@ class ModelResponseIterator:
             cache_creation_tokens=_cache_creation_tokens(self.logging_obj),
         )
 
-        VertexGeminiConfig.set_grounding_usage_counters(usage, self._grounding_through_this_hop(grounding_metadata))
+        VertexGeminiConfig.set_grounding_usage_counters(
+            usage, self._grounding_through_this_hop(grounding_metadata), self._search_queries
+        )
 
         traffic_type: Final = processed_chunk.get("usageMetadata", {}).get("trafficType")
         if traffic_type:
@@ -4344,6 +4354,7 @@ class ModelResponseIterator:
 
             _candidates: Final[list[Candidates] | None] = processed_chunk.get("candidates")
             if _candidates:
+                self._search_queries |= server_side_search_queries(_candidates)
                 (
                     grounding_metadata,
                     url_context_metadata,

tests/unit/llms/vertex_ai/gemini/test_context_circulation.py

@@ -17,10 +17,12 @@ import pytest
 from litellm.llms.vertex_ai.gemini.transformation import (
     gemini_convert_messages_with_history,
 )
+from litellm.llms.vertex_ai.gemini.grounding_requests import server_side_search_queries
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     VertexGeminiConfig,
 )
 from litellm.types.llms.vertex_ai import HttpxPartType
+from litellm.types.utils import PromptTokensDetailsWrapper, Usage
 
 # --- Response extraction tests ---
 
@@ -367,3 +369,50 @@ class TestIncludeServerSideToolInvocationsConfig:
         config = VertexGeminiConfig()
         supported = config.get_supported_openai_params(model="gemini-3-flash-preview")
         assert "include_server_side_tool_invocations" in supported
+
+
+# --- Usage counting for server-side search invocations (#49012) ---
+
+
+def _search_candidates(*query_groups):
+    """Gemini candidates carrying GOOGLE_SEARCH_WEB toolCall parts, one per query group."""
+    return [
+        {"content": {"parts": [{"toolCall": {"toolType": "GOOGLE_SEARCH_WEB", "args": {"queries": list(queries)}}}]}}
+        for queries in query_groups
+    ]
+
+
+class TestServerSideSearchUsageCounting:
+    """Searches reported only as server-side tool invocations are still billed (#49012)."""
+
+    def _web_search_requests(self, usage):
+        return getattr(usage.prompt_tokens_details, "web_search_requests", None)
+
+    def test_invocation_only_searches_counted_when_grounding_metadata_absent(self):
+        candidates = _search_candidates(["weather Buenos Aires", "weather Montevideo"], ["weather Buenos Aires"])
+        usage = Usage(prompt_tokens_details=PromptTokensDetailsWrapper())
+
+        VertexGeminiConfig._set_grounding_usage_counters(
+            usage, grounding_metadata=[], extra_web_search_queries=server_side_search_queries(candidates)
+        )
+
+        assert self._web_search_requests(usage) == 2
+
+    def test_query_reported_in_both_grounding_and_invocation_is_counted_once(self):
+        candidates = _search_candidates(["shared", "invocation only"])
+        grounding_metadata = [{"webSearchQueries": ["shared", "grounding only"], "groundingChunks": [{"web": {}}]}]
+        usage = Usage(prompt_tokens_details=PromptTokensDetailsWrapper())
+
+        VertexGeminiConfig._set_grounding_usage_counters(
+            usage, grounding_metadata, extra_web_search_queries=server_side_search_queries(candidates)
+        )
+
+        assert self._web_search_requests(usage) == 3
+
+    def test_grounding_only_path_is_unchanged_without_invocations(self):
+        grounding_metadata = [{"webSearchQueries": ["q1", "q2"], "groundingChunks": [{"web": {}}]}]
+        usage = Usage(prompt_tokens_details=PromptTokensDetailsWrapper())
+
+        VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)
+
+        assert self._web_search_requests(usage) == 2

tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py

@@ -1,9 +1,14 @@
 from litellm.llms.vertex_ai.gemini.grounding_requests import (
     GroundingRequests,
     calculate_grounding_requests,
+    server_side_search_queries,
 )
 
 
+def _search_candidate(queries, tool_type="GOOGLE_SEARCH_WEB"):
+    return {"content": {"parts": [{"toolCall": {"toolType": tool_type, "args": {"queries": list(queries)}}}]}}
+
+
 def test_search_only_counts_non_empty_queries_as_web_requests():
     result = calculate_grounding_requests(
         [
@@ -100,3 +105,40 @@ def test_has_billable_grounding():
     assert GroundingRequests(web_search_requests=None, google_maps_grounding_requests=1).has_billable_grounding()
     assert GroundingRequests(web_search_requests=1, google_maps_grounding_requests=None).has_billable_grounding()
     assert not GroundingRequests(web_search_requests=None, google_maps_grounding_requests=None).has_billable_grounding()
+
+
+def test_server_side_search_queries_extracts_distinct_google_search_queries():
+    candidates = [
+        _search_candidate(["weather Buenos Aires", "weather Montevideo"]),
+        _search_candidate(["weather Buenos Aires"]),
+    ]
+    assert server_side_search_queries(candidates) == frozenset({"weather Buenos Aires", "weather Montevideo"})
+
+
+def test_server_side_search_queries_ignores_non_search_tool_calls_and_empty_strings():
+    candidates = [
+        _search_candidate(["", "real query"]),
+        _search_candidate(["ignored"], tool_type="CODE_EXECUTION"),
+    ]
+    assert server_side_search_queries(candidates) == frozenset({"real query"})
+
+
+def test_server_side_search_queries_tolerates_missing_and_malformed_shapes():
+    assert server_side_search_queries(None) == frozenset()
+    assert server_side_search_queries([]) == frozenset()
+    assert server_side_search_queries([{"content": {}}, {"content": {"parts": "nope"}}, "not a dict"]) == frozenset()
+    malformed = [{"content": {"parts": [{"toolCall": {"toolType": "GOOGLE_SEARCH_WEB", "args": {"queries": "x"}}}]}}]
+    assert server_side_search_queries(malformed) == frozenset()
+
+
+def test_extra_web_queries_are_counted_when_grounding_metadata_has_none():
+    result = calculate_grounding_requests([], extra_web_queries=frozenset({"q1", "q2"}))
+    assert result == GroundingRequests(web_search_requests=2, google_maps_grounding_requests=None)
+
+
+def test_extra_web_queries_union_with_grounding_metadata_dedupes_shared_query():
+    result = calculate_grounding_requests(
+        [{"webSearchQueries": ["shared", "grounding only"], "groundingChunks": [{"web": {"uri": "https://e.com"}}]}],
+        extra_web_queries=frozenset({"shared", "invocation only"}),
+    )
+    assert result == GroundingRequests(web_search_requests=3, google_maps_grounding_requests=None)