Add UsageLimits.count_tokens_before_request using Gemini count_tokens API

[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.

[kauabh]

@DouweM make sense, let me rework on this. Thanks for detailed input, appreciate your time

[kauabh]

Hey @DouweM I have made changes as per comments, looks like quite a few files got touched, It will would be great if you can provide some feedback on the changes till now. Also if you can share some thoughts on changing "instrumented.py" with count_tokens

[DouweM]

@kauabh Thanks! We're almost there :)

[DouweM]

@kauabh Let me know when this is ready for another round of review! I'm seeing some test failures

[kauabh]

Hey @DouweM just building around test cases, will share soon

[kauabh]

Hey @DouweM could you please have look at changes till now, I made updates as per your comments. Also need your thought on the gemini count token not allowing system instruction and tools. I added details after your PR comment

[kauabh]

Hey @DouweM please share your thoughts

[kauabh]

Hey @DouweM I have made changes as per the comments, the issue is with Vertex ai test case, it works with cassette with i in local, but I tried both adding cassette and removing but getting same error (CI failing),

async def test_google_vertexai_model_usage_limit_exceeded(allow_model_requests: None, google_provider: GoogleProvider):
    model = GoogleModel('gemini-2.5-flash', provider='google-vertex')

    agent = Agent(model, system_prompt='You are a chatbot.')

    @agent.tool_plain
    async def get_user_country() -> str:
        return 'Mexico'

    with pytest.raises(UsageLimitExceeded, match='Exceeded the request_tokens_limit of 9 \\(request_tokens=36\\)'):
        await agent.run(
            'What is the largest city in the user country? Use the get_user_country tool and then your own world knowledge.',
            usage_limits=UsageLimits(request_tokens_limit=50, count_tokens_before_request=True),
        )
        

Cassette

interactions:
- request:
    body: '[REDACTED]'
    headers:
      Accept:
      - '*/*'
      Accept-Encoding:
      - gzip, deflate
      Connection:
      - keep-alive
      Content-Length:
      - '268'
      Content-Type:
      - application/x-www-form-urlencoded
      User-Agent:
      - python-requests/2.32.3
      x-goog-api-client:
      - gl-python/3.12.7 auth/2.38.0 cred-type/u
    method: POST
    uri: https://oauth2.googleapis.com/token
  response:
    body:
      string: "{\n  \"access_token\": \"[REDACTED]\",\n  \"expires_in\": 3599,\n  \"scope\":
        \"openid https://www.googleapis.com/auth/sqlservice.login https://www.googleapis.com/auth/cloud-platform
        https://www.googleapis.com/auth/userinfo.email\",\n  \"token_type\": \"Bearer\",\n
        \ \"id_token\": \"[REDACTED]\"\n}"
    headers:
      Alt-Svc:
      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
      Cache-Control:
      - no-cache, no-store, max-age=0, must-revalidate
      Content-Type:
      - application/json; charset=utf-8
      Date:
      - Thu, 07 Aug 2025 16:48:01 GMT
      Expires:
      - Mon, 01 Jan 1990 00:00:00 GMT
      Pragma:
      - no-cache
      Server:
      - scaffolding on HTTPServer2
      Transfer-Encoding:
      - chunked
      Vary:
      - Origin
      - X-Origin
      - Referer
      X-Content-Type-Options:
      - nosniff
      X-Frame-Options:
      - SAMEORIGIN
      X-XSS-Protection:
      - '0'
      content-length:
      - '1467'
    status:
      code: 200
      message: OK
- request:
    body: '{"contents": [{"parts": [{"text": "What is the largest city in the user
      country? Use the get_user_country tool and then your own world knowledge."}],
      "role": "user"}], "systemInstruction": {"parts": [{"text": "You are a chatbot."}],
      "role": "user"}, "tools": [{"functionDeclarations": [{"description": "", "name":
      "get_user_country", "parameters": {"additional_properties": false, "properties":
      {}, "type": "OBJECT"}}]}]}'
    headers:
      Content-Type:
      - application/json
      user-agent:
      - google-genai-sdk/1.25.0 gl-python/3.12.7
      x-goog-api-client:
      - google-genai-sdk/1.25.0 gl-python/3.12.7
    method: post
    uri: https://us-central1-aiplatform.googleapis.com/v1beta1/projects/[REDACTED]-[REDACTED]/locations/us-central1/publishers/google/models/gemini-2.5-flash:countTokens
  response:
    body:
      string: "{\n  \"totalTokens\": 36,\n  \"totalBillableCharacters\": 123,\n  \"promptTokensDetails\":
        [\n    {\n      \"modality\": \"TEXT\",\n      \"tokenCount\": 36\n    }\n
        \ ]\n}\n"
    headers:
      Alt-Svc:
      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
      Content-Type:
      - application/json; charset=UTF-8
      Date:
      - Thu, 07 Aug 2025 16:48:02 GMT
      Server:
      - scaffolding on HTTPServer2
      Transfer-Encoding:
      - chunked
      Vary:
      - Origin
      - X-Origin
      - Referer
      X-Content-Type-Options:
      - nosniff
      X-Frame-Options:
      - SAMEORIGIN
      X-XSS-Protection:
      - '0'
      content-length:
      - '151'
    status:
      code: 200
      message: OK
