pydantic-ai 0.1.10__tar.gz → 0.1.12__tar.gz

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.1.10
+Version: 0.1.12
 Summary: Agent Framework / shim to use Pydantic with LLMs
 Project-URL: Homepage, https://ai.pydantic.dev
 Project-URL: Source, https://github.com/pydantic/pydantic-ai
@@ -28,9 +28,9 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,evals,groq,mcp,mistral,openai,vertexai]==0.1.10
+Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,evals,groq,mcp,mistral,openai,vertexai]==0.1.12
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.1.10; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.1.12; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/pyproject.toml

@@ -60,14 +60,13 @@ Documentation = "https://ai.pydantic.dev"
 Changelog = "https://github.com/pydantic/pydantic-ai/releases"
 
 [project.scripts]
-pai = "pydantic_ai._cli:app"
+pai = "pydantic_ai._cli:cli_exit"  # TODO remove this when clai has been out for a while
 
 [tool.uv.sources]
 pydantic-ai-slim = { workspace = true }
 pydantic-evals = { workspace = true }
 pydantic-graph = { workspace = true }
 pydantic-ai-examples = { workspace = true }
-mcp-run-python = { workspace = true }
 
 [tool.uv.workspace]
 members = [
@@ -75,6 +74,7 @@ members = [
     "pydantic_evals",
     "pydantic_graph",
     "mcp-run-python",
+    "clai",
     "examples",
 ]
 
@@ -106,6 +106,7 @@ include = [
     "pydantic_graph/**/*.py",
     "mcp-run-python/**/*.py",
     "examples/**/*.py",
+    "clai/**/*.py",
     "tests/**/*.py",
     "docs/**/*.py",
 ]
@@ -114,7 +115,7 @@ include = [
 extend-select = [
     "Q",
     "RUF100",
-    "RUF018",  # https://docs.astral.sh/ruff/rules/assignment-in-assert/
+    "RUF018", # https://docs.astral.sh/ruff/rules/assignment-in-assert/
     "C90",
     "UP",
     "I",
@@ -163,6 +164,7 @@ include = [
     "mcp-run-python",
     "tests",
     "examples",
+    "clai",
 ]
 venvPath = ".venv"
 # see https://github.com/microsoft/pyright/issues/7771 - we don't want to error on decorated functions in tests
@@ -193,7 +195,7 @@ filterwarnings = [
     "ignore:websockets.server.WebSocketServerProtocol is deprecated:DeprecationWarning",
     # random resource warnings; I suspect these are coming from vendor SDKs when running examples..
     "ignore:unclosed <socket:ResourceWarning",
-    "ignore:unclosed event loop:ResourceWarning"
+    "ignore:unclosed event loop:ResourceWarning",
 ]
 
 # https://coverage.readthedocs.io/en/latest/config.html#run
@@ -215,7 +217,10 @@ show_missing = true
 ignore_errors = true
 precision = 2
 exclude_lines = [
+    # `# pragma: no cover` is standard marker for code that's not covered, this will error if code is covered
     'pragma: no cover',
+    # use `# pragma: lax no cover` if you want to ignore cases where (some of) the code is covered
+    'pragma: lax no cover',
     'raise NotImplementedError',
     'if TYPE_CHECKING:',
     'if typing.TYPE_CHECKING:',

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/conftest.py

@@ -24,7 +24,7 @@ import pydantic_ai.models
 from pydantic_ai.messages import BinaryContent
 from pydantic_ai.models import Model, cached_async_http_client
 
-__all__ = 'IsDatetime', 'IsFloat', 'IsNow', 'IsStr', 'TestEnv', 'ClientWithHandler', 'try_import'
+__all__ = 'IsDatetime', 'IsFloat', 'IsNow', 'IsStr', 'IsInt', 'TestEnv', 'ClientWithHandler', 'try_import'
 
 
 pydantic_ai.models.ALLOW_MODEL_REQUESTS = False
@@ -34,10 +34,11 @@ if TYPE_CHECKING:
 
     def IsDatetime(*args: Any, **kwargs: Any) -> datetime: ...
     def IsFloat(*args: Any, **kwargs: Any) -> float: ...
+    def IsInt(*args: Any, **kwargs: Any) -> int: ...
     def IsNow(*args: Any, **kwargs: Any) -> datetime: ...
     def IsStr(*args: Any, **kwargs: Any) -> str: ...
 else:
-    from dirty_equals import IsDatetime, IsFloat, IsNow as _IsNow, IsStr
+    from dirty_equals import IsDatetime, IsFloat, IsInt, IsNow as _IsNow, IsStr
 
     def IsNow(*args: Any, **kwargs: Any):
         # Increase the default value of `delta` to 10 to reduce test flakiness on overburdened machines
@@ -59,7 +60,7 @@ class TestEnv:
     def remove(self, name: str) -> None:
         self.envars[name] = os.environ.pop(name, None)
 
-    def reset(self) -> None:  # pragma: no cover
+    def reset(self) -> None:
         for name, value in self.envars.items():
             if value is None:
                 os.environ.pop(name, None)
@@ -100,7 +101,7 @@ async def client_with_handler() -> AsyncIterator[ClientWithHandler]:
     try:
         yield create_client
     finally:
-        if client:  # pragma: no cover
+        if client:
             await client.aclose()
 
 
@@ -155,7 +156,7 @@ def create_module(tmp_path: Path, request: pytest.FixtureRequest) -> Callable[[s
 
 
 @contextmanager
-def try_import() -> Iterator[Callable[[], bool]]:  # pragma: no cover
+def try_import() -> Iterator[Callable[[], bool]]:
     import_success = False
 
     def check_import() -> bool:
@@ -276,7 +277,7 @@ def mistral_api_key() -> str:
 
 
 @pytest.fixture(scope='session')
-def bedrock_provider():  # pragma: no cover
+def bedrock_provider():
     try:
         import boto3
 
@@ -304,7 +305,7 @@ def model(
     co_api_key: str,
     gemini_api_key: str,
     bedrock_provider: BedrockProvider,
-) -> Model:  # pragma: no cover
+) -> Model:  # pragma: lax no cover
     try:
         if request.param == 'openai':
             from pydantic_ai.models.openai import OpenAIModel

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/evals/test_evaluator_common.py

@@ -40,7 +40,7 @@ if TYPE_CHECKING or imports_successful():
             self.expected_output = expected_output
             self.inputs = inputs
             self.duration = duration
-else:  # pragma: no cover
+else:
     MockContext = object
 
 

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/evals/test_evaluators.py

@@ -10,7 +10,6 @@ from pydantic import BaseModel, TypeAdapter
 from pydantic_ai.messages import ModelMessage, ModelResponse
 from pydantic_ai.models import Model, ModelRequestParameters
 from pydantic_ai.settings import ModelSettings
-from pydantic_ai.usage import Usage
 
 from ..conftest import try_import
 
@@ -122,7 +121,7 @@ async def test_llm_judge_serialization():
             messages: list[ModelMessage],
             model_settings: ModelSettings | None,
             model_request_parameters: ModelRequestParameters,
-        ) -> tuple[ModelResponse, Usage]:
+        ) -> ModelResponse:
             raise NotImplementedError
 
         @property

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/evals/test_otel.py

@@ -879,7 +879,7 @@ async def test_context_subtree_invalid_tracer_provider(mocker: MockerFixture):
 
     mocker.patch('pydantic_evals.otel._context_in_memory_span_exporter.get_tracer_provider', return_value=None)
     with pytest.raises(TypeError) as exc_info:
-        with context_subtree():  # pragma: no cover
+        with context_subtree():
             pass
     assert str(exc_info.value) == snapshot(
         "Expected `tracer_provider` to have an `add_span_processor` method; got an instance of <class 'NoneType'>. For help resolving this, please create an issue at https://github.com/pydantic/pydantic-ai/issues."

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/graph/test_mermaid.py

@@ -259,7 +259,7 @@ def httpx_with_handler() -> Iterator[HttpxWithHandler]:
     try:
         yield create_client
     finally:
-        if client:  # pragma: no cover
+        if client:
             client.close()
 
 

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/models/test_anthropic.py

@@ -6,7 +6,7 @@ from collections.abc import Sequence
 from dataclasses import dataclass, field
 from datetime import timezone
 from functools import cached_property
-from typing import Any, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, Union, cast
 
 import httpx
 import pytest
@@ -52,7 +52,11 @@ with try_import() as imports_successful:
     )
     from anthropic.types.raw_message_delta_event import Delta
 
-    from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+    from pydantic_ai.models.anthropic import (
+        AnthropicModel,
+        AnthropicModelSettings,
+        _map_usage,  # pyright: ignore[reportPrivateUsage]
+    )
     from pydantic_ai.providers.anthropic import AnthropicProvider
 
     # note: we use Union here so that casting works with Python 3.9
@@ -169,12 +173,26 @@ async def test_sync_request_text_response(allow_model_requests: None):
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=5,
+                    response_tokens=10,
+                    total_tokens=15,
+                    details={'input_tokens': 5, 'output_tokens': 10},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=5,
+                    response_tokens=10,
+                    total_tokens=15,
+                    details={'input_tokens': 5, 'output_tokens': 10},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -258,6 +276,13 @@ async def test_request_structured_response(allow_model_requests: None):
                         tool_call_id='123',
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=3,
+                    response_tokens=5,
+                    total_tokens=8,
+                    details={'input_tokens': 3, 'output_tokens': 5},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -320,6 +345,13 @@ async def test_request_tool_call(allow_model_requests: None):
                         tool_call_id='1',
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=2,
+                    response_tokens=1,
+                    total_tokens=3,
+                    details={'input_tokens': 2, 'output_tokens': 1},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -341,6 +373,13 @@ async def test_request_tool_call(allow_model_requests: None):
                         tool_call_id='2',
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=3,
+                    response_tokens=2,
+                    total_tokens=5,
+                    details={'input_tokens': 3, 'output_tokens': 2},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -356,6 +395,13 @@ async def test_request_tool_call(allow_model_requests: None):
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=3,
+                    response_tokens=5,
+                    total_tokens=8,
+                    details={'input_tokens': 3, 'output_tokens': 5},
+                ),
                 model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -692,6 +738,18 @@ async def test_image_as_binary_content_tool_response(
                     TextPart(content='Let me get the image and check what fruit is shown.'),
                     ToolCallPart(tool_name='get_image', args={}, tool_call_id='toolu_01VMGXdexE1Fy5xdWgoom9Te'),
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=372,
+                    response_tokens=49,
+                    total_tokens=421,
+                    details={
+                        'cache_creation_input_tokens': 0,
+                        'cache_read_input_tokens': 0,
+                        'input_tokens': 372,
+                        'output_tokens': 49,
+                    },
+                ),
                 model_name='claude-3-5-sonnet-20241022',
                 timestamp=IsDatetime(),
             ),
@@ -718,6 +776,18 @@ async def test_image_as_binary_content_tool_response(
                         content="The image shows a kiwi fruit that has been cut in half, displaying its characteristic bright green flesh with small black seeds arranged in a circular pattern around a white center core. The kiwi's fuzzy brown skin is visible around the edges of the slice."
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=2025,
+                    response_tokens=57,
+                    total_tokens=2082,
+                    details={
+                        'cache_creation_input_tokens': 0,
+                        'cache_read_input_tokens': 0,
+                        'input_tokens': 2025,
+                        'output_tokens': 57,
+                    },
+                ),
                 model_name='claude-3-5-sonnet-20241022',
                 timestamp=IsDatetime(),
             ),
