fix(google): count tool-use prompt tokens as input tokens

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Bug fixes
 
+* `ChatGoogle()` and `ChatVertex()` now count the tokens of content fetched by built-in tools (URL context, search, ...) as input tokens in `Turn.tokens` and the cost estimate. Gemini reports them as `tool_use_prompt_token_count`, separately from `prompt_token_count`, and they were left out.
 * `ChatOllama()` can now be created without a `model`, so `ChatOllama().list_models()` works. Sending a message without a model still raises an error that lists the available models. (#166)
 * `ChatGoogle()` now reports `finish_reason == "content_filter"` for Gemini's `IMAGE_SAFETY`, `IMAGE_PROHIBITED_CONTENT`, `IMAGE_RECITATION` and `LANGUAGE` finish reasons, instead of passing the raw Gemini string through, and reports `"tool_use"` (rather than `"success"`) when a turn ends with a tool request, as `ChatOpenAI()` does. (#473)
 * The error raised for a tool request with invalid JSON arguments from `ChatOpenAI()` and `ChatOpenAICompletions()` (and the providers built on them) now shows the offending arguments, instead of the literal text `{args}`.

chatlas/_provider_google.py

@@ -565,8 +565,12 @@ def value_tokens(self, completion):
         if usage is None:
             return None
         cached = usage.cached_content_token_count or 0
+        # Content fetched by built-in tools (URL context, search, ...) is
+        # counted as input tokens, but reported separately from the prompt.
         return (
-            (usage.prompt_token_count or 0) - cached,
+            (usage.prompt_token_count or 0)
+            + (usage.tool_use_prompt_token_count or 0)
+            - cached,
             (usage.candidates_token_count or 0) + (usage.thoughts_token_count or 0),
             usage.cached_content_token_count or 0,
         )

tests/test_provider_google.py

@@ -127,6 +127,47 @@ def test_google_stream_turn_tool_request_finishes_as_tool_use():
     assert turn.finish_reason == "tool_use"
 
 
+def test_google_value_tokens_counts_tool_use_prompt_tokens_as_input():
+    from google.genai.types import (
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+    )
+
+    provider = GoogleProvider(
+        model="gemini-2.5-flash", api_key="dummy", name="Google/Gemini", kwargs=None
+    )
+    completion = GenerateContentResponse(
+        usage_metadata=GenerateContentResponseUsageMetadata(
+            prompt_token_count=27,
+            tool_use_prompt_token_count=10309,
+            candidates_token_count=40,
+            thoughts_token_count=200,
+        )
+    )
+
+    assert provider.value_tokens(completion) == (10336, 240, 0)
+
+
+def test_google_value_tokens_without_tool_use_prompt_tokens():
+    from google.genai.types import (
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+    )
+
+    provider = GoogleProvider(
+        model="gemini-2.5-flash", api_key="dummy", name="Google/Gemini", kwargs=None
+    )
+    completion = GenerateContentResponse(
+        usage_metadata=GenerateContentResponseUsageMetadata(
+            prompt_token_count=100,
+            cached_content_token_count=30,
+            candidates_token_count=5,
+        )
+    )
+
+    assert provider.value_tokens(completion) == (70, 5, 30)
+
+
 def test_normalize_finish_reason_passes_through_unknown():
     assert google_normalize_finish_reason("OTHER") == "OTHER"
     assert google_normalize_finish_reason("NEW_REASON") == "NEW_REASON"