import time
from smolagents import CodeAgent, ActionStep

class RateLimiter:
    """
    APIの呼び出し回数を数えて、一定回数の呼び出しが行われたら一定秒数待機する
    """
    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):
        """
        呼び出し回数をインクリメントし、一定回数の呼び出しが行われたら一定秒数待機する
        """
        self._call_count += 1

        if self._call_count >= self.calls_per_wait:
            print(f"呼び出し回数が{self.calls_per_wait}回を超えたので、{self.seconds_to_wait}秒待機します")
            time.sleep(self.seconds_to_wait)
            self._call_count = 0