From dac259af1d0966c94b1d370320be8d84c1b35a63 Mon Sep 17 00:00:00 2001 From: Jeroen Schweitzer Date: Tue, 11 Aug 2026 17:45:48 +0200 Subject: [PATCH] refactor(agents): type the agent and its conversation history MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Partial work on the typecheck gate: 103 mypy errors down to 95, and the two shared roots in agents/tatlock.py are gone. The rest is genuine per-function annotation work and is not attempted here. Five conversation lists were declared bare. mypy infers the element type from the first append, which is a ModelRequest, and then rejects every ModelResponse that follows — five errors from five lists that all hold the same thing: a conversation, which is both kinds of message. Annotated as list[ModelMessage], which is pydantic_ai's own union for exactly this. The agent had no deps type. It is built as Agent(model, system_prompt=...), inferred Agent[None, str], while every tool it registers takes RunContext[ToolCallTracker] and run() is called with a tracker. The declaration now says what was already happening: Agent[ToolCallTracker, str]. Note this is a runtime-visible change — pydantic_ai is now told the deps type it was being handed anyway — so it was verified against the suite rather than reasoned about: 658 passed. _register_tools carries an assert rather than a None check. It is called from _ensure_agent immediately after the agent is constructed, so a None there is a broken invariant, not a case to handle; an `if is None: return` would silently register no tools. Two corrections to my own work in this commit. Declaring `_agent: Agent | None` first made things worse, not better — resolving the bare Agent to Agent[None, str] surfaced four new argument-type errors that the Any had been hiding, which is how the missing deps type became visible at all. And an import fix I thought I had made was a no-op: the target was a multi-line import, my replace matched nothing, and I had asserted the precondition without asserting the result. Ruff caught it. That is the same mistake as a changelog edit earlier today, so the assert now checks what landed. Co-Authored-By: Claude --- src/agents/tatlock.py | 58 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/src/agents/tatlock.py b/src/agents/tatlock.py index 4d04d4b..b3331be 100644 --- a/src/agents/tatlock.py +++ b/src/agents/tatlock.py @@ -157,7 +157,10 @@ class TatlockAgent(AgentInterface): def __init__(self): """Initialize Tatlock (lazy agent creation).""" - self._agent = None # Lazy initialization + # Deps are a ToolCallTracker: every registered tool takes + # RunContext[ToolCallTracker], and run() is called with one. Saying so + # is what lets the tool registrations below type-check at all. + self._agent: Agent[ToolCallTracker, str] | None = None # Lazy initialization def _ensure_agent(self): """Ensure the PydanticAI agent is initialized (lazy initialization).""" @@ -180,13 +183,19 @@ class TatlockAgent(AgentInterface): self._agent = Agent( model, system_prompt=TATLOCK_SYSTEM_PROMPT, + deps_type=ToolCallTracker, ) # Register tools with the agent self._register_tools() - def _register_tools(self): - """Register permanent tools with the PydanticAI agent.""" + def _register_tools(self) -> None: + """Register permanent tools with the PydanticAI agent. + + Called only from _ensure_agent, immediately after the agent is built, so + the assert documents an invariant rather than guarding a real case. + """ + assert self._agent is not None, "_register_tools called before the agent exists" # Calculator tool @self._agent.tool @@ -324,9 +333,15 @@ class TatlockAgent(AgentInterface): # Build message history (all messages except the last user message) # PydanticAI expects history as list of ModelRequest/ModelResponse objects - from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart + from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) - message_history = [] + message_history: list[ModelMessage] = [] for i, msg in enumerate(messages[:-1]): # All messages except the last one role = msg.get("role") content = msg.get("content", "") @@ -498,9 +513,15 @@ class TatlockAgent(AgentInterface): enriched_message = f"{steward_note}\n\n{user_message}" # Convert message history to PydanticAI format - from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart + from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) - pydantic_history = [] + pydantic_history: list[ModelMessage] = [] for msg in message_history: role = msg.get("role") content = msg.get("content", "") @@ -580,9 +601,15 @@ class TatlockAgent(AgentInterface): enriched_message = f"{steward_note}\n\n{user_message}" # Convert message history to PydanticAI format - from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart + from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) - pydantic_history = [] + pydantic_history: list[ModelMessage] = [] for msg in message_history: role = msg.get("role") content = msg.get("content", "") @@ -642,6 +669,7 @@ class TatlockAgent(AgentInterface): - raw_output: The agent's raw text output """ from pydantic_ai.messages import ( + ModelMessage, ModelRequest, ModelResponse, TextPart, @@ -683,7 +711,7 @@ class TatlockAgent(AgentInterface): enriched_message = f"{steward_note}\n\n{user_message}" # Convert message history to PydanticAI format - pydantic_history = [] + pydantic_history: list[ModelMessage] = [] for msg in message_history: role = msg.get("role") content = msg.get("content", "") @@ -781,7 +809,13 @@ class TatlockAgent(AgentInterface): Returns: str: Butler-toned response synthesized from all results """ - from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart + from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) from src.anthropic.model_selector import get_model @@ -840,7 +874,7 @@ class TatlockAgent(AgentInterface): ) # Convert message history to PydanticAI format - pydantic_history = [] + pydantic_history: list[ModelMessage] = [] for msg in message_history: role = msg.get("role") content = msg.get("content", "")