Turn RunContext.retries from a defaultdict into a dict again as the 0 being stored on read broke a test

Author: DouweM

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -1,7 +1,6 @@
 from __future__ import annotations as _annotations
 
 import dataclasses
-from collections import defaultdict
 from collections.abc import Sequence
 from dataclasses import field
 from typing import TYPE_CHECKING, Generic
@@ -34,7 +33,7 @@ class RunContext(Generic[AgentDepsT]):
     """The original user prompt passed to the run."""
     messages: list[_messages.ModelMessage] = field(default_factory=list)
     """Messages exchanged in the conversation so far."""
-    retries: defaultdict[str, int] = field(default_factory=lambda: defaultdict(int))
+    retries: dict[str, int] = field(default_factory=dict)
     """Number of retries for each tool so far."""
     tool_call_id: str | None = None
     """The ID of the tool call."""

pydantic_ai_slim/pydantic_ai/toolsets/_run.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass, replace
@@ -25,7 +24,7 @@ class RunToolset(WrapperToolset[AgentDepsT]):
     ctx: RunContext[AgentDepsT]
     _tool_defs: list[ToolDefinition]
     _tool_names: list[str]
-    _retries: defaultdict[str, int]
+    _retries: dict[str, int]
     _original: AbstractToolset[AgentDepsT]
 
     def __init__(
@@ -108,14 +107,14 @@ def _with_retry(self, name: str, ctx: RunContext[AgentDepsT]) -> Iterator[RunCon
                     msg = 'No tools available.'
                 raise ModelRetry(f'Unknown tool name: {name!r}. {msg}')
 
-            ctx = replace(ctx, tool_name=name, retry=self._retries[name], retries={})
+            ctx = replace(ctx, tool_name=name, retry=self._retries.get(name, 0), retries={})
             yield ctx
         except (ValidationError, ModelRetry, UnexpectedModelBehavior, ToolRetryError) as e:
             try:
                 max_retries = self._max_retries_for_tool(name)
             except Exception:
                 max_retries = 1
-            current_retry = self._retries[name]
+            current_retry = self._retries.get(name, 0)
 
             if isinstance(e, UnexpectedModelBehavior) and e.__cause__ is not None:
                 e = e.__cause__

pydantic_ai_slim/pydantic_ai/toolsets/function.py

@@ -188,7 +188,7 @@ async def prepare_for_run(self, ctx: RunContext[AgentDepsT]) -> RunToolset[Agent
 
     async def _prepare_tool_def(self, ctx: RunContext[AgentDepsT], tool_def: ToolDefinition) -> ToolDefinition | None:
         tool_name = tool_def.name
-        ctx = replace(ctx, tool_name=tool_name, retry=ctx.retries[tool_name])
+        ctx = replace(ctx, tool_name=tool_name, retry=ctx.retries.get(tool_name, 0))
         return await self.tools[tool_name].prepare_tool_def(ctx)
 
     @property

tests/test_tools.py

@@ -1155,7 +1155,7 @@ async def prepare_tool_defs(
         ctx: RunContext[None], tool_defs: list[ToolDefinition]
     ) -> Union[list[ToolDefinition], None]:
         nonlocal prepare_tools_retries
-        retry = ctx.retries['infinite_retry_tool']
+        retry = ctx.retries.get('infinite_retry_tool', 0)
         prepare_tools_retries.append(retry)
         return tool_defs