diff --git a/litellm/llms/vertex_ai/gemini/grounding_requests.py b/litellm/llms/vertex_ai/gemini/grounding_requests.py index 40acd9378df..877f1e517c4 100644 --- a/litellm/llms/vertex_ai/gemini/grounding_requests.py +++ b/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( 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 c506192f253..28b3e7f9b11 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 @@ -108,7 +108,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, @@ -2543,8 +2543,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 @@ -2556,8 +2560,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( @@ -3052,7 +3057,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) @@ -3893,6 +3900,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 @@ -4034,7 +4042,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: @@ -4087,6 +4097,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, diff --git a/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py b/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py index 39c40eeae67..8ab1607f728 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py +++ b/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 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..e6e71f54360 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py +++ b/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)