pydantic-ai-slim 1.0.0b1__py3-none-any.whl → 1.0.2__py3-none-any.whl

pydantic_ai/providers/litellm.py

@@ -0,0 +1,134 @@
+from __future__ import annotations as _annotations
+
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.groq import groq_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.profiles.qwen import qwen_model_profile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the LiteLLM provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class LiteLLMProvider(Provider[AsyncOpenAI]):
+    """Provider for LiteLLM API."""
+
+    @property
+    def name(self) -> str:
+        return 'litellm'
+
+    @property
+    def base_url(self) -> str:
+        return str(self.client.base_url)
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        # Map provider prefixes to their profile functions
+        provider_to_profile = {
+            'anthropic': anthropic_model_profile,
+            'openai': openai_model_profile,
+            'google': google_model_profile,
+            'mistralai': mistral_model_profile,
+            'mistral': mistral_model_profile,
+            'cohere': cohere_model_profile,
+            'amazon': amazon_model_profile,
+            'bedrock': amazon_model_profile,
+            'meta-llama': meta_model_profile,
+            'meta': meta_model_profile,
+            'groq': groq_model_profile,
+            'deepseek': deepseek_model_profile,
+            'moonshotai': moonshotai_model_profile,
+            'x-ai': grok_model_profile,
+            'qwen': qwen_model_profile,
+        }
+
+        profile = None
+
+        # Check if model name contains a provider prefix (e.g., "anthropic/claude-3")
+        if '/' in model_name:
+            provider_prefix, model_suffix = model_name.split('/', 1)
+            if provider_prefix in provider_to_profile:
+                profile = provider_to_profile[provider_prefix](model_suffix)
+
+        # If no profile found, default to OpenAI profile
+        if profile is None:
+            profile = openai_model_profile(model_name)
+
+        # As LiteLLMProvider is used with OpenAIModel, which uses OpenAIJsonSchemaTransformer,
+        # we maintain that behavior
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        http_client: AsyncHTTPClient,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        """Initialize a LiteLLM provider.
+
+        Args:
+            api_key: API key for the model provider. If None, LiteLLM will try to get it from environment variables.
+            api_base: Base URL for the model provider. Use this for custom endpoints or self-hosted models.
+            openai_client: Pre-configured OpenAI client. If provided, other parameters are ignored.
+            http_client: Custom HTTP client to use.
+        """
+        if openai_client is not None:
+            self._client = openai_client
+            return
+
+        # Create OpenAI client that will be used with LiteLLM's completion function
+        # The actual API calls will be intercepted and routed through LiteLLM
+        if http_client is not None:
+            self._client = AsyncOpenAI(
+                base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )
+        else:
+            http_client = cached_async_http_client(provider='litellm')
+            self._client = AsyncOpenAI(
+                base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )

pydantic_ai/retries.py

@@ -13,6 +13,8 @@ The module includes:
 
 from __future__ import annotations
 