@@ -838,8 +908,91 @@ async def test_anthropic_model_instructions(allow_model_requests: None, anthropi
             ),
             ModelResponse(
                 parts=[TextPart(content='The capital of France is Paris.')],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=20,
+                    response_tokens=10,
+                    total_tokens=30,
+                    details={
+                        'cache_creation_input_tokens': 0,
+                        'cache_read_input_tokens': 0,
+                        'input_tokens': 20,
+                        'output_tokens': 10,
+                    },
+                ),
                 model_name='claude-3-opus-20240229',
                 timestamp=IsDatetime(),
             ),
         ]
     )
+
+
+def anth_msg(usage: AnthropicUsage) -> AnthropicMessage:
+    return AnthropicMessage(
+        id='x',
+        content=[],
+        model='claude-3-7-sonnet-latest',
+        role='assistant',
+        type='message',
+        usage=usage,
+    )
+
+
+@pytest.mark.parametrize(
+    'message_callback,usage',
+    [
+        pytest.param(
+            lambda: anth_msg(AnthropicUsage(input_tokens=1, output_tokens=1)),
+            snapshot(
+                Usage(
+                    request_tokens=1, response_tokens=1, total_tokens=2, details={'input_tokens': 1, 'output_tokens': 1}
+                )
+            ),
+            id='AnthropicMessage',
+        ),
+        pytest.param(
+            lambda: anth_msg(
+                AnthropicUsage(
+                    input_tokens=1, output_tokens=1, cache_creation_input_tokens=2, cache_read_input_tokens=3
+                )
+            ),
+            snapshot(
+                Usage(
+                    request_tokens=6,
+                    response_tokens=1,
+                    total_tokens=7,
+                    details={
+                        'cache_creation_input_tokens': 2,
+                        'cache_read_input_tokens': 3,
+                        'input_tokens': 1,
+                        'output_tokens': 1,
+                    },
+                )
+            ),
+            id='AnthropicMessage-cached',
+        ),
+        pytest.param(
+            lambda: RawMessageStartEvent(
+                message=anth_msg(AnthropicUsage(input_tokens=1, output_tokens=1)), type='message_start'
+            ),
+            snapshot(
+                Usage(
+                    request_tokens=1, response_tokens=1, total_tokens=2, details={'input_tokens': 1, 'output_tokens': 1}
+                )
+            ),
+            id='RawMessageStartEvent',
+        ),
+        pytest.param(
+            lambda: RawMessageDeltaEvent(
+                delta=Delta(),
+                usage=MessageDeltaUsage(output_tokens=5),
+                type='message_delta',
+            ),
+            snapshot(Usage(response_tokens=5, total_tokens=5, details={'output_tokens': 5})),
+            id='RawMessageDeltaEvent',
+        ),
+        pytest.param(lambda: RawMessageStopEvent(type='message_stop'), snapshot(Usage()), id='RawMessageStopEvent'),
+    ],
+)
+def test_usage(message_callback: Callable[[], AnthropicMessage | RawMessageStreamEvent], usage: Usage):
+    assert _map_usage(message_callback()) == usage

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/models/test_bedrock.py

