Merge 59cb5aabee into 570f41daad

litellm/llms/vertex_ai/gemini/grounding_requests.py

@@ -2,6 +2,56 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Final
 
+# toolType values Gemini uses for Google Search in server-side tool invocations
+# (includeServerSideToolInvocations / include_server_side_tool_invocations).
+_GOOGLE_SEARCH_INVOCATION_TYPES: Final = frozenset(
+    {
+        "GOOGLE_SEARCH_WEB",
+        "GOOGLE_SEARCH",
+        "GOOGLE_SEARCH_RETRIEVAL",
+        "google_search",
+        "googleSearch",
+        "google_search_retrieval",
+        "googleSearchRetrieval",
+    }
+)
+
+
+def web_search_queries_from_server_side_invocations(
+    invocations: Sequence[Mapping[str, object]] | None,
+) -> frozenset[str]:
+    """Distinct nonempty search queries from Gemini server-side tool invocations.
+
+    When ``include_server_side_tool_invocations`` is on, search evidence often
+    arrives only as ``toolCall`` / ``toolResponse`` parts (``args.queries`` or
+    ``args.query``), not as ``groundingMetadata.webSearchQueries``. Count those
+    the same way grounding metadata is counted so usage and tool cost stay
+    accurate (#49012).
+    """
+    if not invocations:
+        return frozenset()
+    queries: set[str] = set()
+    for invocation in invocations:
+        if not isinstance(invocation, Mapping):
+            continue
+        tool_type = invocation.get("tool_type")
+        if tool_type is None:
+            tool_type = invocation.get("toolType")
+        if tool_type not in _GOOGLE_SEARCH_INVOCATION_TYPES:
+            if not (isinstance(tool_type, str) and "search" in tool_type.lower() and "map" not in tool_type.lower()):
+                continue
+        args = invocation.get("args")
+        if not isinstance(args, Mapping):
+            continue
+        raw_queries = args.get("queries")
+        if isinstance(raw_queries, list):
+            queries.update(q for q in raw_queries if isinstance(q, str) and q)
+        raw_query = args.get("query")
+        if isinstance(raw_query, str) and raw_query:
+            queries.add(raw_query)
+    return frozenset(queries)
+
+
 
 @dataclass(frozen=True, slots=True)
 class GroundingRequests:
