Update run.py

run.py

@@ -38,42 +38,40 @@ append_answer_lock = threading.Lock()
 
 
 class StreamingHandler(logging.Handler):
-    """Custom logging handler that captures agent logs and sends them to callbacks."""
+    """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)
-        # Check if the message is actually different or non-empty after stripping
-        # to avoid sending redundant empty strings, though `highlight_text` in app.py handles empty.
-        if msg.strip(): 
-            for callback in self.callbacks:
-                callback(msg + '\n') # Add newline to ensure distinct lines are processed by app.py's splitter
+        self.buffer.append(msg + '\n')
+        for callback in self.callbacks:
+            callback(msg + '\n')
 
 
-class StreamingCapture(StringIO):
-    """Captures stdout/stderr and sends content to callbacks in real-time."""
+class StreamingCapture:
+    """Captures stdout/stderr and yields content in real-time"""
     def __init__(self):
-        super().__init__()
+        self.content = []
         self.callbacks = []
     
     def add_callback(self, callback):
         self.callbacks.append(callback)
     
-    def write(self, s):
-        # Pass the raw string 's' directly to callbacks immediately
-        if s: # Only send if there's actual content
+    def write(self, text):
+        if text.strip():
+            self.content.append(text)
             for callback in self.callbacks:
-                callback(s)
-        super().write(s) # Still write to the underlying StringIO buffer
-
-
+                callback(text)
+        return len(text)
+    
     def flush(self):
-        super().flush()
+        pass
 
 
 def create_agent(
@@ -92,10 +90,7 @@ def create_agent(
 
     if hf_token:
         print("[DEBUG] Logging into HuggingFace")
-        try:
-            login(hf_token)
-        except Exception as e:
-            print(f"[ERROR] Failed to log into HuggingFace: {e}")
+        login(hf_token)
 
     model_params = {
         "model_id": model_id,
@@ -106,19 +101,10 @@ def create_agent(
     if model_id == "gpt-4o-mini":
         model_params["reasoning_effort"] = "high"
 
-    # Determine which API key to use based on the model_id
-    if "openai" in model_id.lower() and openai_api_key:
-        print("[DEBUG] Using OpenAI API key for OpenAI model")
-        model_params["api_key"] = openai_api_key
-    elif custom_api_endpoint and custom_api_key:
+    if custom_api_endpoint and custom_api_key:
         print("[DEBUG] Using custom API endpoint:", custom_api_endpoint)
         model_params["base_url"] = custom_api_endpoint
         model_params["api_key"] = custom_api_key
-    elif api_endpoint and openai_api_key: # Fallback to default OpenAI if custom not specified
-        print("[DEBUG] Using default API endpoint:", api_endpoint)
-        model_params["base_url"] = api_endpoint
-        model_params["api_key"] = openai_api_key
-    # It's important that if an API key is missing for the chosen model, it fails here or upstream.
 
     model = LiteLLMModel(**model_params)
     print("[DEBUG] Model initialized")
@@ -133,36 +119,23 @@ def create_agent(
             "headers": {"User-Agent": user_agent},
             "timeout": 300,
         },
-        "serpapi_key": serpapi_key, # This will be used by ArchiveSearchTool if SerpAPI is enabled
+        "serpapi_key": serpapi_key,
     }
 
     os.makedirs(f"./{browser_config['downloads_folder']}", exist_ok=True)
     browser = SimpleTextBrowser(**browser_config)
     print("[DEBUG] Browser initialized")
 
-    search_tool = None
+    # Correct tool selection
     if search_provider == "searxng":
-        print("[DEBUG] Using DuckDuckGoSearchTool (acting as a generic web search) for SearxNG context.")
+        print("[DEBUG] Using SearxNG-compatible DuckDuckGoSearchTool with base_url override")
         search_tool = DuckDuckGoSearchTool()
         if custom_search_url:
-            # Note: As mentioned before, DuckDuckGoSearchTool doesn't natively use a custom base_url
-            # for a completely different search engine like SearxNG. This line will likely have no effect.
-            # For true SearxNG integration, you'd need a custom tool or a modified DuckDuckGoSearchTool
-            # that knows how to query SearxNG instances.
-            print(f"[WARNING] DuckDuckGoSearchTool does not directly support 'custom_search_url' for SearxNG. Consider a dedicated SearxNG tool.")
-            # search_tool.base_url = custom_search_url # This line is often not effective for DDCSTool
-    elif search_provider == "serper":
-        print("[DEBUG] Using DuckDuckGoSearchTool (acting as a generic web search) for Serper context.")
-        search_tool = DuckDuckGoSearchTool() # You would need a separate SerperTool for direct Serper API calls.
-        if search_api_key:
-            print("[DEBUG] Serper API Key provided. Ensure your search tool (if custom) uses it.")
-            # If you had a dedicated SerperTool, you'd pass search_api_key to it.
-            # e.g., search_tool = SerperTool(api_key=search_api_key)
+            search_tool.base_url = custom_search_url  # Override default DuckDuckGo URL (only if supported)
     else:
-        print("[DEBUG] No specific search provider selected, or provider not directly supported. Defaulting to DuckDuckGoSearchTool.")
+        print("[DEBUG] Using default DuckDuckGoSearchTool for Serper/standard search")
         search_tool = DuckDuckGoSearchTool()
 
-
     WEB_TOOLS = [
         search_tool,
         VisitTool(browser),
@@ -170,15 +143,15 @@ def create_agent(
         PageDownTool(browser),
         FinderTool(browser),
         FindNextTool(browser),
-        ArchiveSearchTool(browser), # This tool specifically uses serpapi_key from browser_config
+        ArchiveSearchTool(browser),
         TextInspectorTool(model, text_limit),
     ]
 
     text_webbrowser_agent = ToolCallingAgent(
         model=model,
-        tools=[tool for tool in WEB_TOOLS if tool is not None], # Filter out None if search_tool was not set
+        tools=WEB_TOOLS,
         max_steps=20,
-        verbosity_level=3, # Keep this high for detailed output
+        verbosity_level=3,
         planning_interval=4,
         name="search_agent",
         description="A team member that will search the internet to answer your question.",
@@ -193,7 +166,7 @@ Additionally, if after some searching you find out that you need more informatio
         model=model,
         tools=[visualizer, TextInspectorTool(model, text_limit)],
         max_steps=16,
-        verbosity_level=3, # Keep this high for detailed output
+        verbosity_level=3,
         additional_authorized_imports=AUTHORIZED_IMPORTS,
         planning_interval=4,
         managed_agents=[text_webbrowser_agent],
@@ -215,49 +188,35 @@ def run_agent_with_streaming(agent, question, stream_callback=None):
     root_logger = logging.getLogger()
     smolagents_logger = logging.getLogger('smolagents')
     
-    # Store original handlers and levels
-    original_root_handlers = root_logger.handlers[:]
-    original_smolagents_handlers = smolagents_logger.handlers[:]
-    original_root_level = root_logger.level
-    original_smolagents_level = smolagents_logger.level
-
-    # Store original stdout/stderr
-    original_stdout = sys.stdout
-    original_stderr = sys.stderr
+    # Store original handlers
+    original_handlers = root_logger.handlers[:]
+    original_level = root_logger.level
     
-    stdout_capture = StreamingCapture()
-    stderr_capture = StreamingCapture()
-
-    if stream_callback:
-        stdout_capture.add_callback(stream_callback)
-        stderr_capture.add_callback(stream_callback)
-
     try:
         # Configure logging to capture everything
-        # Set logging levels very low to capture all verbose output
         root_logger.setLevel(logging.DEBUG)
-        for handler in root_logger.handlers: # Remove existing handlers to avoid duplicate output
-            root_logger.removeHandler(handler)
         root_logger.addHandler(log_handler)
-
         smolagents_logger.setLevel(logging.DEBUG)
-        for handler in smolagents_logger.handlers: # Remove existing handlers
-            smolagents_logger.removeHandler(handler)
         smolagents_logger.addHandler(log_handler)
         
-        # Redirect stdout/stderr
-        sys.stdout = stdout_capture
-        sys.stderr = stderr_capture
-            
-        if stream_callback:
-            stream_callback(f"[STARTING] Running agent with question: {question}\n")
-            
-        answer = agent.run(question)
+        # Also capture stdout/stderr
+        stdout_capture = StreamingCapture()
+        stderr_capture = StreamingCapture()
         
         if stream_callback:
-            stream_callback(f"[COMPLETED]  {answer}\n")
+            stdout_capture.add_callback(stream_callback)
+            stderr_capture.add_callback(stream_callback)
+        
+        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+            if stream_callback:
+                stream_callback(f"[STARTING] Running agent with question: {question}\n")
+            
+            answer = agent.run(question)
             
-        return 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)}\n"
@@ -266,18 +225,111 @@ def run_agent_with_streaming(agent, question, stream_callback=None):
         raise
     finally:
         # Restore original logging configuration
-        root_logger.handlers = original_root_handlers
-        root_logger.setLevel(original_root_level)
-        smolagents_logger.handlers = original_smolagents_handlers
-        smolagents_logger.setLevel(original_smolagents_level)
-
-        # Restore original stdout/stderr
-        sys.stdout = original_stdout
-        sys.stderr = original_stderr
+        root_logger.handlers = original_handlers
+        root_logger.setLevel(original_level)
+        smolagents_logger.removeHandler(log_handler)
+
+
+def create_gradio_interface():
+    """Create Gradio interface with streaming support"""
+    import gradio as gr
+    import time
+    import threading
+    
+    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,
+            serpapi_key=serpapi_key,
+            api_endpoint=None,
+            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,
+        )
+        
+        # Shared state for streaming
+        output_buffer = []
+        is_complete = False
+        
+        def stream_callback(text):
+            output_buffer.append(text)
+        
+        def run_agent_async():
+            nonlocal is_complete
+            try:
+                answer = run_agent_with_streaming(agent, question, stream_callback)
+                output_buffer.append(f"\n\n**FINAL ANSWER:** {answer}")
+            except Exception as e:
+                output_buffer.append(f"\n\n**ERROR:** {str(e)}")
+            finally:
+                is_complete = True
         
-        # Ensure any remaining buffered output is flushed (especially important for stdout/stderr)
-        stdout_capture.flush()
-        stderr_capture.flush()
+        # Start agent in background thread
+        agent_thread = threading.Thread(target=run_agent_async)
+        agent_thread.start()
+        
+        # Generator that yields updates
+        last_length = 0
+        while not is_complete or agent_thread.is_alive():
+            current_output = "".join(output_buffer)
+            if len(current_output) > last_length:
+                yield current_output
+                last_length = len(current_output)
+            time.sleep(0.1)  # Small delay to prevent excessive updates
+        
+        # Final yield to ensure everything is captured
+        final_output = "".join(output_buffer)
+        if len(final_output) > last_length:
+            yield final_output
+    
+    # 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():
@@ -290,20 +342,19 @@ def main():
     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"))
-    parser.add_argument("--openai-api-key", type=str, default=os.getenv("OPENAI_API_KEY")) # Added
-    parser.add_argument("--api-endpoint", type=str, default=os.getenv("API_ENDPOINT", "https://api.openai.com/v1")) # Added
     parser.add_argument("--custom-api-endpoint", type=str, default=None)
     parser.add_argument("--custom-api-key", type=str, default=None)
-    parser.add_argument("--search-provider", type=str, default="searxng") # Changed default to searxng for consistency
+    parser.add_argument("--search-provider", type=str, default="serper")
     parser.add_argument("--search-api-key", type=str, default=None)
-    parser.add_argument("--custom-search-url", type=str, default="https://search.endorisk.nl/search") # Changed default for consistency
+    parser.add_argument("--custom-search-url", type=str, default=None)
     args = parser.parse_args()
 
     print("[DEBUG] CLI arguments parsed:", args)
 
     if args.gradio:
-        print("Please run `app.py` directly to launch the Gradio interface.")
-        return
+        # Launch Gradio interface
+        demo = create_gradio_interface()
+        demo.launch(share=True)
     else:
         # CLI mode
         if not args.question:
@@ -313,9 +364,9 @@ def main():
         agent = create_agent(
             model_id=args.model_id,
             hf_token=args.hf_token,
-            openai_api_key=args.openai_api_key, 
+            openai_api_key=None,  # Fix: was openai_api_token
             serpapi_key=args.serpapi_key,
-            api_endpoint=args.api_endpoint,
+            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,