Add exception for tool output exceeding token limit

Author: erhuve

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -24,6 +24,7 @@
     FallbackExceptionGroup,
     ModelHTTPError,
     ModelRetry,
+    ToolExceedsTokenLimitError,
     UnexpectedModelBehavior,
     UsageLimitExceeded,
     UserError,
@@ -124,6 +125,7 @@
     'ModelRetry',
     'ModelHTTPError',
     'FallbackExceptionGroup',
+    'ToolExceedsTokenLimitError',
     'UnexpectedModelBehavior',
     'UsageLimitExceeded',
     'UserError',

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -576,6 +576,15 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:  # noqa
                 tool_calls: list[_messages.ToolCallPart] = []
                 files: list[_messages.BinaryContent] = []
 
+                # If the model stopped due to token limit (length) and the last part is a tool call,
+                # the args are incomplete.
+                if self.model_response.finish_reason == 'length':
+                    last_part = self.model_response.parts[-1]
+                    if isinstance(last_part, _messages.ToolCallPart):
+                        raise exceptions.ToolExceedsTokenLimitError(
+                            'Model token limit exceeded while emitting a tool call.'
+                        )
+
                 for part in self.model_response.parts:
                     if isinstance(part, _messages.TextPart):
                         text += part.content

pydantic_ai_slim/pydantic_ai/exceptions.py

@@ -23,6 +23,7 @@
     'UnexpectedModelBehavior',
     'UsageLimitExceeded',
     'ModelHTTPError',
+    'ToolExceedsTokenLimitError',
     'FallbackExceptionGroup',
 )
 
@@ -168,3 +169,7 @@ class ToolRetryError(Exception):
     def __init__(self, tool_retry: RetryPromptPart):
         self.tool_retry = tool_retry
         super().__init__()
+
+
+class ToolExceedsTokenLimitError(AgentRunError):
+    """Error raised when a model stops due to token limit while emitting a tool call."""

tests/test_agent.py

@@ -40,6 +40,7 @@
     SystemPromptPart,
     TextPart,
     ToolCallPart,
+    ToolExceedsTokenLimitError,
     ToolReturn,
     ToolReturnPart,
     UnexpectedModelBehavior,
@@ -2448,6 +2449,18 @@ def empty(m: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
     )
 
 
+def test_tool_exceeds_token_limit_error():
+    def return_incomplete_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        resp = ModelResponse(parts=[ToolCallPart('dummy_tool', args='{"foo": "bar",')])
+        resp.finish_reason = 'length'
+        return resp
+
+    agent = Agent(FunctionModel(return_incomplete_tool), output_type=str)
+
+    with pytest.raises(ToolExceedsTokenLimitError, match='Model token limit exceeded while emitting a tool call.'):
+        agent.run_sync('Hello')
+
+
 def test_model_requests_blocked(env: TestEnv):
     try:
         env.set('GEMINI_API_KEY', 'foobar')