import time
from smolagents import CodeAgent, ActionStep

class RateLimiter:
    """
    responsible for counting the number of API calls 
    and waiting for a certain number of seconds when a certain number of calls have been made.
    """
    def __init__(self, calls_per_wait: int = 10, seconds_to_wait: int = 60):
        self.calls_per_wait = calls_per_wait
        self.seconds_to_wait = seconds_to_wait
        self._call_count = 0

    def increment_and_sleep_if_needed(self, memory_step: ActionStep, agent: CodeAgent):
        """
        Increment the number of calls, and wait for a certain number of seconds 
        after a certain number of calls have been made.
        """
        self._call_count += 1

        if self._call_count >= self.calls_per_wait:
            print(f"Since the number of calls exceeded {self.calls_per_wait} times, wait for {self.seconds_to_wait} seconds")
            time.sleep(self.seconds_to_wait)
            self._call_count = 0