use smolagents' step_callbacks instead of litellm's callback

api_call_counter.py

@@ -1,11 +1,22 @@
-class APICallCounter:
+import time
+from smolagents import CodeAgent, ActionStep
+
+class RateLimiter:
     """
-    litellmのAPI呼び出し回数をカウントする
+    APIの呼び出し回数を数えて、一定回数の呼び出しが行われたら一定秒数待機する
     """
-    def __init__(self):
-        """呼び出し回数を初期化する"""
-        self._call_count = 0 # 呼び出し回数
+    def __init__(self, calls_per_wait: int = 1, seconds_to_wait: int = 10):
+        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
 
-    def reset(self):
-        """呼び出し回数をリセットする"""
-        self._call_count = 0
+        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