Add UsageLimits.count_tokens_before_request using Gemini count_tokens API (#2137)

kauabh · DouweM · web-flow · commit 229359552749 · 2025-08-13T16:16:20.000-06:00
Co-authored-by: Douwe Maan &lt;douwe@pydantic.dev&gt;
diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py
@@ -351,11 +351,6 @@ async def _prepare_request(
     ) -> tuple[ModelSettings | None, models.ModelRequestParameters, list[_messages.ModelMessage], RunContext[DepsT]]:
         ctx.state.message_history.append(self.request)
 
-        # Check usage
-        if ctx.deps.usage_limits:  # pragma: no branch
-            ctx.deps.usage_limits.check_before_request(ctx.state.usage)
-
-        # Increment run_step
         ctx.state.run_step += 1
 
         run_context = build_run_context(ctx)
@@ -367,6 +362,18 @@ async def _prepare_request(
 
         message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
 
+        usage = ctx.state.usage
+        if ctx.deps.usage_limits.count_tokens_before_request:
+            # Copy to avoid modifying the original usage object with the counted usage
+            usage = dataclasses.replace(usage)
+
+            counted_usage = await ctx.deps.model.count_tokens(
+                message_history, ctx.deps.model_settings, model_request_parameters
+            )
+            usage.incr(counted_usage)
+
+        ctx.deps.usage_limits.check_before_request(usage)
+
         return model_settings, model_request_parameters, message_history, run_context
 
     def _finish_handling(
diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py
@@ -413,6 +413,16 @@ async def request(
         """Make a request to the model."""
         raise NotImplementedError()
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> Usage:
+        """Make a request to the model for counting tokens."""
+        # This method is not required, but you need to implement it if you want to support `UsageLimits.count_tokens_before_request`.
+        raise NotImplementedError(f'Token counting ahead of the request is not supported by {self.__class__.__name__}')
+
     @asynccontextmanager
     async def request_stream(
         self,
diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py
@@ -52,13 +52,15 @@
     from google.genai.types import (
         ContentDict,
         ContentUnionDict,
+        CountTokensConfigDict,
         ExecutableCodeDict,
         FunctionCallDict,
         FunctionCallingConfigDict,
         FunctionCallingConfigMode,
         FunctionDeclarationDict,
         GenerateContentConfigDict,
         GenerateContentResponse,
+        GenerationConfigDict,
         GoogleSearchDict,
         HttpOptionsDict,
         MediaResolution,
@@ -188,6 +190,59 @@ async def request(
         response = await self._generate_content(messages, False, model_settings, model_request_parameters)
         return self._process_response(response)
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> usage.Usage:
+        check_allow_model_requests()
+        model_settings = cast(GoogleModelSettings, model_settings or {})
+        contents, generation_config = await self._build_content_and_config(
+            messages, model_settings, model_request_parameters
+        )
+
+        # Annoyingly, the type of `GenerateContentConfigDict.get` is "partially `Unknown`" because `response_schema` includes `typing._UnionGenericAlias`,
+        # so without this we'd need `pyright: ignore[reportUnknownMemberType]` on every line and wouldn't get type checking anyway.
+        generation_config = cast(dict[str, Any], generation_config)
+
+        config = CountTokensConfigDict(
+            http_options=generation_config.get('http_options'),
+        )
+        if self.system != 'google-gla':
+            # The fields are not supported by the Gemini API per https://github.com/googleapis/python-genai/blob/7e4ec284dc6e521949626f3ed54028163ef9121d/google/genai/models.py#L1195-L1214
+            config.update(
+                system_instruction=generation_config.get('system_instruction'),
+                tools=cast(list[ToolDict], generation_config.get('tools')),
+                # Annoyingly, GenerationConfigDict has fewer fields than GenerateContentConfigDict, and no extra fields are allowed.
+                generation_config=GenerationConfigDict(
+                    temperature=generation_config.get('temperature'),
+                    top_p=generation_config.get('top_p'),
+                    max_output_tokens=generation_config.get('max_output_tokens'),
+                    stop_sequences=generation_config.get('stop_sequences'),
+                    presence_penalty=generation_config.get('presence_penalty'),
+                    frequency_penalty=generation_config.get('frequency_penalty'),
+                    thinking_config=generation_config.get('thinking_config'),
+                    media_resolution=generation_config.get('media_resolution'),
+                    response_mime_type=generation_config.get('response_mime_type'),
+                    response_schema=generation_config.get('response_schema'),
+                ),
+            )
+
+        response = await self.client.aio.models.count_tokens(
+            model=self._model_name,
+            contents=contents,
+            config=config,
+        )
+        if response.total_tokens is None:
+            raise UnexpectedModelBehavior(  # pragma: no cover
+                'Total tokens missing from Gemini response', str(response)
+            )
+        return usage.Usage(
+            request_tokens=response.total_tokens,
+            total_tokens=response.total_tokens,
+        )
+
     @asynccontextmanager
     async def request_stream(
         self,
@@ -265,16 +320,23 @@ async def _generate_content(
         model_settings: GoogleModelSettings,
         model_request_parameters: ModelRequestParameters,
     ) -> GenerateContentResponse | Awaitable[AsyncIterator[GenerateContentResponse]]:
-        tools = self._get_tools(model_request_parameters)
+        contents, config = await self._build_content_and_config(messages, model_settings, model_request_parameters)
+        func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content
+        return await func(model=self._model_name, contents=contents, config=config)  # type: ignore
 
+    async def _build_content_and_config(
+        self,
+        messages: list[ModelMessage],
+        model_settings: GoogleModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[list[ContentUnionDict], GenerateContentConfigDict]:
+        tools = self._get_tools(model_request_parameters)
         response_mime_type = None
         response_schema = None
         if model_request_parameters.output_mode == 'native':
             if tools:
                 raise UserError('Gemini does not support structured output and tools at the same time.')
-
             response_mime_type = 'application/json'
-
             output_object = model_request_parameters.output_object
             assert output_object is not None
             response_schema = self._map_response_schema(output_object)
@@ -311,9 +373,7 @@ async def _generate_content(
             response_mime_type=response_mime_type,
             response_schema=response_schema,
         )
-
-        func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content
-        return await func(model=self._model_name, contents=contents, config=config)  # type: ignore
+        return contents, config
 
     def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         if not response.candidates or len(response.candidates) != 1:
diff --git a/pydantic_ai_slim/pydantic_ai/usage.py b/pydantic_ai_slim/pydantic_ai/usage.py
@@ -96,6 +96,10 @@ class UsageLimits:
     """The maximum number of tokens allowed in responses from the model."""
     total_tokens_limit: int | None = None
     """The maximum number of tokens allowed in requests and responses combined."""
+    count_tokens_before_request: bool = False
+    """If True, perform a token counting pass before sending the request to the model,
+    to enforce `request_tokens_limit` ahead of time. This may incur additional overhead
+    (from calling the model's `count_tokens` API before making the actual request) and is disabled by default."""
 
     def has_token_limits(self) -> bool:
         """Returns `True` if this instance places any limits on token counts.
@@ -111,11 +115,23 @@ def has_token_limits(self) -> bool:
         )
 
     def check_before_request(self, usage: Usage) -> None:
-        """Raises a `UsageLimitExceeded` exception if the next request would exceed the request_limit."""
+        """Raises a `UsageLimitExceeded` exception if the next request would exceed any of the limits."""
         request_limit = self.request_limit
         if request_limit is not None and usage.requests >= request_limit:
             raise UsageLimitExceeded(f'The next request would exceed the request_limit of {request_limit}')
 
+        request_tokens = usage.request_tokens or 0
+        if self.request_tokens_limit is not None and request_tokens > self.request_tokens_limit:
+            raise UsageLimitExceeded(
+                f'The next request would exceed the request_tokens_limit of {self.request_tokens_limit} ({request_tokens=})'
+            )
+
+        total_tokens = usage.total_tokens or 0
+        if self.total_tokens_limit is not None and total_tokens > self.total_tokens_limit:
+            raise UsageLimitExceeded(
+                f'The next request would exceed the total_tokens_limit of {self.total_tokens_limit} ({total_tokens=})'
+            )
+
     def check_tokens(self, usage: Usage) -> None:
         """Raises a `UsageLimitExceeded` exception if the usage exceeds any of the token limits."""
         request_tokens = usage.request_tokens or 0
diff --git a/tests/models/cassettes/test_google/test_google_model_usage_limit_exceeded.yaml b/tests/models/cassettes/test_google/test_google_model_usage_limit_exceeded.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"contents": [{"parts": [{"text": "The quick brown fox jumps over the lazydog."}],
+      "role": "user"}]}'
+    headers:
+      Content-Type:
+      - application/json
+      user-agent:
+      - google-genai-sdk/1.26.0 gl-python/3.12.7
+      x-goog-api-client:
+      - google-genai-sdk/1.26.0 gl-python/3.12.7
+    method: post
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:countTokens
+  response:
+    body:
+      string: "{\n  \"totalTokens\": 12,\n  \"promptTokensDetails\": [\n    {\n      \"modality\":
+        \"TEXT\",\n      \"tokenCount\": 12\n    }\n  ]\n}\n"
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Type:
+      - application/json; charset=UTF-8
+      Date:
+      - Fri, 01 Aug 2025 15:59:25 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=1582
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+      content-length:
+      - '117'
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/models/cassettes/test_google/test_google_model_usage_limit_not_exceeded.yaml b/tests/models/cassettes/test_google/test_google_model_usage_limit_not_exceeded.yaml
@@ -0,0 +1,115 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '100'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: The quick brown fox jumps over the lazydog.
+        role: user
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:countTokens
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '117'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=191
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      promptTokensDetails:
+      - modality: TEXT
+        tokenCount: 12
+      totalTokens: 12
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '124'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: The quick brown fox jumps over the lazydog.
+        role: user
+      generationConfig: {}
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '979'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=4808
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - content:
+          parts:
+          - text: |-
+              That's a classic! It's famously known as a **pangram**, which means it's a sentence that contains every letter of the alphabet.
+
+              It's often used for:
+              *   **Typing practice:** To ensure all keys are hit.
+              *   **Displaying font samples:** Because it showcases every character.
+
+              Just a small note, it's typically written as "lazy dog" (two words) and usually ends with a period:
+
+              **The quick brown fox jumps over the lazy dog.**
+          role: model
+        finishReason: STOP
+        index: 0
+      modelVersion: gemini-2.5-flash
+      responseId: ZwudaISALoquqtsP9uCG6Qw
+      usageMetadata:
+        candidatesTokenCount: 109
+        promptTokenCount: 12
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 12
+        thoughtsTokenCount: 806
+        totalTokenCount: 927
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/models/cassettes/test_google/test_google_vertexai_model_usage_limit_exceeded.yaml b/tests/models/cassettes/test_google/test_google_vertexai_model_usage_limit_exceeded.yaml
diff --git a/tests/models/test_google.py b/tests/models/test_google.py
