Update src/agent.py

src/agent.py

@@ -89,73 +89,84 @@ class BasicAgent():
         self.expected_tokens_per_step = 10000
         self.max_retries = 3
         self.base_delay = 5
-        
-    def run(self, question: str, file_content: str = "", file_path: str = ""):
+
+
+    def _read_file(self, file_path: str) -> str:
+        if not os.path.exists(file_path):
+            print(f"File not found: {file_path}")
+            return ""
+        if file_path.endswith(".txt"):
+            with open(file_path, "r", encoding="utf-8") as f:
+                return f.read()
+        elif file_path.endswith(".docx"):
+            doc = docx.Document(file_path)
+            return "\n".join([p.text for p in doc.paragraphs])
+        else:
+            # For unsupported formats, return empty string
+            print(f"Unsupported file type: {file_path}")
+            return ""
+            
+    def run(self, question: str, file_content: str = "", file_path: str = "") -> str:
         final_answer = None
         retry_count = 0
 
-        question = f"Question: {question}"
+        # If file content is empty but file_path exists, read the file
+        if not file_content and file_path:
+            file_content = self._read_file(file_path)
+
+        context = ""
         if file_content:
             context = f"Story content:\n{file_content}"
         elif file_path:
             context = f"File path: {file_path}"
-        else:
-            context = ""
-        print(f"Starting Agent with question text: {question}, {context}")  
+
+        print(f"Starting Agent with question: {question}\nContext length: {len(context)} chars")
+
         while True:
             try:
-                final_input = f"{question}\n\n{context}"
-                # Run the agent
+                final_input = f"Question: {question}\n\n{context}"
                 steps = self.agent.run(final_input)
 
-                # If steps is a string, convert it to a single-item list
+                # Convert string steps to list
                 if isinstance(steps, str):
                     steps = [steps]
 
                 for step in steps:
-                    # Handle string steps
                     if isinstance(step, str):
                         final_answer = step
-                        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: {output}")
+                        final_answer = 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)
 
-                    # Capture the final answer from the final answer step
-                    if step_name == "FinalAnswerStep":
-                        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"Retrying in {delay:.1f}s ... ({e})")
+                    print(f"Retrying in {delay:.1f}s due to rate limit... ({e})")
                     time.sleep(delay)
                     retry_count += 1
                 else:
-                    print(f"Error: {e}")
+                    print(f"Error during agent run: {e}")
                     break
 
-        print(f"\nFinished agent run.\n{'='*60}")
-        print(f"Final Answer: {final_answer}\n")
+        # Ensure a valid answer is always returned
+        if not final_answer:
+            final_answer = "I am unable to answer"
+
+        print(f"Finished agent run. Final Answer: {final_answer}\n{'='*50}")
         return final_answer
 
-    def __call__(self, question: str, file_content: str = "", file_path: str = ""):
-        return self.run(question, file_content)
+    def __call__(self, question: str, file_content: str = "", file_path: str = "") -> str:
+        return self.run(question, file_content, file_path)