agent.py

import logging
import yaml
from pathlib import Path
from typing import List

from prompts import SYSTEM_PROMPT
from utils import extract_json_from_model_markdown_output
from langchain_openai import ChatOpenAI
from toolkit import Toolkit
from memory import Memory
from planner import Planner
from tenacity import retry, stop_after_attempt, wait_exponential
from datetime import datetime


logger = logging.getLogger(__name__)


class Agent():
    def __init__(self, 
                 llm: ChatOpenAI, 
                 toolkit: Toolkit, 
                 max_actions=5, 
                 cwd="",
                 planner:Planner =None):
        self.memory = Memory()
        self.toolkit = toolkit
        self.action_str = self.toolkit.get_prompt_description() 
        self.system_prompt = SYSTEM_PROMPT.format(func_signature=self.action_str, max_actions=max_actions)
        self.llm = llm
        self.cwd = cwd
        self.planner = planner
        
    @retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(min=1, max=10)
    )
    def call_llm(self):
        return self.llm.invoke(self.memory.get_messages()).content
        
        
    def create_state(self, task, prev_action, obs, step_num, bg_process_note):
        state = f"""
## Step {step_num}

### Goal
{task}

### Long-term Memory
{self.memory.get_working_memory()}

### Working Directory
{self.cwd}

### Previous Action
{prev_action}

### Observation
{obs}
""" 
        if(bg_process_note != '(none)'):
            state += """
### Background Process Output
{bg_process_note}
"""
        return state

    def execute(self, task, max_steps=100):
        task = f'{task}'
        if(self.planner is None):
            self.execute_single_task(task, max_steps=max_steps)
        else:
            ### first workout the plan
            ### then deal with the subtasks in the plan one by one
            logger.info('Decompose into procedures ...')
            subtasks = self.planner.plan(task)
            logger.info(f'Goal: {task}')
            for idx, subtask in enumerate(subtasks):
                logger.info(f"#{idx}: {subtask}")
                
            for task_idx, subtask in enumerate(subtasks):
                logger.info(f'Executing Subtask #{task_idx+1}: {subtask}...')
                goal = f"{task}\n\n### Current Task\n{subtask}\n"
                status = self.execute_single_task(goal, task_idx=(task_idx+1), max_steps=max_steps//5)
                logger.info(f'Subtask #{task_idx+1} Final Status: {status}...')   
    
    def execute_single_task(self, task, task_idx=-1, max_steps=100):
        self.memory.clear_task_related_memory()
        self.memory.add(role='system', content=self.system_prompt)
        ## only the first step will append the directory
        self.memory.add(role='user', content=self.create_state(task=task+f'\nInitial Working Directory: {self.cwd}', 
                                                               prev_action='(none)', 
                                                               obs='(none)', 
                                                               step_num=0,
                                                               bg_process_note='(none)'))
        logger.info(f"⚠️ {task}")
        task_status = 'Reach maximal steps. Forwarding to the next task.'
        for step_num in range(1, max_steps+1):
            if(task_idx < 0):
                logger.info('='*10 + f' Step {step_num}/{max_steps} ' + '='*10)
            else:
                logger.info('='*10 + f' Task {task_idx} Step {step_num}/{max_steps} ' + '='*10)
            output = self.call_llm()
            logger.info('+'*10 + f' Agent Output ' + '+'*10)
            logger.info(f'{output}')
            self.memory.add(role='assistant', content=output)
            ### parsing
            parsed_json = extract_json_from_model_markdown_output(output)
            parsed_json['raw_model_output'] = output
            actions = parsed_json['action']
            self.memory.add_working_memory(task_idx, step_num, parsed_json['current_state']['memory'])
            if(parsed_json['current_state']['cwd'] != ''):
                self.cwd = parsed_json['current_state']['cwd']
                self.cwd = self.cwd.replace("'",'').replace("`","")
                for k in self.toolkit.env_registry:
                    env = self.toolkit.env_registry[k]
                    if hasattr(env, "set_cwd") and callable(getattr(env, "set_cwd")):
                        env.set_cwd(self.cwd)
                        
            observations = []
            for i, action in enumerate(actions):
                result = self.toolkit.execute_action(action['action_name'], 
                                                        action['action_params'])
            
                observations.append(result)
                
                
            bg_process_note = '(none)'
            
            if('bash_env' in self.toolkit.env_registry and 
               self.toolkit.env_registry['bash_env'].last_bg_log is not None):
                bg_process_note = f"Execution output from {self.toolkit.env_registry['bash_env'].last_bg_log[0]}\n{open(self.toolkit.env_registry['bash_env'].last_bg_log[1]).read()}"
                self.toolkit.env_registry['bash_env'].last_bg_log = None
                
            state = self.create_state(task=task,
                                      prev_action=parsed_json['current_state']['next_goal'],
                                      obs='\n\n'.join(observations),
                                      step_num=step_num,
                                      bg_process_note=bg_process_note)
            if(action['action_name'] == 'done'):
                task_status = f"{action['action_params']}"
                break
            logger.info('+'*10 + f' Observation ' + '+'*10)
            logger.info(f'{state}')
            self.memory.add(role='user', content=state)
        return task_status
        
        


class AgentConfig:
    def __init__(
        self,
        api_base_url: str,
        model_name: str,
        api_key: str,
        toolkit: List[str],
        workspace_root: Path,
        use_planner: bool
    ):
        self.api_base_url = api_base_url
        self.model_name = model_name
        self.api_key = api_key
        self.toolkit = toolkit  # List of tool names
        self.workspace_root = workspace_root  # Path object
        self.use_planner = use_planner

    @classmethod
    def from_yaml(cls, path: Path) -> "AgentConfig":
        with open(path, "r", encoding="utf-8") as f:
            config_data = yaml.safe_load(f)

        required_fields = ["api_base_url", "api_key", "model_name", "toolkit", "workspace_root", "use_planner"]
        missing = [field for field in required_fields if field not in config_data]
        if missing:
            raise ValueError(f"Missing required config fields: {missing}")

        if not isinstance(config_data["toolkit"], list) or not all(
            isinstance(t, str) for t in config_data["toolkit"]
        ):
            raise ValueError("`toolkit` must be a list of strings")

        workspace_root = Path(config_data["workspace_root"]).expanduser().resolve()

        return cls(
            api_base_url=config_data["api_base_url"],
            model_name=config_data["model_name"],
            api_key=config_data["api_key"],
            toolkit=config_data["toolkit"],
            use_planner=config_data["use_planner"],
            workspace_root=workspace_root,
        )