mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(interactions): bill google_search grounding queries per query
This commit is contained in:
parent
194dca7dd0
commit
e458aa1230
3 changed files with 76 additions and 1 deletions
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue