Update src/agent.py

src/agent.py

@@ -0,0 +1,71 @@
+from settings import Settings
+from smolagents import LiteLLMModel, ToolCallingAgent
+from tools import DuckDuckGoSearchTool, FinalAnswerTool
+from src.utils import InputTokenRateLimiter
+import time
+import random
+
+settings = Settings()
+class GaiaAgent():
+    def __init__(self):
+
+        self.model = LiteLLMModel(
+                model_name=settings.model_name,
+                api_key=settings.api_key
+            ) 
+        
+        self.agent = ToolCallingAgent( 
+                        tools=[
+                            DuckDuckGoSearchTool(), 
+                            FinalAnswerTool()
+                    ],
+                model = self.model,
+                max_steps=5
+                )
+        self.token_rate_limiter = InputTokenRateLimiter()
+        self.expected_tokens_per_step = 10000
+        self.max_retries = 3
+        self.base_delay = 5
+        
+    def run(self, question: str, file_name: str = "", file_content: str = ""):
+        final_answer = None
+        retry_count = 0
+
+        input_text = f"Question: {question}\nFile Name: {file_name}\nFile Content: {file_content}"
+        print(f"Starting Agent with input text: {input_text}")  
+
+        while True:
+            try:
+                for step in self.agent.run(input_text):
+                    step_name = step.__class__.__name__
+                    if step.output:
+                        print(f"Step: {step_name} Output: {step.output}")
+                    print(f"Step: {step_name}")
+
+                    self.token_rate_limiter.maybe_wait(self.expected_tokens_per_step)
+                    tokens_used = getattr(step, "token_usage", None)
+                    if tokens_used:
+                            self.token_rate_limiter.add_tokens(tokens_used.input_tokens)
+
+                    if step_name == "FinalAnswerStep":
+                        final_answer = step.output
+                        print(f"Final Answer: {final_answer}")
+                break
+            except Exception as e:
+                if "overload" in str(e).lower() or "rate limit" in str(e).lower():
+                    print("Rate limit exceeded. Retrying...")
+                    if retry_count >= self.max_retries:
+                        print("Max retries reached. Exiting...")
+                        break
+                    delay = self.base_delay * (2 ** retry_count) + random.random()
+                    print(f"API overload/rate limit. Retrying in {delay:.1f}s ... ({e})")
+                    time.sleep(delay)
+                    retry_count += 1  
+
+                else:
+                    print(f"Error: {e}")
+                    break
+
+        print(f"\nFinished agent run.\n{'='*60}")
+        print(f"Final Answer: {final_answer}\n")
+        return final_answer