src/app.py

"""
Handbook Generator — Gradio UI (legacy fallback).
Primary UI is streamlit_app.py.
Run: python app.py
"""
import asyncio
import shutil
from pathlib import Path

import gradio as gr

from config import GROK_API_KEY, UPLOADS_DIR, BASE_DIR
from handbook_generator import build_handbook
from rag import get_context_for_query, index_pdf, reset_index

HANDBOOK_EXPORT_PATH = BASE_DIR / "handbook_export.md"


def _run_async(coro):
    """Run an async coroutine from sync Gradio code."""
    try:
        return asyncio.run(coro)
    except RuntimeError:
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, coro).result(timeout=300)


def ensure_api_key():
    if not GROK_API_KEY:
        raise gr.Error(
            "GROK_API_KEY is not set. Create a .env file in the ass2 folder with: GROK_API_KEY=your-key"
        )


def _file_path(f):
    """Get path from Gradio file input (path string or object with .name)."""
    if f is None:
        return None
    if isinstance(f, (str, Path)):
        return Path(f)
    return Path(getattr(f, "name", str(f)))


def upload_and_index(files):
    """Handle PDF upload(s) and index into LightRAG."""
    if not files:
        return "No files selected."
    ensure_api_key()
    reset_index()
    saved = []
    for f in (files if isinstance(files, list) else [files]):
        path = _file_path(f)
        if path is None or not path.exists():
            continue
        dest = UPLOADS_DIR / path.name
        try:
            shutil.copy(str(path), str(dest))
        except Exception:
            dest = path
        try:
            n = _run_async(index_pdf(dest, source_name=path.name))
            saved.append(f"{path.name}: indexed")
        except Exception as e:
            saved.append(f"{path.name}: Error - {e}")
    return "\n".join(saved) if saved else "No PDFs processed."


def chat(message, history):
    """RAG chat: retrieve context and answer using Grok via LiteLLM."""
    ensure_api_key()
    from litellm import completion
    from config import CHAT_MODEL

    context = _run_async(get_context_for_query(message))
    if not context or not context.strip():
        context = "No documents have been uploaded yet. Ask the user to upload PDFs first."

    system = (
        "You are a helpful assistant. Answer based ONLY on the following context "
        "from the user's uploaded documents. If the answer is not in the context, say so clearly."
    )
    user_content = f"Context from uploaded documents:\n\n{context}\n\n---\n\nUser question: {message}"

    resp = completion(
        model=CHAT_MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user_content},
        ],
        api_key=GROK_API_KEY,
        max_tokens=1500,
        temperature=0.3,
    )
    return (resp.choices[0].message.content or "").strip()


def run_handbook_simple(topic):
    """Generate handbook and return (status, markdown)."""
    ensure_api_key()
    if not (topic and topic.strip()):
        return "Enter a topic first.", ""
    status_msgs = []
    try:
        full_md = _run_async(build_handbook(topic.strip(), on_progress=status_msgs.append))
        status = "\n".join(status_msgs) if status_msgs else "Done."
        return status, full_md
    except Exception as e:
        return f"Error: {e}", ""


with gr.Blocks(title="Handbook Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Handbook Generator\nUpload PDFs, chat about them, and generate a 20,000+ word handbook.")

    with gr.Tab("Upload PDFs"):
        file_input = gr.File(
            file_count="multiple",
            file_types=[".pdf"],
            label="Upload one or more PDFs",
        )
        index_btn = gr.Button("Index PDFs")
        index_out = gr.Textbox(label="Index result", lines=4)

    with gr.Tab("Chat"):
        chatbot = gr.ChatInterface(
            fn=chat,
            type="messages",
            title="Ask questions about your uploaded documents",
        )

    with gr.Tab("Generate Handbook"):
        gr.Markdown(
            "Enter a topic (e.g. *Create a handbook on Retrieval-Augmented Generation*). "
            "Generation may take several minutes."
        )
        topic_in = gr.Textbox(
            label="Handbook topic",
            placeholder="e.g. Retrieval-Augmented Generation",
            lines=1,
        )
        gen_btn = gr.Button("Generate 20k-word handbook")
        status_out = gr.Textbox(label="Status", lines=4, interactive=False)
        handbook_out = gr.Markdown(label="Handbook (Markdown)")
        export_btn = gr.DownloadButton("Export as Markdown", visible=False)

    index_btn.click(
        fn=lambda files: upload_and_index(files) if files else "No files selected.",
        inputs=[file_input],
        outputs=[index_out],
    )

    def do_handbook(topic):
        status, md = run_handbook_simple(topic)
        if md:
            HANDBOOK_EXPORT_PATH.write_text(md, encoding="utf-8")
        return (
            status,
            md,
            gr.update(visible=bool(md), value=str(HANDBOOK_EXPORT_PATH) if md else None),
        )

    gen_btn.click(
        fn=do_handbook,
        inputs=[topic_in],
        outputs=[status_out, handbook_out, export_btn],
    )

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