apps/web/demo.py

from __future__ import annotations

import json
import os
from functools import lru_cache
from pathlib import Path

import gradio as gr

from run_demo_pipeline import main as run_demo_pipeline
from services.solve.solver import MathSolver
from services.verify.verifier import verify_solution

INDEX_PATH = Path(os.getenv("INDEX_PATH", "index/hybrid_index.pkl"))
SOLUTIONS_PATH = Path(os.getenv("SOLUTIONS_PATH", "data/normalized/solutions.jsonl"))
_ASSETS_READY = False


def _is_assets_ready() -> bool:
    return INDEX_PATH.exists() and SOLUTIONS_PATH.exists()


def _ensure_assets_ready() -> None:
    global _ASSETS_READY
    if _ASSETS_READY or _is_assets_ready():
        _ASSETS_READY = True
        return
    run_demo_pipeline()
    _ASSETS_READY = _is_assets_ready()
    if not _ASSETS_READY:
        raise RuntimeError(
            "Pipeline ran but required assets are still missing: "
            f"{INDEX_PATH} and/or {SOLUTIONS_PATH}."
        )


def _load_canonical() -> dict:
    if not SOLUTIONS_PATH.exists():
        return {}
    rows = [json.loads(line) for line in SOLUTIONS_PATH.read_text().splitlines()]
    return {r["question_id"]: r for r in rows}


@lru_cache(maxsize=1)
def _get_solver() -> MathSolver:
    _ensure_assets_ready()
    return MathSolver(str(INDEX_PATH))


def solve_question(question_text: str, question_id: str = "") -> str:
    if not question_text.strip():
        return "Please provide a question."

    solver = _get_solver()
    result = solver.solve(question_text.strip())
    canonical_map = _load_canonical()

    report_text = "Verification skipped (no question_id supplied)."
    if question_id and question_id in canonical_map:
        report = verify_solution(
            {"question_text": question_text},
            canonical_map[question_id],
            result.answer,
        )
        report_text = (
            f"passed={report.passed}, mp_coverage={report.mp_coverage:.2f}, "
            f"final_answer_correct={report.final_answer_correct}, "
            f"format={report.format_compliant}"
        )

    retrieved = "\n\n".join(
        (f"[{r['chunk'].chunk_type}] score={r['score']:.3f}\n{r['chunk'].text[:200]}")
        for r in result.retrieved
    )
    return (
        "## Plan\n- "
        + "\n- ".join(result.plan)
        + "\n\n## Answer\n"
        + result.answer
        + "\n\n## Verification\n"
        + report_text
        + "\n\n## Retrieved\n"
        + retrieved
    )


def build_app() -> gr.Blocks:
    with gr.Blocks(title="A-Level Math RAG Solver (HF)") as demo:
        gr.Markdown("# A-Level Math Exam Answering Demo (RAG + Tools + Verifier)")
        qid = gr.Textbox(label="Question ID (optional)", placeholder="e.g., 2022-P1-Q4a")
        question = gr.Textbox(label="Question", lines=5)
        output = gr.Markdown(label="Output")
        run = gr.Button("Solve")
        run.click(fn=solve_question, inputs=[question, qid], outputs=output)
    return demo


def get_launch_config() -> dict:
    return {
        "server_name": os.getenv("HOST", "0.0.0.0"),
        "server_port": int(os.getenv("PORT", "7860")),
    }


if __name__ == "__main__":
    build_app().launch(**get_launch_config())