Put error in increment_retries to allow attempts, check lists for elements, check if args actually are incomplete or not

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

@@ -95,6 +95,19 @@ class GraphAgentState:
     def increment_retries(self, max_result_retries: int, error: BaseException | None = None) -> None:
         self.retries += 1
         if self.retries > max_result_retries:
+            if (
+                self.message_history
+                and isinstance(model_response := self.message_history[-1], _messages.ModelResponse)
+                and model_response.finish_reason == 'length'
+                and model_response.parts
+                and isinstance(tool_call := model_response.parts[-1], _messages.ToolCallPart)
+            ):
+                try:
+                    tool_call.args_as_dict()
+                except Exception:
+                    raise exceptions.ToolExceedsTokenLimitError(
+                        'Model token limit exceeded while emitting a tool call. Increase max tokens or simplify tool call arguments.'
+                    )
             message = f'Exceeded maximum retries ({max_result_retries}) for output validation'
             if error:
                 if isinstance(error, exceptions.UnexpectedModelBehavior) and error.__cause__ is not None:

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,39 @@ 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. Increase max tokens or simplify tool call arguments.',
+    ):
+        agent.run_sync('Hello')
+
+
+def test_tool_exceeds_token_limit_but_complete_args():
+    def return_complete_tool_but_hit_limit(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        if len(messages) == 1:
+            resp = ModelResponse(parts=[ToolCallPart('dummy_tool', args='{"foo": "bar"}')])
+            resp.finish_reason = 'length'
+            return resp
+        return ModelResponse(parts=[TextPart('done')])
+
+    agent = Agent(FunctionModel(return_complete_tool_but_hit_limit), output_type=str)
+
+    @agent.tool_plain
+    def dummy_tool(foo: str) -> str:
+        return 'tool-ok'
+
+    result = agent.run_sync('Hello')
+    assert result.output == 'done'
+
+
 def test_model_requests_blocked(env: TestEnv):
     try:
         env.set('GEMINI_API_KEY', 'foobar')