feat: add GoogleGenAITokenCounter to google-genai integration

[pma1999]

Related Issues

• fixes feat: Add GoogleGenAITokenCounter to google-genai integration #3857

Proposed Changes:

Adds GoogleGenAITokenCounter to the google_genai integration, implementing the TokenCounter protocol introduced in Haystack 3.1 on top of Google's countTokens endpoint — the Gemini analogue of OpenAITokenCounter.

It follows the OpenAITokenCounter pattern: model is required and never assumed, everything after it is keyword-only, the client is built lazily in warm_up(), count() auto-warms, and there is a close() lifecycle method. The constructor mirrors GoogleGenAIChatGenerator (api, vertex_ai_project, vertex_ai_location, timeout, max_retries) so a counter can be configured exactly like the generator it measures.

Inputs are assembled the way GoogleGenAIChatGenerator sends them: a leading system message becomes the system instruction, the remaining messages become the request contents.

from haystack.dataclasses import ChatMessage
from haystack_integrations.token_counters.google_genai import GoogleGenAITokenCounter

counter = GoogleGenAITokenCounter("gemini-3.7-flash")
print(counter.count([ChatMessage.from_user("Hello, how are you?")]))

One design decision I'd like your call on

The Google Gen AI SDK cannot measure a system instruction or tool schemas on the Gemini Developer API. It is not a silent drop — _CountTokensConfig_to_mldev raises explicitly:

if getv(from_object, ['tools']) is not None:
    raise ValueError(
        'tools parameter is only supported in Gemini Enterprise Agent Platform'
        ' mode, not in Gemini Developer API mode.'
    )

_CountTokensConfig_to_vertex maps both, so Vertex is fine. Worth noting the REST API itself could do it via generateContentRequest, which accepts system instructions and function declarations — but the Python SDK does not expose that field anywhere, so there is no supported path to it today.

Since count(messages, tools) has to answer for both, I chose to raise a ValueError naming the limitation and pointing at api="vertex", rather than return a count that silently omits tool schemas — those can be worth hundreds of tokens, and a counter that quietly under-reports seemed worse than one that refuses. Counting plain messages works on either backend.

Happy to switch to a warning plus a partial count if you'd rather the component always return a number.

How did you test it?

• hatch run test:unit — 9 new unit tests pass. They cover serde round-trip, the empty-input short circuit, the total_tokens: None case, the rejection path for both a system message and tools, that Vertex actually forwards the system instruction and tool schemas, and close().
• hatch run test:integration — the live test passes against the real Gemini API.
• hatch run test:types — mypy clean (the new package is registered in the types command).
• hatch run fmt-check — clean.

Two failures in the full unit run (test_init_fail_wo_api_key, test_extract_sources_info_rejects_symlink_escaping_root) are unrelated to this change: they reproduce identically on a clean checkout of main on my machine — the first because I have GEMINI_API_KEY exported and the test only clears GOOGLE_API_KEY, the second because of Windows symlink permissions.

The live test deliberately covers a user turn, an assistant tool call and a tool result, but not a system message or tools, since the Gemini Developer API these tests authenticate against cannot measure those.

[CLAassistant]

All committers have signed the CLA.

[github-actions]

Heads-up for maintainers

This PR is from a fork and touches integrations whose integration tests require API keys.
Those tests are skipped in CI because fork PRs don't have access to repo secrets for security reasons.

Affected integrations:

• google_genai

Please run the integration tests locally (hatch run test:integration inside each folder) before approving.

[github-actions]

Coverage report (google_genai)

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
integrations/google_genai/src/haystack_integrations/token_counters/google_genai
token_counter.py					91, 142-143
Project Total

This report was generated by python-coverage-comment-action