chore(app): static typing

Author: fred913

app/agent/base.py

@@ -1,6 +1,5 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
-from typing import List, Optional
 
 from pydantic import BaseModel, Field, model_validator
 
@@ -19,13 +18,13 @@ class BaseAgent(BaseModel, ABC):
 
     # Core attributes
     name: str = Field(..., description="Unique name of the agent")
-    description: Optional[str] = Field(None, description="Optional agent description")
+    description: str | None = Field(None, description="Optional agent description")
 
     # Prompts
-    system_prompt: Optional[str] = Field(
+    system_prompt: str | None = Field(
         None, description="System-level instruction prompt"
     )
-    next_step_prompt: Optional[str] = Field(
+    next_step_prompt: str | None = Field(
         None, description="Prompt for determining next action"
     )
 
@@ -85,7 +84,7 @@ def update_memory(
         self,
         role: ROLE_TYPE,  # type: ignore
         content: str,
-        base64_image: Optional[str] = None,
+        base64_image: str | None = None,
         **kwargs,
     ) -> None:
         """Add a message to the agent's memory.
@@ -113,7 +112,7 @@ def update_memory(
         kwargs = {"base64_image": base64_image, **(kwargs if role == "tool" else {})}
         self.memory.add_message(message_map[role](content, **kwargs))
 
-    async def run(self, request: Optional[str] = None) -> str:
+    async def run(self, request: str | None = None) -> str:
         """Execute the agent's main loop asynchronously.
 
         Args:
@@ -131,7 +130,7 @@ async def run(self, request: Optional[str] = None) -> str:
         if request:
             self.update_memory("user", request)
 
-        results: List[str] = []
+        results: list[str] = []
         async with self.state_context(AgentState.RUNNING):
             while (
                 self.current_step < self.max_steps and self.state != AgentState.FINISHED
@@ -186,11 +185,11 @@ def is_stuck(self) -> bool:
         return duplicate_count >= self.duplicate_threshold
 
     @property
-    def messages(self) -> List[Message]:
+    def messages(self) -> list[Message]:
         """Retrieve a list of messages from the agent's memory."""
         return self.memory.messages
 
     @messages.setter
-    def messages(self, value: List[Message]):
+    def messages(self, value: list[Message]):
         """Set the list of messages in the agent's memory."""
         self.memory.messages = value

app/agent/browser.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Optional
+from typing import Any
 
 from pydantic import Field
 
@@ -18,14 +18,14 @@ class BrowserAgent(ToolCallAgent):
     extract content, and perform other browser-based actions to accomplish tasks.
     """
 
-    name: str = "browser"
-    description: str = "A browser agent that can control a browser to accomplish tasks"
+    name = "browser"
+    description = "A browser agent that can control a browser to accomplish tasks"
 
-    system_prompt: str = SYSTEM_PROMPT
-    next_step_prompt: str = NEXT_STEP_PROMPT
+    system_prompt = SYSTEM_PROMPT
+    next_step_prompt = NEXT_STEP_PROMPT
 
-    max_observe: int = 10000
-    max_steps: int = 20
+    max_observe = 10000
+    max_steps = 20
 
     # Configure the available tools
     available_tools: ToolCollection = Field(
@@ -36,18 +36,18 @@ class BrowserAgent(ToolCallAgent):
     tool_choices: ToolChoice = ToolChoice.AUTO
     special_tool_names: list[str] = Field(default_factory=lambda: [Terminate().name])
 
-    _current_base64_image: Optional[str] = None
+    _current_base64_image: str | None = None
 
     async def _handle_special_tool(self, name: str, result: Any, **kwargs):
         if not self._is_special_tool(name):
             return
         else:
-            await self.available_tools.get_tool(BrowserUseTool().name).cleanup()
+            await self.available_tools.get_tool(BrowserUseTool).cleanup()
             await super()._handle_special_tool(name, result, **kwargs)
 
-    async def get_browser_state(self) -> Optional[dict]:
+    async def get_browser_state(self) -> dict | None:
         """Get the current browser state for context in next steps."""
-        browser_tool = self.available_tools.get_tool(BrowserUseTool().name)
+        browser_tool = self.available_tools.get_tool(BrowserUseTool)
         if not browser_tool:
             return None
 

app/agent/planning.py

@@ -1,9 +1,9 @@
 import time
-from typing import Dict, List, Optional
+from typing import Sequence
 
 from pydantic import Field, model_validator
 
-from app.agent.toolcall import ToolCallAgent
+from app.agent.toolcall import CompatibleToolCallObject, ToolCallAgent
 from app.logger import logger
 from app.prompt.planning import NEXT_STEP_PROMPT, PLANNING_SYSTEM_PROMPT
 from app.schema import TOOL_CHOICE_TYPE, Message, ToolCall, ToolChoice
@@ -28,14 +28,16 @@ class PlanningAgent(ToolCallAgent):
         default_factory=lambda: ToolCollection(PlanningTool(), Terminate())
     )
     tool_choices: TOOL_CHOICE_TYPE = ToolChoice.AUTO  # type: ignore
-    special_tool_names: List[str] = Field(default_factory=lambda: [Terminate().name])
+    special_tool_names: list[str] = Field(default_factory=lambda: [Terminate().name])
 
-    tool_calls: List[ToolCall] = Field(default_factory=list)
-    active_plan_id: Optional[str] = Field(default=None)
+    tool_calls: Sequence[ToolCall | CompatibleToolCallObject] = Field(
+        default_factory=list
+    )
+    active_plan_id: str | None = Field(default=None)
 
     # Add a dictionary to track the step status for each tool call
-    step_execution_tracker: Dict[str, Dict] = Field(default_factory=dict)
-    current_step_index: Optional[int] = None
+    step_execution_tracker: dict[str, dict] = Field(default_factory=dict)
+    current_step_index: int | None = None
 
     max_steps: int = 20
 
@@ -113,7 +115,7 @@ async def get_plan(self) -> str:
         )
         return result.output if hasattr(result, "output") else str(result)
 
-    async def run(self, request: Optional[str] = None) -> str:
+    async def run(self, request: str | None = None) -> str:
         """Run the agent with an optional initial request."""
         if request:
             await self.create_initial_plan(request)
@@ -155,7 +157,7 @@ async def update_plan_status(self, tool_call_id: str) -> None:
         except Exception as e:
             logger.warning(f"Failed to update plan status: {e}")
 
-    async def _get_current_step_index(self) -> Optional[int]:
+    async def _get_current_step_index(self) -> int | None:
         """
         Parse the current plan to identify the first non-completed step's index.
         Returns None if no active step is found.
@@ -209,19 +211,22 @@ async def create_initial_plan(self, request: str) -> None:
         ]
         self.memory.add_messages(messages)
         response = await self.llm.ask_tool(
-            messages=messages,
+            messages=list(messages),
             system_msgs=[Message.system_message(self.system_prompt)],
             tools=self.available_tools.to_params(),
             tool_choice=ToolChoice.AUTO,
         )
+        if not response:
+            raise RuntimeError("Model refused to create plan")
+        tool_calls = response.tool_calls or []
         assistant_msg = Message.from_tool_calls(
-            content=response.content, tool_calls=response.tool_calls
+            content=response.content or "", tool_calls=tool_calls
         )
 
         self.memory.add_message(assistant_msg)
 
         plan_created = False
-        for tool_call in response.tool_calls:
+        for tool_call in tool_calls:
             if tool_call.function.name == "planning":
                 result = await self.execute_tool(tool_call)
                 logger.info(

app/agent/react.py

@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import Optional
 
 from pydantic import Field
 
@@ -10,12 +9,12 @@
 
 class ReActAgent(BaseAgent, ABC):
     name: str
-    description: Optional[str] = None
+    description: str | None = None
 
-    system_prompt: Optional[str] = None
-    next_step_prompt: Optional[str] = None
+    system_prompt: str | None = None
+    next_step_prompt: str | None = None
 
-    llm: Optional[LLM] = Field(default_factory=LLM)
+    llm: LLM = Field(default_factory=LLM)
     memory: Memory = Field(default_factory=Memory)
     state: AgentState = AgentState.IDLE
 

app/agent/swe.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from pydantic import Field
 
 from app.agent.toolcall import ToolCallAgent
@@ -19,7 +17,7 @@ class SWEAgent(ToolCallAgent):
     available_tools: ToolCollection = ToolCollection(
         Bash(), StrReplaceEditor(), Terminate()
     )
-    special_tool_names: List[str] = Field(default_factory=lambda: [Terminate().name])
+    special_tool_names: list[str] = Field(default_factory=lambda: [Terminate().name])
 
     max_steps: int = 30
 
@@ -29,7 +27,7 @@ class SWEAgent(ToolCallAgent):
     async def think(self) -> bool:
         """Process current state and decide next action"""
         # Update working directory
-        self.working_dir = await self.bash.execute("pwd")
+        self.working_dir = (await self.bash.execute("pwd")).output.strip()
         self.next_step_prompt = self.next_step_prompt.format(
             current_dir=self.working_dir
         )

app/agent/toolcall.py

@@ -1,6 +1,9 @@
 import json
-from typing import Any, List, Optional, Union
+from typing import Any, Sequence, TypeAlias
 
+from openai.types.chat.chat_completion_message_tool_call import (
+    ChatCompletionMessageToolCall,
+)
 from pydantic import Field
 
 from app.agent.react import ReActAgent
@@ -11,29 +14,33 @@
 from app.tool import CreateChatCompletion, Terminate, ToolCollection
 
 
+CompatibleToolCallObject: TypeAlias = ChatCompletionMessageToolCall
+
 TOOL_CALL_REQUIRED = "Tool calls required but none provided"
 
 
 class ToolCallAgent(ReActAgent):
     """Base agent class for handling tool/function calls with enhanced abstraction"""
 
-    name: str = "toolcall"
-    description: str = "an agent that can execute tool calls."
+    name = "toolcall"
+    description = "an agent that can execute tool calls."
 
-    system_prompt: str = SYSTEM_PROMPT
-    next_step_prompt: str = NEXT_STEP_PROMPT
+    system_prompt = SYSTEM_PROMPT
+    next_step_prompt = NEXT_STEP_PROMPT
 
     available_tools: ToolCollection = ToolCollection(
         CreateChatCompletion(), Terminate()
     )
     tool_choices: TOOL_CHOICE_TYPE = ToolChoice.AUTO  # type: ignore
-    special_tool_names: List[str] = Field(default_factory=lambda: [Terminate().name])
+    special_tool_names: list[str] = Field(default_factory=lambda: [Terminate().name])
 
-    tool_calls: List[ToolCall] = Field(default_factory=list)
-    _current_base64_image: Optional[str] = None
+    tool_calls: Sequence[ToolCall | CompatibleToolCallObject] = Field(
+        default_factory=list
+    )
+    _current_base64_image: str | None = None
 
     max_steps: int = 30
-    max_observe: Optional[Union[int, bool]] = None
+    max_observe: int | None = None
 
     async def think(self) -> bool:
         """Process current state and decide next actions using tools"""
@@ -44,7 +51,7 @@ async def think(self) -> bool:
         try:
             # Get response with tool options
             response = await self.llm.ask_tool(
-                messages=self.messages,
+                messages=list(self.messages),
                 system_msgs=(
                     [Message.system_message(self.system_prompt)]
                     if self.system_prompt
@@ -162,7 +169,7 @@ async def act(self) -> str:
 
         return "\n\n".join(results)
 
-    async def execute_tool(self, command: ToolCall) -> str:
+    async def execute_tool(self, command: ToolCall | CompatibleToolCallObject) -> str:
         """Execute a single tool call with robust error handling"""
         if not command or not command.function or not command.function.name:
             return "Error: Invalid command format"