diff --git a/CHANGELOG.md b/CHANGELOG.md index 040d34c8..ef79d967 100644 --- a/CHANGELOG.md +++ b/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}`. diff --git a/chatlas/_provider_google.py b/chatlas/_provider_google.py index 4c750a76..0e2a59d4 100644 --- a/chatlas/_provider_google.py +++ b/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, ) diff --git a/tests/test_provider_google.py b/tests/test_provider_google.py index b57e06ad..7bee3544 100644 --- a/tests/test_provider_google.py +++ b/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"