test: update usage assertions to include tool_calls in various tests

- Modified assertions in multiple test files to include the new `tool_calls` attribute in `RunUsage`.
- Added new test cases to validate tool call limits and ensure correct behavior when exceeding limits.
- Updated existing tests to reflect changes in tool call tracking across different models and scenarios.

Author: tradeqvest

tests/models/test_anthropic.py

@@ -659,7 +659,11 @@ async def my_tool(first: str, second: str) -> int:
                 requests=2,
                 input_tokens=20,
                 output_tokens=5,
-                details={'input_tokens': 20, 'output_tokens': 5},
+                tool_calls=1,
+                details={
+                    'input_tokens': 20,
+                    'output_tokens': 5,
+                },
             )
         )
         assert tool_called

tests/models/test_bedrock.py

@@ -111,7 +111,7 @@ async def temperature(city: str, date: datetime.date) -> str:
 
     result = await agent.run('What was the temperature in London 1st January 2022?', output_type=Response)
     assert result.output == snapshot({'temperature': '30°C', 'date': datetime.date(2022, 1, 1), 'city': 'London'})
-    assert result.usage() == snapshot(RunUsage(requests=2, input_tokens=1236, output_tokens=298))
+    assert result.usage() == snapshot(RunUsage(requests=2, input_tokens=1236, output_tokens=298, tool_calls=2))
     assert result.all_messages() == snapshot(
         [
             ModelRequest(

tests/models/test_cohere.py

@@ -325,6 +325,7 @@ async def get_location(loc_name: str) -> str:
             input_tokens=5,
             output_tokens=3,
             details={'input_tokens': 4, 'output_tokens': 2},
+            tool_calls=2,
         )
     )
 

tests/models/test_gemini.py

@@ -783,7 +783,7 @@ async def get_location(loc_name: str) -> str:
             ),
         ]
     )
-    assert result.usage() == snapshot(RunUsage(requests=3, input_tokens=3, output_tokens=6))
+    assert result.usage() == snapshot(RunUsage(requests=3, input_tokens=3, output_tokens=6, tool_calls=3))
 
 
 async def test_unexpected_response(client_with_handler: ClientWithHandler, env: TestEnv, allow_model_requests: None):
@@ -932,7 +932,7 @@ async def bar(y: str) -> str:
     async with agent.run_stream('Hello') as result:
         response = await result.get_output()
         assert response == snapshot((1, 2))
-    assert result.usage() == snapshot(RunUsage(requests=2, input_tokens=2, output_tokens=4))
+    assert result.usage() == snapshot(RunUsage(requests=2, input_tokens=2, output_tokens=4, tool_calls=2))
     assert result.all_messages() == snapshot(
         [
             ModelRequest(parts=[UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc))]),

tests/models/test_google.py

@@ -147,6 +147,7 @@ async def temperature(city: str, date: datetime.date) -> str:
             requests=2,
             input_tokens=224,
             output_tokens=35,
+            tool_calls=2,
             details={'text_prompt_tokens': 224, 'text_candidates_tokens': 35},
         )
     )

tests/models/test_openai.py

@@ -416,7 +416,9 @@ async def get_location(loc_name: str) -> str:
             ),
         ]
     )
-    assert result.usage() == snapshot(RunUsage(requests=3, cache_read_tokens=3, input_tokens=5, output_tokens=3))
+    assert result.usage() == snapshot(
+        RunUsage(requests=3, cache_read_tokens=3, input_tokens=5, output_tokens=3, tool_calls=2)
+    )
 
 
 FinishReason = Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call']

tests/test_agent.py

@@ -1954,7 +1954,7 @@ async def ret_a(x: str) -> str:
     assert result2.output == snapshot(Response(a=0))
     assert result2._new_message_index == snapshot(5)  # pyright: ignore[reportPrivateUsage]
     assert result2._output_tool_name == snapshot('final_result')  # pyright: ignore[reportPrivateUsage]
-    assert result2.usage() == snapshot(RunUsage(requests=1, input_tokens=59, output_tokens=13))
+    assert result2.usage() == snapshot(RunUsage(requests=1, input_tokens=59, output_tokens=13, tool_calls=1))
     new_msg_part_kinds = [(m.kind, [p.part_kind for p in m.parts]) for m in result2.all_messages()]
     assert new_msg_part_kinds == snapshot(
         [

tests/test_examples.py

@@ -369,6 +369,7 @@ async def call_tool(
         'The capital of Italy is Rome (Roma, in Italian), which has been a cultural and political center for centuries.'
         'Rome is known for its rich history, stunning architecture, and delicious cuisine.'
     ),
+    'Please call the tool twice': ToolCallPart(tool_name='do_work', args={}, tool_call_id='pyd_ai_tool_call_id'),
     'Begin infinite retry loop!': ToolCallPart(
         tool_name='infinite_retry_tool', args={}, tool_call_id='pyd_ai_tool_call_id'
     ),
@@ -626,6 +627,9 @@ async def model_logic(  # noqa: C901
         return ModelResponse(
             parts=[ToolCallPart(tool_name='final_result', args=args, tool_call_id='pyd_ai_tool_call_id')]
         )
+    elif isinstance(m, ToolReturnPart) and m.tool_name == 'do_work':
+        # Trigger a second tool call to cause tool_calls_limit to be exceeded in the docs example
+        return ModelResponse(parts=[ToolCallPart(tool_name='do_work', args={}, tool_call_id='pyd_ai_tool_call_id')])
     elif isinstance(m, RetryPromptPart) and m.tool_name == 'calc_volume':
         return ModelResponse(
             parts=[ToolCallPart(tool_name='calc_volume', args={'size': 6}, tool_call_id='pyd_ai_tool_call_id')]

tests/test_streaming.py

@@ -82,6 +82,7 @@ async def ret_a(x: str) -> str:
                 requests=2,
                 input_tokens=103,
                 output_tokens=5,
+                tool_calls=1,
             )
         )
         response = await result.get_output()
@@ -117,6 +118,7 @@ async def ret_a(x: str) -> str:
                 requests=2,
                 input_tokens=103,
                 output_tokens=11,
+                tool_calls=1,
             )
         )
 

