Update app.py

app.py

@@ -4,10 +4,11 @@
 
 import gradio.utils, os, sys
 import gradio as gr
+import threading
+import queue
+import time
 from agents.crew import run_crew
-from contextlib import contextmanager
 from huggingface_hub import whoami
-from io import StringIO
 from utils.utils import (
     DATASET_TYPE_GAIA,
     DATASET_TYPE_HLE,
@@ -21,39 +22,56 @@ SPACE_ID = os.environ.get("SPACE_ID")
 
 BASE_URL = f"https://huggingface.co/spaces/{SPACE_ID}/blob/main"
 
-# Console output capture
+# Streaming console output capture
 
-@contextmanager
-def capture_output():
-    old_stdout = sys.stdout
-    old_stderr = sys.stderr
-
-    class Tee:
-        def __init__(self, buffer, original):
-            self.buffer = buffer
-            self.original = original
-
-        def write(self, data):
-            self.buffer.write(data)
-            self.original.write(data)
-
-        def flush(self):
-            self.buffer.flush()
-            self.original.flush()
-
-    try:
-        stdout_buffer = StringIO()
-        stderr_buffer = StringIO()
-        sys.stdout = Tee(stdout_buffer, old_stdout)
-        sys.stderr = Tee(stderr_buffer, old_stderr)
-        yield stdout_buffer, stderr_buffer
-    finally:
-        sys.stdout = old_stdout
-        sys.stderr = old_stderr
+class StreamingCapture:
+    def __init__(self):
+        self.queue = queue.Queue()
+        self.old_stdout = None
+        self.old_stderr = None
+        self.stopped = False
+        
+    def start(self):
+        self.old_stdout = sys.stdout
+        self.old_stderr = sys.stderr
+        
+        parent = self
+        
+        class QueueWriter:
+            def __init__(self, original, q):
+                self.original = original
+                self.q = q
+                
+            def write(self, data):
+                if data:
+                    self.q.put(data)
+                    self.original.write(data)
+                    
+            def flush(self):
+                self.original.flush()
+        
+        sys.stdout = QueueWriter(self.old_stdout, self.queue)
+        sys.stderr = QueueWriter(self.old_stderr, self.queue)
+        
+    def stop(self):
+        self.stopped = True
+        if self.old_stdout:
+            sys.stdout = self.old_stdout
+        if self.old_stderr:
+            sys.stderr = self.old_stderr
+            
+    def get_new_output(self):
+        output = []
+        while not self.queue.empty():
+            try:
+                output.append(self.queue.get_nowait())
+            except queue.Empty:
+                break
+        return "".join(output)
 
 # MCP server functions
 
-def ask(oauth_token: gr.OAuthToken, question: str, openai_api_key: str, gemini_api_key: str, anthropic_api_key: str, file_name : str = ""):
+def ask(oauth_token: gr.OAuthToken, question: str, openai_api_key: str, gemini_api_key: str, anthropic_api_key: str, file_name: str = ""):
     """
     Ask General AI Assistant a question to answer.
     Args:
@@ -62,38 +80,44 @@ def ask(oauth_token: gr.OAuthToken, question: str, openai_api_key: str, gemini_a
         gemini_api_key (str): Gemini API key (always used)
         anthropic_api_key (str): Anthropic API key (only used by Stagehand tool)
         file_name (str): Optional file name
-    Returns:
-        tuple: (answer, console_logs) - The answer to the question and captured console output
+    Yields:
+        tuple: (answer, console_logs) - The answer to the question and captured console output (streamed)
     """
     if not question:
         msg = "Question is required."
         gr.Warning(msg)
-        return None, msg
+        yield None, msg
+        return
 
     if not openai_api_key:
         msg = "OpenAI API Key is required."
         gr.Warning(msg)
-        return None, msg
+        yield None, msg
+        return
     
     if not gemini_api_key:
         msg = "Gemini API Key is required."
         gr.Warning(msg)
-        return None, msg
+        yield None, msg
+        return
 
     if not anthropic_api_key:
         msg = "Anthropic API Key is required."
         gr.Warning(msg)
-        return None, msg
+        yield None, msg
+        return
 
     is_valid = validate_input(question, openai_api_key, gemini_api_key, anthropic_api_key)
     
     if not is_valid:
         msg = "Invalid input"
         gr.Warning(msg)
-        return None, msg
+        yield None, msg
+        return
     
-    if file_name:
-        file_name = f"files/{file_name}"
+    task_file_name = file_name
+    if task_file_name:
+        task_file_name = f"files/{task_file_name}"
     
     try:
         # API keys are provided thanks to sponsor credit.
@@ -106,18 +130,49 @@ def ask(oauth_token: gr.OAuthToken, question: str, openai_api_key: str, gemini_a
         os.environ["GEMINI_API_KEY"] = gemini_key
         os.environ["ANTHROPIC_API_KEY"] = anthropic_key
 
-        with capture_output() as (stdout, stderr):
-            answer = run_crew(question, file_name)
-            
-        stdout_logs = stdout.getvalue()
-        stderr_logs = stderr.getvalue()
-        console_logs = stdout_logs + (stderr_logs if stderr_logs else "")
+        # Streaming capture
+        capture = StreamingCapture()
+        capture.start()
+        
+        result = {"answer": None, "error": None}
+        all_logs = ""
+        
+        def run_task():
+            try:
+                result["answer"] = run_crew(question, task_file_name)
+            except Exception as e:
+                result["error"] = str(e)
         
-        return answer, console_logs
+        # Run in background thread
+        thread = threading.Thread(target=run_task)
+        thread.start()
+        
+        # Stream logs while waiting
+        while thread.is_alive():
+            new_output = capture.get_new_output()
+            if new_output:
+                all_logs += new_output
+                yield None, all_logs  # Update console logs, answer still None
+            time.sleep(0.1)  # Poll interval
+        
+        thread.join()
+        capture.stop()
+        
+        # Get any remaining output
+        new_output = capture.get_new_output()
+        if new_output:
+            all_logs += new_output
+        
+        if result["error"]:
+            gr.Warning(result["error"])
+            yield None, all_logs + f"\n\nError: {result['error']}"
+        else:
+            yield result["answer"], all_logs
+            
     except Exception as e:
         msg = str(e)
         gr.Warning(msg)
-        return None, f"Error: {msg}"
+        yield None, f"Error: {msg}"
 
 # Helper functions
 
@@ -247,12 +302,13 @@ with gr.Blocks(elem_classes=["full-width-app"]) as gaia:
                     max_lines=5
                 )
 
-        with gr.Accordion("Console Logs", open=False):
+        with gr.Accordion("Console Logs", open=True):
             console_logs = gr.Textbox(
                 label="Output",
                 interactive=False,
                 lines=25,
-                max_lines=25
+                max_lines=25,
+                autoscroll=True
             )
         
         with gr.Row():