@@ -34,18 +84,25 @@ 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:
+def calculate_grounding_requests(
+    grounding_metadata: Sequence[Mapping[str, object]],
+    server_side_invocations: Sequence[Mapping[str, object]] | None = None,
+) -> 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.
+
+    When search evidence is only in server-side tool invocations (#49012), those
+    queries are unioned with grounding-metadata queries so each distinct string
+    still counts once.
     """
     items: Final = tuple(item for item in grounding_metadata if isinstance(item, Mapping))
     web_queries: Final = frozenset(
         query for item in items if not _attributes_queries_to_maps(item) for query in _queries(item)
-    )
+    ) | web_search_queries_from_server_side_invocations(server_side_invocations)
     maps_queries: Final = frozenset(
         query for item in items if _attributes_queries_to_maps(item) for query in _queries(item)
     )

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -109,7 +109,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, web_search_queries_from_server_side_invocations
 from .long_decode import (
     HOP_FAILURES,
     AsyncStreamResume,
@@ -2608,8 +2608,14 @@ 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]],
+        server_side_invocations: Sequence[Mapping[str, object]] | None = None,
+    ) -> None:
+        grounding_requests: Final = calculate_grounding_requests(
+            grounding_metadata, server_side_invocations=server_side_invocations
+        )
         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
@@ -2621,8 +2627,33 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         cls,
         usage: Usage,
         grounding_metadata: Sequence[Mapping[str, object]],
+        server_side_invocations: Sequence[Mapping[str, object]] | None = None,
     ) -> None:
-        return cls._set_grounding_usage_counters(usage, grounding_metadata)
+        return cls._set_grounding_usage_counters(
+            usage, grounding_metadata, server_side_invocations=server_side_invocations
+        )
+
+    @staticmethod
+    def _collect_server_side_tool_invocations_from_response(
+        model_response: Union[ModelResponse, "ModelResponseStream"],
+    ) -> list[dict[str, object]]:
+        """Gather server_side_tool_invocations already placed on choices."""
+        from litellm.types.utils import ModelResponseStream
+
+        collected: list[dict[str, object]] = []
+        for choice in getattr(model_response, "choices", None) or []:
+            message = None
+            if isinstance(model_response, ModelResponseStream):
+                message = getattr(choice, "delta", None)
+            else:
+                message = getattr(choice, "message", None)
+            fields = getattr(message, "provider_specific_fields", None) if message is not None else None
+            if not isinstance(fields, dict):
+                continue
+            invocations = fields.get("server_side_tool_invocations")
+            if isinstance(invocations, list):
+                collected.extend(inv for inv in invocations if isinstance(inv, dict))
+        return collected
 
     @staticmethod
     def _create_streaming_choice(
@@ -3131,7 +3162,13 @@ 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_invocations=VertexGeminiConfig._collect_server_side_tool_invocations_from_response(
+                    model_response
+                ),
+            )
 
             setattr(model_response, "usage", usage)
 
@@ -3985,6 +4022,10 @@ class ModelResponseIterator:
         self._initial_prompt_tokens: int | None = None
         self._grounding_offset: tuple[Mapping[str, object], ...] = ()
         self._hop_grounding: tuple[Mapping[str, object], ...] = ()
+        # Distinct Google Search queries seen in server-side tool invocations
+        # across stream chunks (usage often arrives in a later chunk without
+        # groundingMetadata when include_server_side_tool_invocations is on).
+        self._server_side_search_queries: set[str] = set()
         self._answer_offset: int = 0
         self._hop_answer_bytes: int = 0
 
@@ -4092,6 +4133,12 @@ class ModelResponseIterator:
             citation_metadata,
         )
 
+        self._server_side_search_queries.update(
+            web_search_queries_from_server_side_invocations(
+                VertexGeminiConfig._collect_server_side_tool_invocations_from_response(model_response)
+            )
+        )
+
         return (
             grounding_metadata,
             url_context_metadata,
@@ -4127,7 +4174,23 @@ class ModelResponseIterator:
             cache_creation_tokens=_cache_creation_tokens(self.logging_obj),
         )
 
-        VertexGeminiConfig.set_grounding_usage_counters(usage, self._grounding_through_this_hop(grounding_metadata))
+        # Rebuild minimal invocation records so calculate_grounding_requests can
+        # union stream-accumulated search queries with any groundingMetadata.
+        accumulated_invocations: Final = (
+            [
+                {
+                    "tool_type": "GOOGLE_SEARCH_WEB",
+                    "args": {"queries": sorted(self._server_side_search_queries)},
+                }
+            ]
+            if self._server_side_search_queries
+            else None
+        )
+        VertexGeminiConfig.set_grounding_usage_counters(
+            usage,
+            self._grounding_through_this_hop(grounding_metadata),
+            server_side_invocations=accumulated_invocations,
+        )
 
         traffic_type: Final = processed_chunk.get("usageMetadata", {}).get("trafficType")
         if traffic_type:

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

@@ -1,6 +1,7 @@
 from litellm.llms.vertex_ai.gemini.grounding_requests import (
     GroundingRequests,
     calculate_grounding_requests,
+    web_search_queries_from_server_side_invocations,
 )
 
 
@@ -100,3 +101,71 @@ 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_web_search_queries_from_server_side_invocations_queries_list():
+    invocations = [
+        {
+            "tool_type": "GOOGLE_SEARCH_WEB",
+            "id": "call_1",
+            "args": {"queries": ["top news headlines today"]},
+        },
+        {
+            "tool_type": "GOOGLE_SEARCH_WEB",
+            "id": "call_2",
+            "args": {
+                "queries": [
+                    '"MV Odyssey" suspended voyage fuel costs CBS News',
+                    "Cambodia US Cruise Ship MV Odyssey fuel costs",
+                    "top news headlines today",  # duplicate across calls
+                ]
+            },
+        },
+        {
+            # function tool — must be ignored
+            "tool_type": "FUNCTION",
+            "args": {"queries": ["should not count"]},
+        },
+        {
+            "tool_type": "google_search",
+            "args": {"query": "single query form"},
+        },
+        {"tool_type": "GOOGLE_SEARCH_WEB", "args": {"queries": ["", "kept"]}},
+    ]
+
+    result = web_search_queries_from_server_side_invocations(invocations)
+    assert result == frozenset(
+        {
+            "top news headlines today",
+            '"MV Odyssey" suspended voyage fuel costs CBS News',
+            "Cambodia US Cruise Ship MV Odyssey fuel costs",
+            "single query form",
+            "kept",
+        }
+    )
+
+
+def test_calculate_grounding_requests_unions_invocation_queries():
+    result = calculate_grounding_requests(
+        [{"webSearchQueries": ["from metadata", "shared"]}],
+        server_side_invocations=[
+            {
+                "tool_type": "GOOGLE_SEARCH_WEB",
+                "args": {"queries": ["shared", "from invocation"]},
+            }
+        ],
+    )
+    assert result == GroundingRequests(web_search_requests=3, google_maps_grounding_requests=None)
+
+
+def test_calculate_grounding_requests_invocations_only():
+    result = calculate_grounding_requests(
+        [],
+        server_side_invocations=[
+            {
+                "tool_type": "GOOGLE_SEARCH_WEB",
+                "args": {"queries": ["a", "b", "a"]},
+            }
+        ],
+    )
+    assert result == GroundingRequests(web_search_requests=2, google_maps_grounding_requests=None)