version: 1

[kauabh]

Hey @DouweM I have made changes as per the comments, the issue is with Vertex ai test case, it works with cassette with i in local, but I tried both adding cassette and removing but getting same error (CI failing),

async def test_google_vertexai_model_usage_limit_exceeded(allow_model_requests: None, google_provider: GoogleProvider):

    model = GoogleModel('gemini-2.5-flash', provider='google-vertex')



    agent = Agent(model, system_prompt='You are a chatbot.')



    @agent.tool_plain

    async def get_user_country() -> str:

        return 'Mexico'



    with pytest.raises(UsageLimitExceeded, match='Exceeded the request_tokens_limit of 9 \\(request_tokens=36\\)'):

        await agent.run(

            'What is the largest city in the user country? Use the get_user_country tool and then your own world knowledge.',

            usage_limits=UsageLimits(request_tokens_limit=50, count_tokens_before_request=True),

        )

        

Cassette

interactions:

- request:

    body: '[REDACTED]'

    headers:

      Accept:

      - '*/*'

      Accept-Encoding:

      - gzip, deflate

      Connection:

      - keep-alive

      Content-Length:

      - '268'

      Content-Type:

      - application/x-www-form-urlencoded

      User-Agent:

      - python-requests/2.32.3

      x-goog-api-client:

      - gl-python/3.12.7 auth/2.38.0 cred-type/u

    method: POST

    uri: https://oauth2.googleapis.com/token

  response:

    body:

      string: "{\n  \"access_token\": \"[REDACTED]\",\n  \"expires_in\": 3599,\n  \"scope\":

        \"openid https://www.googleapis.com/auth/sqlservice.login https://www.googleapis.com/auth/cloud-platform

        https://www.googleapis.com/auth/userinfo.email\",\n  \"token_type\": \"Bearer\",\n

        \ \"id_token\": \"[REDACTED]\"\n}"

    headers:

      Alt-Svc:

      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000

      Cache-Control:

      - no-cache, no-store, max-age=0, must-revalidate

      Content-Type:

      - application/json; charset=utf-8

      Date:

      - Thu, 07 Aug 2025 16:48:01 GMT

      Expires:

      - Mon, 01 Jan 1990 00:00:00 GMT

      Pragma:

      - no-cache

      Server:

      - scaffolding on HTTPServer2

      Transfer-Encoding:

      - chunked

      Vary:

      - Origin

      - X-Origin

      - Referer

      X-Content-Type-Options:

      - nosniff

      X-Frame-Options:

      - SAMEORIGIN

      X-XSS-Protection:

      - '0'

      content-length:

      - '1467'

    status:

      code: 200

      message: OK

- request:

    body: '{"contents": [{"parts": [{"text": "What is the largest city in the user

      country? Use the get_user_country tool and then your own world knowledge."}],

      "role": "user"}], "systemInstruction": {"parts": [{"text": "You are a chatbot."}],

      "role": "user"}, "tools": [{"functionDeclarations": [{"description": "", "name":

      "get_user_country", "parameters": {"additional_properties": false, "properties":

      {}, "type": "OBJECT"}}]}]}'

    headers:

      Content-Type:

      - application/json

      user-agent:

      - google-genai-sdk/1.25.0 gl-python/3.12.7

      x-goog-api-client:

      - google-genai-sdk/1.25.0 gl-python/3.12.7

    method: post

    uri: https://us-central1-aiplatform.googleapis.com/v1beta1/projects/[REDACTED]-[REDACTED]/locations/us-central1/publishers/google/models/gemini-2.5-flash:countTokens

  response:

    body:

      string: "{\n  \"totalTokens\": 36,\n  \"totalBillableCharacters\": 123,\n  \"promptTokensDetails\":

        [\n    {\n      \"modality\": \"TEXT\",\n      \"tokenCount\": 36\n    }\n

        \ ]\n}\n"

    headers:

      Alt-Svc:

      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000

      Content-Type:

      - application/json; charset=UTF-8

      Date:

      - Thu, 07 Aug 2025 16:48:02 GMT

      Server:

      - scaffolding on HTTPServer2

      Transfer-Encoding:

      - chunked

      Vary:

      - Origin

      - X-Origin

      - Referer

      X-Content-Type-Options:

      - nosniff

      X-Frame-Options:

      - SAMEORIGIN

      X-XSS-Protection:

      - '0'

      content-length:

      - '151'

    status:

      code: 200

      message: OK

version: 1

Hey @DouweM please share your thoughts on changes and also on VCR for this VERTEXAI limit test, not able to run it in CI

[DouweM]

@kauabh I was able to fix it up, thanks a lot!

[kauabh]

Hey @DouweM thank you so much, just wanted to ask what mistake I was making in creating that VERTEX AI cassette

[DouweM]

@kauabh Vertex is a bit tricky, it took me some time to make it work as well.

I had to comment out these lines:

pydantic-ai/tests/conftest.py

Lines 440 to 441 in eeeb32d

	if not os.getenv('CI', False):
	pytest.skip('Requires properly configured local google vertex config to pass')

, and make sure I authenticate with a service account just like happens in CI, so that the VCR requests are exactly the same.