mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(vertex_ai): count server-side Google Search invocations in web_search_requests
When include_server_side_tool_invocations is enabled, Gemini reports completed Google searches as GOOGLE_SEARCH_WEB toolCall parts carrying the queries it ran, instead of (or in addition to) webSearchQueries in groundingMetadata. Only the grounding-metadata queries were counted, so these searches were billed at zero. Extract the distinct queries from the server-side invocations and union them with the grounding-metadata queries before setting usage.prompt_tokens_details. web_search_requests, in both the non-streaming and streaming paths. A query reported in both places is still counted once. Fixes #49012
This commit is contained in:
parent
e63fc85b7b
commit
217ca5b379
4 changed files with 157 additions and 11 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue