app.py

import gradio as gr
import asyncio
import os
from dotenv import load_dotenv
from backend.app.core.orchestrator import ResearchManager
from backend.app.core import ResearchMode
from backend.app.data import init_db
from datetime import datetime
from pathlib import Path


load_dotenv(override=True)


def is_valid_openai_key(key: str) -> tuple[bool, str]:
    if not key:
        return (False, "You must insert an OpenAI API key.")
    if not key.startswith("sk-"):
        return (False, 'You must insert a valid OpenAI API key (it starts with "sk-").')
    return (True, "")


DEFAULT_OPENAI_KEY = os.environ.get("OPENAI_API_KEY", "")
DEFAULT_SERPER_KEY = os.environ.get("SERPER_API_KEY", "")
DEFAULT_SENDGRID_KEY = os.environ.get("SENDGRID_API_KEY", "")
DEFAULT_SENDGRID_FROM = os.environ.get("SENDGRID_FROM", "")
DEFAULT_SENDGRID_TO = os.environ.get("SENDGRID_TO", "")

print("Initializing database...")
try:
    init_db()
    print("✓ Database initialized successfully")
except Exception as e:
    print(f"⚠️ Database initialization failed: {e}")
    print("   Research will continue but reports won't be saved")

manager = ResearchManager()

last_report = {"content": None, "query": None}


async def run_research(
    query: str,
    mode: str,
    send_email: bool,
    search_provider: str,
    openai_key: str,
    serper_key: str,
    sendgrid_key: str,
    sendgrid_from: str,
    sendgrid_to: str
):
    global last_report

    original_env = {
        "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY"),
        "SERPER_API_KEY": os.environ.get("SERPER_API_KEY"),
        "SENDGRID_API_KEY": os.environ.get("SENDGRID_API_KEY"),
        "SENDGRID_FROM": os.environ.get("SENDGRID_FROM"),
        "SENDGRID_TO": os.environ.get("SENDGRID_TO"),
    }

    try:
        if openai_key:
            os.environ["OPENAI_API_KEY"] = openai_key
        if serper_key:
            os.environ["SERPER_API_KEY"] = serper_key

        if send_email:
            if sendgrid_key:
                os.environ["SENDGRID_API_KEY"] = sendgrid_key
            if sendgrid_from:
                os.environ["SENDGRID_FROM"] = sendgrid_from
            if sendgrid_to:
                os.environ["SENDGRID_TO"] = sendgrid_to
        else:
            os.environ.pop("SENDGRID_API_KEY", None)

        os.environ["SEARCH_PROVIDER"] = search_provider.lower()

        research_mode = ResearchMode.QUICK if mode == "Quick" else ResearchMode.DEEP

        full_output = ""
        async for chunk in manager.run(query, mode=research_mode):
            full_output = chunk
            yield chunk

        if full_output and len(full_output) > 100:
            last_report["content"] = full_output
            last_report["query"] = query

    finally:
        for key, value in original_env.items():
            if value is not None:
                os.environ[key] = value
            else:
                os.environ.pop(key, None)


async def run_research_with_status(
    query: str,
    mode: str,
    send_email: bool,
    search_provider: str,
    openai_key: str,
    serper_key: str,
    sendgrid_key: str,
    sendgrid_from: str,
    sendgrid_to: str
):
    yield "🔄 Running research...", ""

    try:
        async for chunk in run_research(
            query, mode, send_email, search_provider,
            openai_key, serper_key, sendgrid_key,
            sendgrid_from, sendgrid_to
        ):
            yield "🔄 Running research...", chunk

        yield "✅ Research complete", chunk

    except asyncio.CancelledError:
        yield "⚠️ Stopped by user", chunk

    except Exception as e:
        yield f"❌ Error: {type(e).__name__}", chunk


async def run_research_with_validation(
    query: str,
    mode: str,
    send_email: bool,
    search_provider: str,
    openai_key: str,
    serper_key: str,
    sendgrid_key: str,
    sendgrid_from: str,
    sendgrid_to: str
):
    is_valid, error_msg = is_valid_openai_key(openai_key)

    if not is_valid:
        yield ("❌ API Key Error", error_msg)
        return

    async for status, report in run_research_with_status(
        query, mode, send_email, search_provider,
        openai_key, serper_key, sendgrid_key,
        sendgrid_from, sendgrid_to
    ):
        yield status, report


def stop_research() -> tuple[str, str]:
    manager.request_stop()
    return "⚠️ Stopping...", "Stopping research at next checkpoint..."


