app.py

from __future__ import annotations

import json
import os
import sys
import textwrap
from typing import List

import gradio as gr

# Ensure the local src/ directory is on sys.path (for Hugging Face Spaces)
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.join(ROOT_DIR, "src")
if SRC_DIR not in sys.path:
    sys.path.append(SRC_DIR)

from agentic_multiwriter.state import AgentState, ResearchSnippet
from agentic_multiwriter.agents import (
    researcher_node,
    aggregator_node,
    writer_node,
    critic_node,
    formatter_node,
)
from agentic_multiwriter.tools import get_logger

logger = get_logger()


def _format_sources(snippets: List[ResearchSnippet]) -> str:
    if not snippets:
        return "No web sources were retrieved."

    lines = []
    for s in snippets:
        title = s["title"] or s["url"]
        url = s["url"]
        snippet = s["snippet"]
        if url:
            lines.append(f"- [{title}]({url})\n  \n  > {snippet}")
        else:
            lines.append(f"- {title}\n  \n  > {snippet}")
    return "\n\n".join(lines)


def generate(topic: str, mode: str, progress=gr.Progress()):
    """Gradio callback to run the pipeline step-by-step with progress."""
    topic = topic.strip()
    if not topic:
        return (
            "Please enter a topic.",
            "",
            "",
            "",
            "",
        )

    # Initial state
    state: AgentState = {
        "topic": topic,
        "mode": mode,
        "research_snippets": [],
        "outline": [],
        "draft": "",
        "revised_draft": "",
        "final_output": "",
        "meta": {},
    }

    # 1. Research
    progress(0.1, "Researching the web...")
    logger.info("UI: starting researcher_node")
    state = researcher_node(state)

    # 2. Aggregate
    progress(0.25, "Aggregating and cleaning snippets...")
    logger.info("UI: starting aggregator_node")
    state = aggregator_node(state)

    # 3. Write draft
    progress(0.5, "Writing first draft...")
    logger.info("UI: starting writer_node")
    state = writer_node(state)
    initial_draft = state.get("draft", "") or ""

    # 4. Critic / edit
    progress(0.7, "Reviewing and improving draft...")
    logger.info("UI: starting critic_node")
    state = critic_node(state)
    revised_draft = state.get("revised_draft", "") or initial_draft

    # 5. Format final output
    progress(0.9, f"Formatting final output as {mode}...")
    logger.info("UI: starting formatter_node")
    state = formatter_node(state)
    final_output = state.get("final_output", "") or revised_draft

    # 6. Prepare outline, meta, sources
    outline = state.get("outline", []) or []
    meta = state.get("meta", {}) or {}
    snippets = state.get("research_snippets", []) or []

    outline_text = "\n".join(f"- {item}" for item in outline)
    meta_text = json.dumps(meta, indent=2)
    sources_md = _format_sources(snippets)

    progress(1.0, "Done.")
    return final_output, initial_draft, revised_draft, sources_md, meta_text


def build_interface() -> gr.Blocks:
    with gr.Blocks(title="Agentic Multiwriter") as demo:
        gr.Markdown(
            """
            # 🧠 Agentic Multiwriter

            Multi-agent research & writing system built with **LangGraph**.

            1. Researches your topic on the web  
            2. Aggregates and cleans snippets  
            3. Writes a draft  
            4. Critiques and improves it  
            5. Formats it as a blog, research summary, or LinkedIn-style post
            """
        )

        with gr.Row():
            topic_input = gr.Textbox(
                label="Topic",
                placeholder="e.g. Future of agentic AI",
                lines=2,
            )
        mode_input = gr.Radio(
            choices=["blog", "research_summary", "linkedin_post"],
            value="blog",
            label="Output mode",
        )
        run_button = gr.Button("Generate", variant="primary")

        with gr.Tab("Final Output"):
            final_output_box = gr.Markdown(label="Final Output")

        with gr.Tab("Initial Draft (Writer)"):
            initial_draft_box = gr.Markdown(label="Initial Draft")

        with gr.Tab("Revised Draft (Critic)"):
            revised_draft_box = gr.Markdown(label="Revised Draft")

        with gr.Tab("Sources"):
            sources_box = gr.Markdown(label="Web Sources Used")

        with gr.Tab("Meta"):
            meta_box = gr.Textbox(label="Meta (timings, counts, etc.)", lines=10)

        run_button.click(
            fn=generate,
            inputs=[topic_input, mode_input],
            outputs=[
                final_output_box,
                initial_draft_box,
                revised_draft_box,
                sources_box,
                meta_box,
            ],
        )

        gr.Markdown(
            textwrap.dedent(
                """
                ---
                ⚠️ **Note**: First run may take longer while the model loads or if you are
                using a local model. On Spaces, this app uses an open-source model via
                Hugging Face Inference API.
                """
            )
        )

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860)