Update run.py

run.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import threading
 import sys
+import logging
 from io import StringIO
 from contextlib import redirect_stdout, redirect_stderr
 
@@ -36,6 +37,23 @@ AUTHORIZED_IMPORTS = [
 append_answer_lock = threading.Lock()
 
 
+class StreamingHandler(logging.Handler):
+    """Custom logging handler that captures agent logs"""
+    def __init__(self):
+        super().__init__()
+        self.callbacks = []
+        self.buffer = []
+    
+    def add_callback(self, callback):
+        self.callbacks.append(callback)
+    
+    def emit(self, record):
+        msg = self.format(record)
+        self.buffer.append(msg + '\n')
+        for callback in self.callbacks:
+            callback(msg + '\n')
+
+
 class StreamingCapture:
     """Captures stdout/stderr and yields content in real-time"""
     def __init__(self):
@@ -161,26 +179,55 @@ Additionally, if after some searching you find out that you need more informatio
 def run_agent_with_streaming(agent, question, stream_callback=None):
     """Run agent and stream output in real-time"""
     
-    # Capture stdout and stderr
-    stdout_capture = StreamingCapture()
-    stderr_capture = StreamingCapture()
-    
+    # Set up logging capture
+    log_handler = StreamingHandler()
     if stream_callback:
-        stdout_capture.add_callback(stream_callback)
-        stderr_capture.add_callback(stream_callback)
+        log_handler.add_callback(stream_callback)
+    
+    # Add handler to root logger and smolagents loggers
+    root_logger = logging.getLogger()
+    smolagents_logger = logging.getLogger('smolagents')
+    
+    # Store original handlers
+    original_handlers = root_logger.handlers[:]
+    original_level = root_logger.level
     
     try:
+        # Configure logging to capture everything
+        root_logger.setLevel(logging.DEBUG)
+        root_logger.addHandler(log_handler)
+        smolagents_logger.setLevel(logging.DEBUG)
+        smolagents_logger.addHandler(log_handler)
+        
+        # Also capture stdout/stderr
+        stdout_capture = StreamingCapture()
+        stderr_capture = StreamingCapture()
+        
+        if stream_callback:
+            stdout_capture.add_callback(stream_callback)
+            stderr_capture.add_callback(stream_callback)
+        
         with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
-            print(f"[STARTING] Running agent with question: {question}")
+            if stream_callback:
+                stream_callback(f"[STARTING] Running agent with question: {question}\n")
+            
             answer = agent.run(question)
-            print(f"[COMPLETED] Final answer: {answer}")
+            
+            if stream_callback:
+                stream_callback(f"[COMPLETED] Final answer: {answer}\n")
+            
             return answer
+            
     except Exception as e:
-        error_msg = f"[ERROR] Exception occurred: {str(e)}"
-        print(error_msg)
+        error_msg = f"[ERROR] Exception occurred: {str(e)}\n"
         if stream_callback:
             stream_callback(error_msg)
         raise
+    finally:
+        # Restore original logging configuration
+        root_logger.handlers = original_handlers
+        root_logger.setLevel(original_level)
+        smolagents_logger.removeHandler(log_handler)
 
 
 def create_gradio_interface():