@@ -75,6 +75,7 @@ async def test_bedrock_model(allow_model_requests: None, bedrock_provider: Bedro
                         content="Hello! How can I assist you today? Whether you have questions, need information, or just want to chat, I'm here to help."
                     )
                 ],
+                usage=Usage(requests=1, request_tokens=7, response_tokens=30, total_tokens=37),
                 model_name='us.amazon.nova-micro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -132,6 +133,7 @@ async def test_bedrock_model_structured_response(allow_model_requests: None, bed
                         tool_call_id='tooluse_5WEci1UmQ8ifMFkUcy2gHQ',
                     ),
                 ],
+                usage=Usage(requests=1, request_tokens=551, response_tokens=132, total_tokens=683),
                 model_name='us.amazon.nova-micro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -156,6 +158,7 @@ async def test_bedrock_model_structured_response(allow_model_requests: None, bed
                         tool_call_id='tooluse_9AjloJSaQDKmpPFff-2Clg',
                     ),
                 ],
+                usage=Usage(requests=1, request_tokens=685, response_tokens=166, total_tokens=851),
                 model_name='us.amazon.nova-micro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -255,6 +258,7 @@ async def test_bedrock_model_retry(allow_model_requests: None, bedrock_provider:
                         tool_call_id='tooluse_F8LnaCMtQ0-chKTnPhNH2g',
                     ),
                 ],
