fix(interactions): bill google_search grounding queries per query

This commit is contained in:
mateo-berri 2026-07-15 08:12:49 -07:00
parent 194dca7dd0
commit e458aa1230
3 changed files with 76 additions and 1 deletions

View file

@ -62,6 +62,14 @@ def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, i
}
def _google_search_query_count(usage_object: Mapping[str, Any]) -> int:
return sum(
_token_count(entry.get("count"))
for entry in tuple(usage_object.get("grounding_tool_count") or ())
if isinstance(entry, Mapping) and entry.get("type") == "google_search"
)
def _subtract_cached_from_input(
input_sums: Mapping[str, int],
cached_sums: Mapping[str, int],
@ -116,12 +124,14 @@ class InteractionsUsageObjectTransformation:
completion_tokens = _token_count(usage_object.get("total_output_tokens")) + reasoning_tokens
total_tokens = _token_count(usage_object.get("total_tokens")) or (prompt_tokens + completion_tokens)
web_search_requests = _google_search_query_count(usage_object)
prompt_tokens_details = (
PromptTokensDetailsWrapper(
cached_tokens=total_cached_tokens or None,
web_search_requests=web_search_requests or None,
**input_sums,
)
if input_sums or total_cached_tokens
if input_sums or total_cached_tokens or web_search_requests
else None
)
completion_tokens_details = (

View file

@ -106,6 +106,35 @@ def test_tool_use_tokens_billed_as_input():
assert usage.prompt_tokens_details.text_tokens == 140
def test_google_search_grounding_count_maps_to_web_search_requests():
usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
{
"total_input_tokens": 103,
"input_tokens_by_modality": [{"modality": "text", "tokens": 103}],
"total_output_tokens": 226,
"total_thought_tokens": 351,
"grounding_tool_count": [
{"type": "google_search", "count": 3},
{"type": "url_context", "count": 2},
],
}
)
assert usage.prompt_tokens_details is not None
assert usage.prompt_tokens_details.web_search_requests == 3
def test_no_grounding_leaves_web_search_requests_unset():
usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
{
"total_input_tokens": 10,
"input_tokens_by_modality": [{"modality": "text", "tokens": 10}],
"total_output_tokens": 5,
}
)
assert usage.prompt_tokens_details is not None
assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None
def test_document_modality_folds_into_text():
usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
{

View file

@ -3514,6 +3514,42 @@ def test_completion_cost_bills_interactions_api_response():
assert cost > 0
def test_completion_cost_bills_interactions_google_search_per_query():
from litellm.types.interactions import InteractionsAPIResponse
model_info = litellm.get_model_info(model="gemini-3-flash-preview", custom_llm_provider="gemini")
response = InteractionsAPIResponse(
id="interactions/search123",
model="gemini-3-flash-preview",
status="completed",
steps=[],
usage={
"total_tokens": 680,
"total_input_tokens": 103,
"input_tokens_by_modality": [{"modality": "text", "tokens": 103}],
"total_cached_tokens": 0,
"total_output_tokens": 226,
"total_tool_use_tokens": 0,
"total_thought_tokens": 351,
"grounding_tool_count": [{"type": "google_search", "count": 3}],
},
)
cost = completion_cost(completion_response=response, custom_llm_provider="gemini")
per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"]
reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"]
expected = (
103 * model_info["input_cost_per_token"]
+ 226 * model_info["output_cost_per_token"]
+ 351 * reasoning_rate
+ 3 * per_query_cost
)
assert model_info.get("web_search_billing_unit") == "per_query"
assert cost == pytest.approx(expected)
assert cost > 3 * per_query_cost
def test_completion_cost_bills_interactions_video_output_at_video_rate():
from litellm.types.interactions import InteractionsAPIResponse