tests/test_usage_limits.py

@@ -120,6 +120,7 @@ async def ret_a(x: str) -> str:
                     requests=2,
                     input_tokens=103,
                     output_tokens=5,
+                    tool_calls=1,
                 )
             )
             succeeded = True
@@ -151,26 +152,26 @@ async def delegate_to_other_agent1(ctx: RunContext[None], sentence: str) -> int:
         delegate_result = await delegate_agent.run(sentence)
         delegate_usage = delegate_result.usage()
         run_1_usages.append(delegate_usage)
-        assert delegate_usage == snapshot(RunUsage(requests=1, input_tokens=51, output_tokens=4))
+        assert delegate_usage == snapshot(RunUsage(requests=1, input_tokens=51, output_tokens=4, tool_calls=1))
         return delegate_result.output
 
     result1 = await controller_agent1.run('foobar')
     assert result1.output == snapshot('{"delegate_to_other_agent1":0}')
     run_1_usages.append(result1.usage())
-    assert result1.usage() == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=13))
+    assert result1.usage() == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=13, tool_calls=1))
 
     controller_agent2 = Agent(TestModel())
 
     @controller_agent2.tool
     async def delegate_to_other_agent2(ctx: RunContext[None], sentence: str) -> int:
         delegate_result = await delegate_agent.run(sentence, usage=ctx.usage)
         delegate_usage = delegate_result.usage()
-        assert delegate_usage == snapshot(RunUsage(requests=2, input_tokens=102, output_tokens=9))
+        assert delegate_usage == snapshot(RunUsage(requests=2, input_tokens=102, output_tokens=9, tool_calls=2))
         return delegate_result.output
 
     result2 = await controller_agent2.run('foobar')
     assert result2.output == snapshot('{"delegate_to_other_agent2":0}')
-    assert result2.usage() == snapshot(RunUsage(requests=3, input_tokens=154, output_tokens=17))
+    assert result2.usage() == snapshot(RunUsage(requests=3, input_tokens=154, output_tokens=17, tool_calls=2))
 
     # confirm the usage from result2 is the sum of the usage from result1
     assert result2.usage() == functools.reduce(operator.add, run_1_usages)
@@ -197,7 +198,7 @@ def delegate_to_other_agent(ctx: RunContext[None], sentence: str) -> int:
 
     result = await controller_agent.run('foobar')
     assert result.output == snapshot('{"delegate_to_other_agent":0}')
-    assert result.usage() == snapshot(RunUsage(requests=7, input_tokens=105, output_tokens=16))
+    assert result.usage() == snapshot(RunUsage(requests=7, input_tokens=105, output_tokens=16, tool_calls=1))
 
 
 def test_request_usage_basics():
@@ -215,6 +216,7 @@ def test_add_usages():
         cache_write_tokens=40,
         input_audio_tokens=50,
         cache_audio_read_tokens=60,
+        tool_calls=3,
         details={
             'custom1': 10,
             'custom2': 20,
@@ -229,13 +231,27 @@ def test_add_usages():
             cache_read_tokens=60,
             input_audio_tokens=100,
             cache_audio_read_tokens=120,
+            tool_calls=6,
             details={'custom1': 20, 'custom2': 40},
         )
     )
     assert usage + RunUsage() == usage
     assert RunUsage() + RunUsage() == RunUsage()
 
 
+async def test_tool_call_limit() -> None:
+    test_agent = Agent(TestModel())
+
+    @test_agent.tool_plain
+    async def ret_a(x: str) -> str:
+        return f'{x}-apple'
+
+    with pytest.raises(
+        UsageLimitExceeded, match=re.escape('The next tool call would exceed the tool_calls_limit of 0 (tool_calls=0)')
+    ):
+        await test_agent.run('Hello', usage_limits=UsageLimits(tool_calls_limit=0))
+
+
 def test_deprecated_usage_limits():
     with warns(
         snapshot(['DeprecationWarning: `request_tokens_limit` is deprecated, use `input_tokens_limit` instead'])