+                usage=Usage(requests=1, request_tokens=417, response_tokens=69, total_tokens=486),
                 model_name='us.amazon.nova-micro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -278,6 +282,7 @@ I'm sorry, but the tool I have does not support retrieving the capital of France
 """
                     )
                 ],
+                usage=Usage(requests=1, request_tokens=509, response_tokens=108, total_tokens=617),
                 model_name='us.amazon.nova-micro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -358,7 +363,7 @@ async def test_bedrock_model_iter_stream(allow_model_requests: None, bedrock_pro
         Args:
             city: The city name.
         """
-        return '30°C'  # pragma: no cover
+        return '30°C'
 
     event_parts: list[Any] = []
     async with agent.iter(user_prompt='What is the temperature of the capital of France?') as agent_run:
@@ -544,6 +549,7 @@ async def test_bedrock_model_instructions(allow_model_requests: None, bedrock_pr
                         content='The capital of France is Paris. Paris is not only the political and economic hub of the country but also a major center for culture, fashion, art, and tourism. It is renowned for its rich history, iconic landmarks such as the Eiffel Tower, Notre-Dame Cathedral, and the Louvre Museum, as well as its influence on global culture and cuisine.'
                     )
                 ],
+                usage=Usage(requests=1, request_tokens=13, response_tokens=71, total_tokens=84),
                 model_name='us.amazon.nova-pro-v1:0',
                 timestamp=IsDatetime(),
             ),
