picogent-py/picogent/agent.py

"""
Main agent with ReAct loop
"""

import asyncio
import uuid
from typing import Dict, Any, List, Optional

from .config import Config
from .session import Session
from .context import ContextBuilder
from .providers.anthropic import AnthropicProvider
from .tools.registry import ToolRegistry
from .tools.read import ReadTool
from .tools.write import WriteTool
from .tools.edit import EditTool
from .tools.bash import BashTool


class Agent:
    """AI coding agent with ReAct loop"""
    
    def __init__(self, config: Config):
        self.config = config
        self.session = Session()
        self.context_builder = ContextBuilder(config.workspace)
        
        # Initialize provider
        if config.provider == "anthropic":
            self.provider = AnthropicProvider(config.api_key, config.model)
        else:
            raise ValueError(f"Unknown provider: {config.provider}")
        
        # Initialize tools
        self.tool_registry = ToolRegistry()
        self._register_default_tools()
    
    def _register_default_tools(self):
        """Register default tools"""
        self.tool_registry.register(ReadTool())
        self.tool_registry.register(WriteTool())
        self.tool_registry.register(EditTool())
        self.tool_registry.register(BashTool())
    
    async def run(self, user_message: str, session_file: Optional[str] = None) -> str:
        """Run the agent with ReAct loop"""
        # Load session if specified
        if session_file:
            self.session.session_file = session_file
            self.session.load()
        
        # Add user message to session
        self.session.add_message("user", user_message)
        
        # Build system prompt
        system_prompt = self.context_builder.build_system_prompt(self.config.system_prompt)
        
        # Add working directory info to system prompt
        system_prompt += "\n\n" + self.context_builder.get_working_directory_info()
        
        # ReAct loop
        iteration = 0
        max_iterations = self.config.max_iterations
        
        while iteration < max_iterations:
            iteration += 1
            
            # Get messages for API
            messages = self.session.get_anthropic_messages()
            
            # Get available tools
            tools = self.tool_registry.get_anthropic_tools()
            
            try:
                # Generate response
                response = await self.provider.generate_response(
                    messages=messages,
                    system_prompt=system_prompt,
                    tools=tools,
                    max_tokens=self.config.max_tokens
                )
                
                # Check if there are tool calls
                tool_calls = response.get("tool_calls", [])
                
                if not tool_calls:
                    # No tool calls, return text response
                    response_text = response.get("content", "").strip()
                    if response_text:
                        self.session.add_message("assistant", response_text)
                    
                    # Save session if file specified
                    if session_file:
                        self.session.save()
                    
                    return response_text or "I apologize, but I didn't generate a response."
                
                # Process tool calls
                # First, add the assistant message with tool calls
                if response.get("raw_content"):
                    self.session.add_message("assistant", response["raw_content"])
                else:
                    # Fallback for when we don't have raw content
                    self.session.add_message("assistant", f"I'll use the following tools: {[tc['name'] for tc in tool_calls]}")
                
                # Execute each tool call
                for tool_call in tool_calls:
                    tool_id = tool_call.get("id", str(uuid.uuid4()))
                    tool_name = tool_call.get("name")
                    tool_input = tool_call.get("input", {})
                    
                    try:
                        # Execute the tool
                        result = await self.tool_registry.execute_tool(tool_name, tool_input)
                        
                        # Add tool result to session
                        self.session.add_tool_result(tool_id, result)
                        
                    except Exception as e:
                        error_result = f"Error executing {tool_name}: {str(e)}"
                        self.session.add_tool_result(tool_id, error_result)
                
                # Continue the loop to get the next response
                
            except Exception as e:
                error_msg = f"Error during agent execution: {str(e)}"
                
                # Save session if file specified
                if session_file:
                    self.session.save()
                
                return error_msg
        
        # Max iterations reached
        final_message = f"Maximum iterations ({max_iterations}) reached. The agent may need more steps to complete the task."
        
        # Save session if file specified
        if session_file:
            self.session.save()
        
        return final_message
    
    async def chat(self, user_message: str) -> str:
        """Simple chat interface without session persistence"""
        return await self.run(user_message)
    
    def clear_session(self):
        """Clear the current session"""
        self.session.clear()
    
    def get_session_messages(self) -> List[Dict[str, Any]]:
        """Get all messages in the current session"""
        return self.session.get_messages()