add token counting

Author: wirthual

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -14,6 +14,7 @@
 from types import GenericAlias
 from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeGuard, TypeVar, get_args, get_origin, overload
 
+import tiktoken
 from anyio.to_thread import run_sync
 from pydantic import BaseModel, TypeAdapter
 from pydantic.json_schema import JsonSchemaValue
@@ -32,10 +33,14 @@
 AbstractSpan = AbstractSpan
 
 if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletionMessageParam
+    from openai.types.responses.response_input_item_param import ResponseInputItemParam
+
     from pydantic_ai.agent import AgentRun, AgentRunResult
     from pydantic_graph import GraphRun, GraphRunResult
 
     from . import messages as _messages
+    from .models.openai import OpenAIModelName
     from .tools import ObjectJsonSchema
 
 _P = ParamSpec('_P')
@@ -507,3 +512,48 @@ def get_event_loop():
         event_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(event_loop)
     return event_loop
+
+
+def num_tokens_from_messages(
+    messages: list[ChatCompletionMessageParam] | list[ResponseInputItemParam],
+    model: OpenAIModelName = 'gpt-4o-mini-2024-07-18',
+) -> int:
+    """Return the number of tokens used by a list of messages."""
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+    except KeyError:
+        print('Warning: model not found. Using o200k_base encoding.')  # TODO: How to handle warnings?
+        encoding = tiktoken.get_encoding('o200k_base')
+    if model in {
+        'gpt-3.5-turbo-0125',
+        'gpt-4-0314',
+        'gpt-4-32k-0314',
+        'gpt-4-0613',
+        'gpt-4-32k-0613',
+        'gpt-4o-mini-2024-07-18',
+        'gpt-4o-2024-08-06',
+    }:
+        tokens_per_message = 3
+        tokens_per_name = 1
+    elif 'gpt-3.5-turbo' in model:
+        return num_tokens_from_messages(messages, model='gpt-3.5-turbo-0125')
+    elif 'gpt-4o-mini' in model:
+        return num_tokens_from_messages(messages, model='gpt-4o-mini-2024-07-18')
+    elif 'gpt-4o' in model:
+        return num_tokens_from_messages(messages, model='gpt-4o-2024-08-06')
+    elif 'gpt-4' in model:
+        return num_tokens_from_messages(messages, model='gpt-4-0613')
+    else:
+        raise NotImplementedError(
+            f"""num_tokens_from_messages() is not implemented for model {model}."""
+        )  # TODO: How to handle other models?
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            if isinstance(value, str):
+                num_tokens += len(encoding.encode(value))
+            if key == 'name':
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
+    return num_tokens

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -17,7 +17,12 @@
 from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
 from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
-from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
+from .._utils import (
+    guard_tool_call_id as _guard_tool_call_id,
+    now_utc as _now_utc,
+    num_tokens_from_messages,
+    number_to_datetime,
+)
 from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, MCPServerTool, WebSearchTool
 from ..exceptions import UserError
 from ..messages import (
@@ -907,6 +912,20 @@ def _inline_text_file_part(text: str, *, media_type: str, identifier: str) -> Ch
         )
         return ChatCompletionContentPartTextParam(text=text, type='text')
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> usage.RequestUsage:
+        """Make a request to the model for counting tokens."""
+        openai_messages = await self._map_messages(messages, model_request_parameters)
+        token_count = num_tokens_from_messages(openai_messages, self.model_name)
+
+        return usage.RequestUsage(
+            input_tokens=token_count,
+        )
+
 
 @deprecated(
     '`OpenAIModel` was renamed to `OpenAIChatModel` to clearly distinguish it from `OpenAIResponsesModel` which '
@@ -1701,6 +1720,22 @@ async def _map_user_prompt(part: UserPromptPart) -> responses.EasyInputMessagePa
                     assert_never(item)
         return responses.EasyInputMessageParam(role='user', content=content)
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> usage.RequestUsage:
+        """Make a request to the model for counting tokens."""
+        _, openai_messages = await self._map_messages(
+            messages, cast(OpenAIResponsesModelSettings, model_settings or {}), model_request_parameters
+        )
+        token_count = num_tokens_from_messages(openai_messages, self.model_name)
+
+        return usage.RequestUsage(
+            input_tokens=token_count,
+        )
+
 
 @dataclass
 class OpenAIStreamedResponse(StreamedResponse):

pydantic_ai_slim/pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
 # WARNING if you add optional groups, please update docs/install.md
 logfire = ["logfire[httpx]>=3.14.1"]
 # Models
-openai = ["openai>=1.107.2"]
+openai = ["openai>=1.107.2","tiktoken>=0.12.0"]
 cohere = ["cohere>=5.18.0; platform_system != 'Emscripten'"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
 google = ["google-genai>=1.50.1"]

tests/models/test_openai.py

@@ -21,6 +21,7 @@
     DocumentUrl,
     ImageUrl,
     ModelHTTPError,
+    ModelMessage,
     ModelProfile,
     ModelRequest,
     ModelResponse,
@@ -3085,3 +3086,57 @@ async def test_cache_point_filtering_responses_model():
     assert len(msg['content']) == 2
     assert msg['content'][0]['text'] == 'text before'  # type: ignore[reportUnknownArgumentType]
     assert msg['content'][1]['text'] == 'text after'  # type: ignore[reportUnknownArgumentType]
+
+
+@pytest.mark.vcr()
+@pytest.mark.parametrize(
+    'model_name,expected_token_count',
+    [
+        ('gpt-3.5-turbo', 115),
+        ('gpt-4-0613', 115),
+        ('gpt-4', 115),
+        ('gpt-4o', 110),
+        ('gpt-4o-mini', 110),
+    ],
+)
+async def test_count_tokens(model_name: str, expected_token_count: int):
+    """Test token counting with OpenAI Chat and Response models."""
+    test_messages: list[ModelMessage] = [
+        ModelRequest(
+            parts=[
+                SystemPromptPart(
+                    content='You are a helpful, pattern-following assistant that translates corporate jargon into plain English.',
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                SystemPromptPart(
+                    content='New synergies will help drive top-line growth.',
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                SystemPromptPart(
+                    content='Things working well together will increase revenue.',
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                SystemPromptPart(
+                    content="Let's circle back when we have more bandwidth to touch base on opportunities for increased leverage.",
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                SystemPromptPart(
+                    content="Let's talk later when we're less busy about how to do better.",
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                UserPromptPart(
+                    content="This late pivot means we don't have time to boil the ocean for the client deliverable.",
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+            ],
+            run_id=IsStr(),
+        )
+    ]
+
+    chat_model = OpenAIChatModel(model_name, provider=OpenAIProvider(api_key='foobar'))
+    usage_result: RequestUsage = await chat_model.count_tokens(test_messages, {}, ModelRequestParameters())
+    assert usage_result.input_tokens == expected_token_count
+
+    responses_model = OpenAIResponsesModel(model_name, provider=OpenAIProvider(api_key='foobar'))
+    usage_result: RequestUsage = await responses_model.count_tokens(test_messages, {}, ModelRequestParameters())
+    assert usage_result.input_tokens == expected_token_count