refactor(agents): type the agent and its conversation history
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 <noreply@anthropic.com>
This commit is contained in:
+46
-12
@@ -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", "")
|
||||
|
||||
Reference in New Issue
Block a user