Adding CountToken to Gemini

[kauabh]

Gemini Provides an endpoint to count tokens https://ai.google.dev/api/tokens#method:-models.counttokens.
I think it will be useful and address some concerns in this issue #1794 (at least for gemini).

@DouweM Wanted to check if this will be helpful. If yes and if the approach is right, wanted to know if you can share some pointers around adding it in usage_limits for gemini. Happy to work on other models too, if this one make it through.

[DouweM]

@kauabh I agree that if a model API has a method to count tokens, it would be nice to expose that on the Model class.

But I don't think we should automatically use it when UsageLimits(request_tokens_limit=...) is used, as it adds an extra request and the overhead and latency that comes with that, unlike OpenAI's tiktoken which was mentioned in #1794 and can be run locally. So if we'd like to give users the option to better enforce request_tokens_limit by doing a separate count-tokens request ahead of the actual LLM request, that should be opt-in with some flag on UsageLimits and appropriate warnings in the docs about the extra overhead.

That check would need to be implemented here, just before we call model.request, once we have the messages, model settings, and model request params ready:

pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py

Lines 379 to 393 in b31c77d

	async def _make_request(
	self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
	) -> CallToolsNode[DepsT, NodeRunEndT]:
	if self._result is not None:
	return self._result # pragma: no cover
	model_settings, model_request_parameters = await self._prepare_request(ctx)
	model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
	message_history = await _process_message_history(
	ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
	)
	model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
	ctx.state.usage.incr(_usage.Usage())
	return self._finish_handling(ctx, model_response)

This would require a method that exists on every model, so it'd be implemented as an abstract method on the base Model class with a default implementation of raise NotImplementedError(...), and only models that have a count-tokens method would override it with a concrete implementation.

As for that concrete implementation, I recommend adding it to GoogleModel instead of GeminiModel, as you can directly use the google-genai library there, and reducing the duplication with the request-preparation logic in _generate_content as much as possible.