+from types import TracebackType
+
 from httpx import (
     AsyncBaseTransport,
     AsyncHTTPTransport,
@@ -185,11 +187,30 @@ class TenacityTransport(BaseTransport):
             response.request = req
 
             if self.validate_response:
-                self.validate_response(response)
+                try:
+                    self.validate_response(response)
+                except Exception:
+                    response.close()
+                    raise
             return response
 
         return handle_request(request)
 
+    def __enter__(self) -> TenacityTransport:
+        self.wrapped.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        self.wrapped.__exit__(exc_type, exc_value, traceback)
+
+    def close(self) -> None:
+        self.wrapped.close()  # pragma: no cover
+
 
 class AsyncTenacityTransport(AsyncBaseTransport):
     """Asynchronous HTTP transport with tenacity-based retry functionality.
@@ -263,11 +284,30 @@ class AsyncTenacityTransport(AsyncBaseTransport):
             response.request = req
 
             if self.validate_response:
-                self.validate_response(response)
+                try:
+                    self.validate_response(response)
+                except Exception:
+                    await response.aclose()
+                    raise
             return response
 
         return await handle_async_request(request)
 
+    async def __aenter__(self) -> AsyncTenacityTransport:
+        await self.wrapped.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        await self.wrapped.__aexit__(exc_type, exc_value, traceback)
+
+    async def aclose(self) -> None:
+        await self.wrapped.aclose()
+
 
 def wait_retry_after(
     fallback_strategy: Callable[[RetryCallState], float] | None = None, max_wait: float = 300

pydantic_ai/tools.py

@@ -70,7 +70,7 @@ Usage `ToolFuncEither[AgentDepsT, ToolParams]`.
 ToolPrepareFunc: TypeAlias = Callable[[RunContext[AgentDepsT], 'ToolDefinition'], Awaitable['ToolDefinition | None']]
 """Definition of a function that can prepare a tool definition at call time.
 
-See [tool docs](../tools.md#tool-prepare) for more information.
+See [tool docs](../tools-advanced.md#tool-prepare) for more information.
 
 Example — here `only_if_42` is valid as a `ToolPrepareFunc`:
 
@@ -140,7 +140,7 @@ class DeferredToolRequests:
 
     Results can be passed to the next agent run using a [`DeferredToolResults`][pydantic_ai.tools.DeferredToolResults] object with the same tool call IDs.
 
-    See [deferred tools docs](../tools.md#deferred-tools) for more information.
+    See [deferred tools docs](../deferred-tools.md#deferred-tools) for more information.
     """
 
     calls: list[ToolCallPart] = field(default_factory=list)
@@ -204,7 +204,7 @@ class DeferredToolResults:
 
     The tool call IDs need to match those from the [`DeferredToolRequests`][pydantic_ai.output.DeferredToolRequests] output object from the previous run.
 
-    See [deferred tools docs](../tools.md#deferred-tools) for more information.
+    See [deferred tools docs](../deferred-tools.md#deferred-tools) for more information.
     """
 
     calls: dict[str, DeferredToolCallResult | Any] = field(default_factory=dict)
@@ -253,6 +253,7 @@ class Tool(Generic[AgentDepsT]):
     docstring_format: DocstringFormat
     require_parameter_descriptions: bool
     strict: bool | None
+    sequential: bool
     requires_approval: bool
     function_schema: _function_schema.FunctionSchema
     """
@@ -274,6 +275,7 @@ class Tool(Generic[AgentDepsT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
         function_schema: _function_schema.FunctionSchema | None = None,
     ):
@@ -327,8 +329,9 @@ class Tool(Generic[AgentDepsT]):
             schema_generator: The JSON schema generator class to use. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
-                See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
+                See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
             function_schema: The function schema to use for the tool. If not provided, it will be generated.
         """
         self.function = function
@@ -347,6 +350,7 @@ class Tool(Generic[AgentDepsT]):
         self.docstring_format = docstring_format
         self.require_parameter_descriptions = require_parameter_descriptions
         self.strict = strict
+        self.sequential = sequential
         self.requires_approval = requires_approval
 
     @classmethod
@@ -357,6 +361,7 @@ class Tool(Generic[AgentDepsT]):
         description: str | None,
         json_schema: JsonSchemaValue,
         takes_ctx: bool = False,
+        sequential: bool = False,
     ) -> Self:
         """Creates a Pydantic tool from a function and a JSON schema.
 
@@ -370,6 +375,7 @@ class Tool(Generic[AgentDepsT]):
             json_schema: The schema for the function arguments
             takes_ctx: An optional boolean parameter indicating whether the function
                 accepts the context object as an argument.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
 
         Returns:
             A Pydantic tool that calls the function
@@ -389,6 +395,7 @@ class Tool(Generic[AgentDepsT]):
             name=name,
             description=description,
             function_schema=function_schema,
+            sequential=sequential,
         )
 
     @property
@@ -398,6 +405,7 @@ class Tool(Generic[AgentDepsT]):
             description=self.description,
             parameters_json_schema=self.function_schema.json_schema,
             strict=self.strict,
+            sequential=self.sequential,
         )
 
     async def prepare_tool_def(self, ctx: RunContext[AgentDepsT]) -> ToolDefinition | None:
@@ -466,22 +474,25 @@ class ToolDefinition:
     Note: this is currently only supported by OpenAI models.
     """
 
+    sequential: bool = False
+    """Whether this tool requires a sequential/serial execution environment."""
+
     kind: ToolKind = field(default='function')
     """The kind of tool:
 
     - `'function'`: a tool that will be executed by Pydantic AI during an agent run and has its result returned to the model
     - `'output'`: a tool that passes through an output value that ends the run
     - `'external'`: a tool whose result will be produced outside of the Pydantic AI agent run in which it was called, because it depends on an upstream service (or user) or could take longer to generate than it's reasonable to keep the agent process running.
-        See the [tools documentation](../tools.md#deferred-tools) for more info.
+        See the [tools documentation](../deferred-tools.md#deferred-tools) for more info.
     - `'unapproved'`: a tool that requires human-in-the-loop approval.
-        See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
+        See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
     """
 
     @property
     def defer(self) -> bool:
         """Whether calls to this tool will be deferred.
 
-        See the [tools documentation](../tools.md#deferred-tools) for more info.
+        See the [tools documentation](../deferred-tools.md#deferred-tools) for more info.
         """
         return self.kind in ('external', 'unapproved')
 

pydantic_ai/toolsets/combined.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 import asyncio
+from asyncio import Lock
 from collections.abc import Callable, Sequence
 from contextlib import AsyncExitStack
 from dataclasses import dataclass, field, replace
 from typing import Any
 
-import anyio
 from typing_extensions import Self
 
 from .._run_context import AgentDepsT, RunContext
@@ -31,7 +31,7 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
 
     toolsets: Sequence[AbstractToolset[AgentDepsT]]
 
-    _enter_lock: anyio.Lock = field(compare=False, init=False, default_factory=anyio.Lock)
+    _enter_lock: Lock = field(compare=False, init=False, default_factory=Lock)
     _entered_count: int = field(init=False, default=0)
     _exit_stack: AsyncExitStack | None = field(init=False, default=None)
 

pydantic_ai/toolsets/function.py

@@ -33,9 +33,12 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
     See [toolset docs](../toolsets.md#function-toolset) for more information.
     """
 
-    max_retries: int
     tools: dict[str, Tool[Any]]
+    max_retries: int
     _id: str | None
+    docstring_format: DocstringFormat
+    require_parameter_descriptions: bool
+    schema_generator: type[GenerateJsonSchema]
 
     def __init__(
         self,
@@ -43,16 +46,30 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         *,
         max_retries: int = 1,
         id: str | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ):
         """Build a new function toolset.
 
         Args:
             tools: The tools to add to the toolset.
             max_retries: The maximum number of retries for each tool during a run.
-            id: An optional unique ID for the toolset. A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the toolset's activities within the workflow.
+            id: An optional unique ID for the toolset. A toolset needs to have an ID in order to be used in a durable execution environment like Temporal,
+                in which case the ID will be used to identify the toolset's activities within the workflow.
+            docstring_format: Format of tool docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+                Applies to all tools, unless overridden when adding a tool.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+                Applies to all tools, unless overridden when adding a tool.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                Applies to all tools, unless overridden when adding a tool.
         """
         self.max_retries = max_retries
         self._id = id
+        self.docstring_format = docstring_format
+        self.require_parameter_descriptions = require_parameter_descriptions
+        self.schema_generator = schema_generator
 
         self.tools = {}
         for tool in tools:
@@ -76,10 +93,11 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Callable[[ToolFuncEither[AgentDepsT, ToolParams]], ToolFuncEither[AgentDepsT, ToolParams]]: ...
 
@@ -91,10 +109,11 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
@@ -137,13 +156,16 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
-                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
-            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
-            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                If `None`, the default value is determined by the toolset.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing.
+                If `None`, the default value is determined by the toolset.
+            schema_generator: The JSON schema generator class to use for this tool.
+                If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
-                See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
+                See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
 
         def tool_decorator(
@@ -160,6 +182,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                sequential,
                 requires_approval,
             )
             return func_
@@ -173,10 +196,11 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         name: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
-        docstring_format: DocstringFormat = 'auto',
-        require_parameter_descriptions: bool = False,
-        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        docstring_format: DocstringFormat | None = None,
+        require_parameter_descriptions: bool | None = None,
+        schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> None:
         """Add a function as a tool to the toolset.
@@ -196,14 +220,24 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
-                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
-            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
-            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+                If `None`, the default value is determined by the toolset.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing.
+                If `None`, the default value is determined by the toolset.
+            schema_generator: The JSON schema generator class to use for this tool.
+                If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
-                See the [tools documentation](../tools.md#human-in-the-loop-tool-approval) for more info.
+                See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
+        if docstring_format is None:
+            docstring_format = self.docstring_format
+        if require_parameter_descriptions is None:
+            require_parameter_descriptions = self.require_parameter_descriptions
+        if schema_generator is None:
+            schema_generator = self.schema_generator
+
         tool = Tool[AgentDepsT](
             func,
             takes_ctx=takes_ctx,
@@ -214,6 +248,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
             require_parameter_descriptions=require_parameter_descriptions,
             schema_generator=schema_generator,
             strict=strict,
+            sequential=sequential,
             requires_approval=requires_approval,
         )
         self.add_tool(tool)

pydantic_ai/usage.py

@@ -3,7 +3,9 @@ from __future__ import annotations as _annotations
 import dataclasses
 from copy import copy
 from dataclasses import dataclass, fields
+from typing import Annotated
 
+from pydantic import AliasChoices, BeforeValidator, Field
 from typing_extensions import deprecated, overload
 
 from . import _utils
@@ -14,7 +16,11 @@ __all__ = 'RequestUsage', 'RunUsage', 'Usage', 'UsageLimits'
 
 @dataclass(repr=False, kw_only=True)
 class UsageBase:
-    input_tokens: int = 0
+    input_tokens: Annotated[
+        int,
+        # `request_tokens` is deprecated, but we still want to support deserializing model responses stored in a DB before the name was changed
+        Field(validation_alias=AliasChoices('input_tokens', 'request_tokens')),
+    ] = 0
     """Number of input/prompt tokens."""
 
     cache_write_tokens: int = 0
@@ -22,7 +28,11 @@ class UsageBase:
     cache_read_tokens: int = 0
     """Number of tokens read from the cache."""
 
-    output_tokens: int = 0
+    output_tokens: Annotated[
+        int,
+        # `response_tokens` is deprecated, but we still want to support deserializing model responses stored in a DB before the name was changed
+        Field(validation_alias=AliasChoices('output_tokens', 'response_tokens')),
+    ] = 0
     """Number of output/completion tokens."""
 
     input_audio_tokens: int = 0
@@ -32,7 +42,11 @@ class UsageBase:
     output_audio_tokens: int = 0
     """Number of audio output tokens."""
 
-    details: dict[str, int] = dataclasses.field(default_factory=dict)
+    details: Annotated[
+        dict[str, int],
+        # `details` can not be `None` any longer, but we still want to support deserializing model responses stored in a DB before this was changed
+        BeforeValidator(lambda d: d or {}),
+    ] = dataclasses.field(default_factory=dict)
     """Any extra details returned by the model."""
 
     @property
@@ -117,6 +131,9 @@ class RunUsage(UsageBase):
     requests: int = 0
     """Number of requests made to the LLM API."""
 
+    tool_calls: int = 0
+    """Number of successful tool calls executed during the run."""
+
     input_tokens: int = 0
     """Total number of text input/prompt tokens."""
 
@@ -146,6 +163,7 @@ class RunUsage(UsageBase):
         """
         if isinstance(incr_usage, RunUsage):
             self.requests += incr_usage.requests
+            self.tool_calls += incr_usage.tool_calls
         return _incr_usage_tokens(self, incr_usage)
 
     def __add__(self, other: RunUsage | RequestUsage) -> RunUsage:
@@ -194,6 +212,8 @@ class UsageLimits:
 
     request_limit: int | None = 50
     """The maximum number of requests allowed to the model."""
+    tool_calls_limit: int | None = None
+    """The maximum number of successful tool calls allowed to be executed."""
     input_tokens_limit: int | None = None
     """The maximum number of input/prompt tokens allowed."""
     output_tokens_limit: int | None = None
@@ -220,12 +240,14 @@ class UsageLimits:
         self,
         *,
         request_limit: int | None = 50,
+        tool_calls_limit: int | None = None,
         input_tokens_limit: int | None = None,
         output_tokens_limit: int | None = None,
         total_tokens_limit: int | None = None,
         count_tokens_before_request: bool = False,
     ) -> None:
         self.request_limit = request_limit
+        self.tool_calls_limit = tool_calls_limit
         self.input_tokens_limit = input_tokens_limit
         self.output_tokens_limit = output_tokens_limit
         self.total_tokens_limit = total_tokens_limit
@@ -239,12 +261,14 @@ class UsageLimits:
         self,
         *,
         request_limit: int | None = 50,
+        tool_calls_limit: int | None = None,
         request_tokens_limit: int | None = None,
         response_tokens_limit: int | None = None,
         total_tokens_limit: int | None = None,
         count_tokens_before_request: bool = False,
     ) -> None:
         self.request_limit = request_limit
+        self.tool_calls_limit = tool_calls_limit
         self.input_tokens_limit = request_tokens_limit
         self.output_tokens_limit = response_tokens_limit
         self.total_tokens_limit = total_tokens_limit
@@ -254,6 +278,7 @@ class UsageLimits:
         self,
         *,
         request_limit: int | None = 50,
+        tool_calls_limit: int | None = None,
         input_tokens_limit: int | None = None,
         output_tokens_limit: int | None = None,
         total_tokens_limit: int | None = None,
@@ -263,6 +288,7 @@ class UsageLimits:
         response_tokens_limit: int | None = None,
     ):
         self.request_limit = request_limit
+        self.tool_calls_limit = tool_calls_limit
         self.input_tokens_limit = input_tokens_limit or request_tokens_limit
         self.output_tokens_limit = output_tokens_limit or response_tokens_limit
         self.total_tokens_limit = total_tokens_limit
@@ -314,4 +340,12 @@ class UsageLimits:
         if self.total_tokens_limit is not None and total_tokens > self.total_tokens_limit:
             raise UsageLimitExceeded(f'Exceeded the total_tokens_limit of {self.total_tokens_limit} ({total_tokens=})')
 
+    def check_before_tool_call(self, usage: RunUsage) -> None:
+        """Raises a `UsageLimitExceeded` exception if the next tool call would exceed the tool call limit."""
+        tool_calls_limit = self.tool_calls_limit
+        if tool_calls_limit is not None and usage.tool_calls >= tool_calls_limit:
+            raise UsageLimitExceeded(
+                f'The next tool call would exceed the tool_calls_limit of {tool_calls_limit} (tool_calls={usage.tool_calls})'
+            )
+
     __repr__ = _utils.dataclasses_no_defaults_repr