import time
from smolagents import (
    CodeAgent,
    ToolCallingAgent,
)

class PausingCodeAgent(CodeAgent):
    """
    A custom CodeAgent that pauses for a specified duration after each LLM call.
    """
    def __init__(self, pause_duration=20, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.pause_duration = pause_duration

    def step(self, memory_step):
        """
        Executes a single step of the agent's reasoning, then pauses.
        """
        # Pause for the specified duration
        print(f"Pausing for {self.pause_duration} seconds after LLM call...")
        time.sleep(self.pause_duration)
        
        # Call the original step method to perform the LLM call and get the action
        result = super().step(memory_step)


        return result
    
class PausingToolCallingAgent(ToolCallingAgent):
    """
    A custom CodeAgent that pauses for a specified duration after each LLM call.
    """
    def __init__(self, pause_duration=20, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.pause_duration = pause_duration

    def step(self, memory_step):
        """
        Executes a single step of the agent's reasoning, then pauses.
        """
        # Call the original step method to perform the LLM call and get the action
        result = super().step(memory_step)

        # Pause for the specified duration
        print(f"Pausing for {self.pause_duration} seconds after LLM call...")
        time.sleep(self.pause_duration)

        return result