def export_report() -> str:
    global last_report

    if not last_report["content"]:
        return None

    exports_dir = Path("exports")
    exports_dir.mkdir(exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    query_slug = last_report["query"][:50].replace(" ", "_").replace("/", "_") if last_report["query"] else "report"
    query_slug = "".join(c for c in query_slug if c.isalnum() or c in ('_', '-'))
    filename = f"research_{query_slug}_{timestamp}.md"
    filepath = exports_dir / filename

    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(last_report["content"])
        print(f"Report exported to: {filepath}")
        return str(filepath)
    except Exception as e:
        print(f"Export failed: {e}")
        return None


with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as ui:
    gr.Markdown("# Deep Research")
    gr.Markdown("""
    **Research Modes:**
    - **Quick**: Faster research with 4-6 sources (~2 minutes)
    - **Deep**: Comprehensive research with 10-14 sources (~8 minutes)
    """)

    query_textbox = gr.Textbox(
        label="What topic would you like to research?",
        placeholder="Enter your research query..."
    )

    mode_radio = gr.Radio(
        choices=["Quick", "Deep"],
        value="Quick",
        label="Research Mode",
        info="Choose between quick research (4-6 sources) or deep research (10-14 sources)"
    )

    with gr.Accordion("⚙️ Configuration", open=False):
        gr.Markdown("### API Keys")

        openai_key_input = gr.Textbox(
            label="OpenAI API Key",
            placeholder="sk-...",
            value=DEFAULT_OPENAI_KEY,
            type="password",
            info="Required for AI agents and OpenAI web search"
        )

        gr.Markdown("### Web Search Provider")
        search_provider_radio = gr.Radio(
            choices=["OpenAI", "Serper"],
            value="OpenAI",
            label="Search Provider",
            info="Choose which search API to use"
        )

        serper_key_input = gr.Textbox(
            label="Serper API Key",
            placeholder="Enter Serper API key...",
            value=DEFAULT_SERPER_KEY,
            type="password",
            visible=False,
            info="Required when using Serper search provider"
        )

        gr.Markdown("### Email Settings")
        send_email_checkbox = gr.Checkbox(
            label="Send Email Report",
            value=False,
            info="Send the completed report via email"
        )

        sendgrid_key_input = gr.Textbox(
            label="SendGrid API Key",
            placeholder="SG...",
            value=DEFAULT_SENDGRID_KEY,
            type="password",
            visible=False,
            info="Required for email delivery"
        )

        sendgrid_from_input = gr.Textbox(
            label="From Email",
            placeholder="research@example.com",
            value=DEFAULT_SENDGRID_FROM,
            visible=False,
            info="Sender email (must be verified in SendGrid)"
        )

        sendgrid_to_input = gr.Textbox(
            label="To Email",
            placeholder="recipient@example.com",
            value=DEFAULT_SENDGRID_TO,
            visible=False,
            info="Recipient email address"
        )

    with gr.Row():
        run_button = gr.Button("Run", variant="primary")
        stop_button = gr.Button("Stop", variant="stop")
        export_button = gr.Button("📥 Export", variant="secondary")

    status_box = gr.Textbox(
        label="Status",
        value="Ready",
        interactive=False,
        max_lines=1
    )

    report = gr.Markdown(label="Report")

    export_file = gr.File(label="Exported Report", visible=False)

    def update_serper_visibility(provider):
        return gr.update(visible=(provider == "Serper"))

    def update_email_fields_visibility(send_email):
        return [gr.update(visible=send_email)] * 3

    search_provider_radio.change(
        fn=update_serper_visibility,
        inputs=[search_provider_radio],
        outputs=[serper_key_input]
    )

    send_email_checkbox.change(
        fn=update_email_fields_visibility,
        inputs=[send_email_checkbox],
        outputs=[sendgrid_key_input, sendgrid_from_input, sendgrid_to_input]
    )

    run_inputs = [
        query_textbox,
        mode_radio,
        send_email_checkbox,
        search_provider_radio,
        openai_key_input,
        serper_key_input,
        sendgrid_key_input,
        sendgrid_from_input,
        sendgrid_to_input
    ]

    run_button.click(fn=run_research_with_validation, inputs=run_inputs, outputs=[status_box, report])
    query_textbox.submit(fn=run_research_with_validation, inputs=run_inputs, outputs=[status_box, report])

    stop_button.click(fn=stop_research, outputs=[status_box, report])

    export_button.click(fn=export_report, outputs=export_file)

ui.launch(inbrowser=True)