@@ -581,3 +587,45 @@ async def test_bedrock_multiple_documents_in_history(
     assert result.output == snapshot(
         'Based on the documents you\'ve shared, both Document 1.pdf and Document 2.pdf contain the text "Dummy PDF file". These appear to be placeholder or sample PDF documents rather than files with substantial content.'
     )
+
+
+async def test_bedrock_group_consecutive_tool_return_parts(bedrock_provider: BedrockProvider):
+    """
+    Test that consecutive ToolReturnPart objects are grouped into a single user message for Bedrock.
+    """
+    model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider)
+    now = datetime.datetime.now()
+    # Create a ModelRequest with 3 consecutive ToolReturnParts
+    req = [
+        ModelRequest(parts=[UserPromptPart(content=['Hello'])]),
+        ModelResponse(parts=[TextPart(content='Hi')]),
+        ModelRequest(parts=[UserPromptPart(content=['How are you?'])]),
+        ModelResponse(parts=[TextPart(content='Cloudy')]),
+        ModelRequest(
+            parts=[
+                ToolReturnPart(tool_name='tool1', content='result1', tool_call_id='id1', timestamp=now),
+                ToolReturnPart(tool_name='tool2', content='result2', tool_call_id='id2', timestamp=now),
+                ToolReturnPart(tool_name='tool3', content='result3', tool_call_id='id3', timestamp=now),
+            ]
+        ),
+    ]
+
+    # Call the mapping function directly
+    _, bedrock_messages = await model._map_messages(req)  # type: ignore[reportPrivateUsage]
+
+    assert bedrock_messages == snapshot(
+        [
+            {'role': 'user', 'content': [{'text': 'Hello'}]},
+            {'role': 'assistant', 'content': [{'text': 'Hi'}]},
+            {'role': 'user', 'content': [{'text': 'How are you?'}]},
+            {'role': 'assistant', 'content': [{'text': 'Cloudy'}]},
+            {
+                'role': 'user',
+                'content': [
+                    {'toolResult': {'toolUseId': 'id1', 'content': [{'text': 'result1'}], 'status': 'success'}},
+                    {'toolResult': {'toolUseId': 'id2', 'content': [{'text': 'result2'}], 'status': 'success'}},
+                    {'toolResult': {'toolUseId': 'id3', 'content': [{'text': 'result3'}], 'status': 'success'}},
+                ],
+            },
+        ]
+    )

