shulapy's picture
final-agent-code (#1)
6e31f04 verified
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