Update src/agent.py

src/agent.py

@@ -1,30 +1,31 @@
 from src.settings import Settings
 from smolagents import LiteLLMModel, ToolCallingAgent
-from tools import DuckDuckGoSearchTool, FinalAnswerTool
+from tools.FinalAnswerTool import FinalAnswerTool
 from src.utils import InputTokenRateLimiter
 import time
 import random
+from litellm import completion
 
 settings = Settings()
 class GaiaAgent():
     def __init__(self):
-
         self.model = LiteLLMModel(
-                model_name=settings.model_name,
-                api_key=settings.api_key,
-                temperature=0.1,
-                max_tokens=150,
-                system_prompt=settings.system_prompt,
-            ) 
-        
-        self.agent = ToolCallingAgent( 
-                        tools=[
-                            DuckDuckGoSearchTool(), 
-                            FinalAnswerTool()
-                    ],
-                model = self.model,
-                max_steps=5
-                )
+            model_id=settings.llm_model_id, 
+            api_key=settings.llm_api_key
+        )
+        self.agent = ToolCallingAgent(
+            tools=[
+                FinalAnswerTool(),
+                # TODO: MP3 interpretation
+            ],
+            max_steps=10,
+            planning_interval=5,
+            model=self.model
+        )
+        self.token_rate_limiter = InputTokenRateLimiter()
+        self.expected_tokens_per_step = 10000
+        self.max_retries = 3
+        self.base_delay = 5
         self.token_rate_limiter = InputTokenRateLimiter()
         self.expected_tokens_per_step = 10000
         self.max_retries = 3
@@ -40,10 +41,13 @@ class GaiaAgent():
         while True:
             try:
                 for step in self.agent.run(input_text):
+                    if isinstance(step, str):
+                        continue
                     step_name = step.__class__.__name__
-                    if step.output:
+                    output = getattr(step, "output", None)
+                    if output:
                         print(f"Step: {step_name} Output: {step.output}")
-                    print(f"Step: {step_name}")
+                    print(f"Step: {step_name} Output: {step.output}")
 
                     self.token_rate_limiter.maybe_wait(self.expected_tokens_per_step)
                     tokens_used = getattr(step, "token_usage", None)
@@ -71,5 +75,4 @@ class GaiaAgent():
 
         print(f"\nFinished agent run.\n{'='*60}")
         print(f"Final Answer: {final_answer}\n")
-        return final_answer
-
+        return final_answer