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", "")