run.py

import argparse
import os
import threading
import sys
import logging
from io import StringIO
from contextlib import redirect_stdout, redirect_stderr

from dotenv import load_dotenv
from huggingface_hub import login
from scripts.text_inspector_tool import TextInspectorTool
from scripts.text_web_browser import (
    ArchiveSearchTool,
    FinderTool,
    FindNextTool,
    PageDownTool,
    PageUpTool,
    SimpleTextBrowser,
    VisitTool,
)
from scripts.visual_qa import visualizer

from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    LiteLLMModel,
    DuckDuckGoSearchTool,
    Tool,
)

AUTHORIZED_IMPORTS = [
    "shell_gpt", "sgpt", "openai", "requests", "zipfile", "os", "pandas", "numpy", "sympy", "json", "bs4",
    "pubchempy", "xml", "yahoo_finance", "Bio", "sklearn", "scipy", "pydub",
    "yaml", "string", "secrets", "io", "PIL", "chess", "PyPDF2", "pptx", "torch", "datetime", "fractions", "csv",
]

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):
        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,
    openai_api_key=None,
    serpapi_key=None,
    api_endpoint=None,
    custom_api_endpoint=None,
    custom_api_key=None,
    search_provider="serper",
    search_api_key=None,
    custom_search_url=None
):
    print("[DEBUG] Creating agent with model_id:", model_id)

    if hf_token:
        print("[DEBUG] Logging into HuggingFace")
        login(hf_token)

    model_params = {
        "model_id": model_id,
        "custom_role_conversions": {"tool-call": "assistant", "tool-response": "user"},
        "max_completion_tokens": 8192,
    }

    if model_id == "gpt-4o-mini":
        model_params["reasoning_effort"] = "high"

    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

    model = LiteLLMModel(**model_params)
    print("[DEBUG] Model initialized")

    text_limit = 100000
    user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"

    browser_config = {
        "viewport_size": 1024 * 5,
        "downloads_folder": "downloads_folder",
        "request_kwargs": {
            "headers": {"User-Agent": user_agent},
            "timeout": 300,
        },
        "serpapi_key": serpapi_key,
    }

    os.makedirs(f"./{browser_config['downloads_folder']}", exist_ok=True)
    browser = SimpleTextBrowser(**browser_config)
    print("[DEBUG] Browser initialized")

    # Correct tool selection
    if search_provider == "searxng":
        print("[DEBUG] Using SearxNG-compatible DuckDuckGoSearchTool with base_url override")
        search_tool = DuckDuckGoSearchTool()
        if custom_search_url:
            search_tool.base_url = custom_search_url  # Override default DuckDuckGo URL (only if supported)
    else:
        print("[DEBUG] Using default DuckDuckGoSearchTool for Serper/standard search")
        search_tool = DuckDuckGoSearchTool()

    WEB_TOOLS = [
        search_tool,
        VisitTool(browser),
        PageUpTool(browser),
        PageDownTool(browser),
        FinderTool(browser),
        FindNextTool(browser),
        ArchiveSearchTool(browser),
        TextInspectorTool(model, text_limit),
    ]

    text_webbrowser_agent = ToolCallingAgent(
        model=model,
        tools=WEB_TOOLS,
        max_steps=20,
        verbosity_level=3,
        planning_interval=4,
        name="search_agent",
        description="A team member that will search the internet to answer your question.",
        provide_run_summary=True,
    )

    text_webbrowser_agent.prompt_templates["managed_agent"]["task"] += """You can navigate to .txt online files.
If a non-html page is in another format, especially .pdf or a Youtube video, use tool 'inspect_file_as_text' to inspect it.
Additionally, if after some searching you find out that you need more information to answer the question, you can use `final_answer` with your request for clarification as argument to request for more information."""

    manager_agent = CodeAgent(
        model=model,
        tools=[visualizer, TextInspectorTool(model, text_limit)],
        max_steps=16,
        verbosity_level=3,
        additional_authorized_imports=AUTHORIZED_IMPORTS,
        planning_interval=4,
        managed_agents=[text_webbrowser_agent],
    )

    print("[DEBUG] Agent fully initialized")
    return manager_agent


def run_agent_with_streaming(agent, question, stream_callback=None):
    """Run agent and stream output in real-time"""
    
    # Set up logging capture
    log_handler = StreamingHandler()
    if 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):
            if stream_callback:
                stream_callback(f"[STARTING] Running agent with question: {question}\n")
            
            answer = agent.run(question)
            
            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"
        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():
    """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
        
        # 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():
    print("[DEBUG] Loading environment variables")
    load_dotenv(override=True)

    parser = argparse.ArgumentParser()
    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"))
    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="serper")
    parser.add_argument("--search-api-key", type=str, default=None)
    parser.add_argument("--custom-search-url", type=str, default=None)
    args = parser.parse_args()

    print("[DEBUG] CLI arguments parsed:", args)

    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...")
        
        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()