diff --git a/litellm/llms/vertex_ai/gemini/grounding_requests.py b/litellm/llms/vertex_ai/gemini/grounding_requests.py index 40acd9378df..1af7d64c104 100644 --- a/litellm/llms/vertex_ai/gemini/grounding_requests.py +++ b/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) ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 32d7b39a20e..37e35d34fd9 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/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: diff --git a/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py b/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py index ec9642bdef6..f7d7a8e67d8 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py +++ b/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)