Update src/agent.py

src/agent.py

@@ -40,35 +40,50 @@ class BasicAgent():
 
         while True:
             try:
-                for step in self.agent.run(input_text):
+                # Run the agent
+                steps = self.agent.run(input_text)
+
+                # If steps is a string, convert it to a single-item list
+                if isinstance(steps, str):
+                    steps = [steps]
+
+                for step in steps:
+                    # Handle string steps
                     if isinstance(step, str):
                         final_answer = step
-                        print(f"Final Answer (string step): {final_answer}")
+                        print(f"Step: String Output: {final_answer}")
                         continue
+
+                    # Handle object steps
                     step_name = step.__class__.__name__
                     output = getattr(step, "output", None)
+
                     if output:
-                        print(f"Step: {step_name} Output: {step.output}")
+                        print(f"Step: {step_name} Output: {output}")
+
                     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)
+                        self.token_rate_limiter.add_tokens(tokens_used.input_tokens)
 
+                    # Capture the final answer from the final answer step
                     if step_name == "FinalAnswerStep":
-                        final_answer = step.output
-                        print(f"Final Answer (FinalAnswerStep): {final_answer}")
-                break
+                        final_answer = output
+                        print(f"Captured Final Answer from step: {final_answer}")
+
+                break  # Exit retry loop if successful
+
             except Exception as e:
+                # Handle API overload/rate limits
                 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})")
+                    print(f"Retrying in {delay:.1f}s ... ({e})")
                     time.sleep(delay)
-                    retry_count += 1  
-
+                    retry_count += 1
                 else:
                     print(f"Error: {e}")
                     break