Update run.py

run.py

@@ -1,6 +1,9 @@
 import argparse
 import os
 import threading
+import sys
+from io import StringIO
+from contextlib import redirect_stdout, redirect_stderr
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -33,6 +36,26 @@ AUTHORIZED_IMPORTS = [
 append_answer_lock = threading.Lock()
 
 
+class StreamingCapture:
+    """Captures stdout/stderr and yields content in real-time"""
+    def __init__(self):
+        self.content = []
+        self.callbacks = []
+    
+    def add_callback(self, callback):
+        self.callbacks.append(callback)
+    
+    def write(self, text):
+        if text.strip():
+            self.content.append(text)
+            for callback in self.callbacks:
+                callback(text)
+        return len(text)
+    
+    def flush(self):
+        pass
+
+
 def create_agent(
     model_id="gpt-4o-mini",
     hf_token=None,
@@ -135,12 +158,120 @@ Additionally, if after some searching you find out that you need more informatio
     return manager_agent
 
 
+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()
+    
+    if stream_callback:
+        stdout_capture.add_callback(stream_callback)
+        stderr_capture.add_callback(stream_callback)
+    
+    try:
+        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+            print(f"[STARTING] Running agent with question: {question}")
+            answer = agent.run(question)
+            print(f"[COMPLETED] Final answer: {answer}")
+            return answer
+    except Exception as e:
+        error_msg = f"[ERROR] Exception occurred: {str(e)}"
+        print(error_msg)
+        if stream_callback:
+            stream_callback(error_msg)
+        raise
+
+
+def create_gradio_interface():
+    """Create Gradio interface with streaming support"""
+    import gradio as gr
+    
+    def process_question(question, model_id, hf_token, serpapi_key, custom_api_endpoint, 
+                        custom_api_key, search_provider, search_api_key, custom_search_url):
+        
+        # Create agent
+        agent = create_agent(
+            model_id=model_id,
+            hf_token=hf_token,
+            openai_api_key=None,  # Add if needed
+            serpapi_key=serpapi_key,
+            api_endpoint=None,  # Add if needed
+            custom_api_endpoint=custom_api_endpoint,
+            custom_api_key=custom_api_key,
+            search_provider=search_provider,
+            search_api_key=search_api_key,
+            custom_search_url=custom_search_url,
+        )
+        
+        # Stream output
+        full_output = []
+        
+        def stream_callback(text):
+            full_output.append(text)
+            return "".join(full_output)
+        
+        # Generator function for streaming
+        def generate():
+            try:
+                answer = run_agent_with_streaming(agent, question, stream_callback)
+                yield "".join(full_output) + f"\n\n**FINAL ANSWER:** {answer}"
+            except Exception as e:
+                yield "".join(full_output) + f"\n\n**ERROR:** {str(e)}"
+        
+        return generate()
+    
+    # Create Gradio interface
+    with gr.Blocks(title="Streaming Agent Chat") as demo:
+        gr.Markdown("# Streaming Agent Chat Interface")
+        
+        with gr.Row():
+            with gr.Column():
+                question_input = gr.Textbox(label="Question", placeholder="Enter your question here...")
+                model_id_input = gr.Textbox(label="Model ID", value="gpt-4o-mini")
+                hf_token_input = gr.Textbox(label="HuggingFace Token", type="password")
+                serpapi_key_input = gr.Textbox(label="SerpAPI Key", type="password")
+                custom_api_endpoint_input = gr.Textbox(label="Custom API Endpoint")
+                custom_api_key_input = gr.Textbox(label="Custom API Key", type="password")
+                search_provider_input = gr.Dropdown(
+                    choices=["serper", "searxng"], 
+                    value="serper", 
+                    label="Search Provider"
+                )
+                search_api_key_input = gr.Textbox(label="Search API Key", type="password")
+                custom_search_url_input = gr.Textbox(label="Custom Search URL")
+                
+                submit_btn = gr.Button("Submit", variant="primary")
+            
+            with gr.Column():
+                output = gr.Textbox(
+                    label="Agent Output (Streaming)", 
+                    lines=30, 
+                    max_lines=50,
+                    interactive=False
+                )
+        
+        submit_btn.click(
+            fn=process_question,
+            inputs=[
+                question_input, model_id_input, hf_token_input, serpapi_key_input,
+                custom_api_endpoint_input, custom_api_key_input, search_provider_input,
+                search_api_key_input, custom_search_url_input
+            ],
+            outputs=output,
+            show_progress=True
+        )
+    
+    return demo
+
+
 def main():
     print("[DEBUG] Loading environment variables")
     load_dotenv(override=True)
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("question", type=str)
+    parser.add_argument("--gradio", action="store_true", help="Launch Gradio interface")
+    parser.add_argument("question", type=str, nargs='?', help="Question to ask (CLI mode)")
     parser.add_argument("--model-id", type=str, default="gpt-4o-mini")
     parser.add_argument("--hf-token", type=str, default=os.getenv("HF_TOKEN"))
     parser.add_argument("--serpapi-key", type=str, default=os.getenv("SERPAPI_API_KEY"))
@@ -153,24 +284,37 @@ def main():
 
     print("[DEBUG] CLI arguments parsed:", args)
 
-    agent = create_agent(
-        model_id=args.model_id,
-        hf_token=args.hf_token,
-        openai_api_token=openai_api_token,
-        serpapi_key=args.serpapi_key,
-        api_endpoint=api_endpoint,
-        custom_api_endpoint=args.custom_api_endpoint,
-        custom_api_key=args.custom_api_key,
-        search_provider=args.search_provider,
-        search_api_key=args.search_api_key,
-        custom_search_url=args.custom_search_url,
-    )
+    if args.gradio:
+        # Launch Gradio interface
+        demo = create_gradio_interface()
+        demo.launch(share=True)
+    else:
+        # CLI mode
+        if not args.question:
+            print("Error: Question required for CLI mode")
+            return
+        
+        agent = create_agent(
+            model_id=args.model_id,
+            hf_token=args.hf_token,
+            openai_api_key=None,  # Fix: was openai_api_token
+            serpapi_key=args.serpapi_key,
+            api_endpoint=None,  # Fix: was api_endpoint
+            custom_api_endpoint=args.custom_api_endpoint,
+            custom_api_key=args.custom_api_key,
+            search_provider=args.search_provider,
+            search_api_key=args.search_api_key,
+            custom_search_url=args.custom_search_url,
+        )
 
-    print("[DEBUG] Running agent...")
-    answer = agent.run(args.question)
-    print(f"Got this answer: {answer}")
+        print("[DEBUG] Running agent...")
+        
+        def print_stream(text):
+            print(text, end='', flush=True)
+        
+        answer = run_agent_with_streaming(agent, args.question, print_stream)
+        print(f"\n\nGot this answer: {answer}")
 
 
 if __name__ == "__main__":
-    main()
-
+    main()