Create src/utils.py

src/utils.py

@@ -0,0 +1,51 @@
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from typing import Any
+
+SECONDS_IN_MINUTE = 60
+
+class BaseAgent(ABC):
+    @abstractmethod
+    def run(self, question: str, file_name: str = "", file_content: str = "") -> Any:
+        """Override this method in your agent implementation"""
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        return "BaseAgent default answer"
+
+class InputTokenRateLimiter:
+    """Ensures we don’t exceed tokens per minute limits for LLMs"""
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, max_tpm=50000):
+        self.max_tpm = max_tpm
+        if not hasattr(self, "_initialized"):
+            self.token_window = deque()  # stores (timestamp, tokens_used)
+            self._initialized = True
+
+    def _update_queue(self, time_now):
+        while self.token_window and time_now - self.token_window[0][0] > SECONDS_IN_MINUTE:
+            self.token_window.popleft()
+
+    def tokens_used_last_minute(self):
+        now = time.time()
+        self._update_queue(now)
+        return sum(tokens for _, tokens in self.token_window)
+
+    def maybe_wait(self, tokens_expected_to_use):
+        ctr = 0
+        while self.tokens_used_last_minute() + tokens_expected_to_use > self.max_tpm:
+            if ctr % 10 == 0:
+                print("Sleeping 0.5s to respect token limits...")
+            time.sleep(0.5)
+            self._update_queue(time.time())
+            ctr += 1
+
+    def add_tokens(self, tokens):
+        now = time.time()
+        self.token_window.append((now, tokens))
+        self._update_queue(now)