{pydantic_ai-0.1.10 → pydantic_ai-0.1.12}/tests/models/test_cohere.py

@@ -66,9 +66,7 @@ class MockAsyncClientV2:
     def create_mock(cls, completions: MockChatResponse | Sequence[MockChatResponse]) -> AsyncClientV2:
         return cast(AsyncClientV2, cls(completions=completions))
 
-    async def chat(  # pragma: no cover
-        self, *_args: Any, **_kwargs: Any
-    ) -> ChatResponse:
+    async def chat(self, *_args: Any, **_kwargs: Any) -> ChatResponse:
         assert self.completions is not None
         if isinstance(self.completions, Sequence):
             raise_if_exception(self.completions[self.index])
@@ -115,11 +113,17 @@ async def test_request_simple_success(allow_model_requests: None):
         [
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
-                parts=[TextPart(content='world')], model_name='command-r7b-12-2024', timestamp=IsNow(tz=timezone.utc)
+                parts=[TextPart(content='world')],
+                usage=Usage(requests=1),
+                model_name='command-r7b-12-2024',
+                timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
-                parts=[TextPart(content='world')], model_name='command-r7b-12-2024', timestamp=IsNow(tz=timezone.utc)
+                parts=[TextPart(content='world')],
+                usage=Usage(requests=1),
+                model_name='command-r7b-12-2024',
+                timestamp=IsNow(tz=timezone.utc),
             ),
         ]
     )
@@ -187,6 +191,7 @@ async def test_request_structured_response(allow_model_requests: None):
                         tool_call_id='123',
                     )
                 ],
+                usage=Usage(requests=1),
                 model_name='command-r7b-12-2024',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -273,6 +278,7 @@ async def test_request_tool_call(allow_model_requests: None):
                         tool_call_id='1',
                     )
                 ],
+                usage=Usage(requests=1, total_tokens=0, details={}),
                 model_name='command-r7b-12-2024',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -294,6 +300,13 @@ async def test_request_tool_call(allow_model_requests: None):
                         tool_call_id='2',
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=5,
+                    response_tokens=3,
+                    total_tokens=8,
+                    details={'input_tokens': 4, 'output_tokens': 2},
+                ),
                 model_name='command-r7b-12-2024',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -309,6 +322,7 @@ async def test_request_tool_call(allow_model_requests: None):
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
+                usage=Usage(requests=1),
                 model_name='command-r7b-12-2024',
                 timestamp=IsNow(tz=timezone.utc),
             ),
@@ -386,6 +400,13 @@ async def test_cohere_model_instructions(allow_model_requests: None, co_api_key:
                         content="The capital of France is Paris. It is the country's largest city and serves as the economic, cultural, and political center of France. Paris is known for its rich history, iconic landmarks such as the Eiffel Tower and the Louvre Museum, and its significant influence on fashion, cuisine, and the arts."
                     )
                 ],
+                usage=Usage(
+                    requests=1,
+                    request_tokens=542,
+                    response_tokens=63,
+                    total_tokens=605,
+                    details={'input_tokens': 13, 'output_tokens': 61},
+                ),
                 model_name='command-r7b-12-2024',
                 timestamp=IsDatetime(),
             ),