app.py

"""FreeCAD RAG Assistant — Gradio Blocks UI."""
import gradio as gr

from src.generate import generate_response
from src.retrieve import HybridRetriever, indices_ready
from src.config import HF_MODELS, DEFAULT_MODEL

# ── example queries ───────────────────────────────────────────────────────────

_EXAMPLES = [
    ["Create a parametric box width=50 height=30 depth=20 with 5mm fillets on all vertical edges"],
    ["Make a flange OD=80mm ID=40mm thickness=10mm with 4 M6 bolt holes on a 60mm PCD using PolarPattern"],
    ["Create a hex nut for M10 thread conforming to ISO 4032 dimensions"],
    ["Generate an L-bracket 60x40x4mm with two 6mm countersunk holes"],
    ["Create a 20mm-diameter shaft with M20 thread for 30mm at one end"],
    ["Make a parametric gear blank where the number of teeth is driven by a Spreadsheet cell"],
    ["Create a wine-glass shape by revolving a profile around the Z axis"],
    ["How do I add a coincident constraint between two endpoints in a Sketcher script?"],
    ["What is the topological naming problem and how should I avoid it in generated scripts?"],
    ["Linear pattern of 5 pockets along X with 15mm spacing"],
    ["Sweep a circle along a helical path to make a spring"],
    ["Create a Pad with a sketch containing an interior circular hole (multi-loop sketch)"],
]

_INDEX_WARNING = (
    "> **Index not found.** Run `python build_index.py --repo <path-to-freecad-docs>` "
    "to build the retrieval index before using this app."
)

# ── generation handler ────────────────────────────────────────────────────────

def run(
    prompt: str,
    use_bm25: bool,
    use_dense: bool,
    use_rerank: bool,
    top_n: int,
    model: str,
):
    if not prompt.strip():
        return "", "Please enter a request.", []

    if not indices_ready():
        return "", _INDEX_WARNING, []

    retriever  = HybridRetriever(
        use_bm25=use_bm25,
        use_dense=use_dense,
        use_rerank=use_rerank,
        top_n=top_n,
    )

    try:
        citations = retriever.retrieve(prompt)
    except Exception as exc:  # noqa: BLE001
        return "", f"Retrieval error: {exc}", []

    code, explain, err = generate_response(
        query=prompt,
        citations=citations,
        model=model,
    )

    if err:
        return "", f"**Error:** {err}", []

    chunk_rows = [
        [c.id, c.page_title, c.section, c.source_url, f"{c.score:.4f}"]
        for c in citations
    ]
    return code, explain, chunk_rows


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="FreeCAD RAG Assistant", analytics_enabled=False) as demo:
    gr.Markdown(
        "# FreeCAD Python Code Generator\n"
        "Describe a parametric part and get a complete, runnable FreeCAD 1.1 Python script "
        "retrieved from the official FreeCAD wiki documentation.\n\n"
        "> Source: [FreeCAD Wiki](https://wiki.freecad.org), CC-BY 3.0"
    )

    with gr.Row():
        # ── left column: inputs ───────────────────────────────────────────────
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Describe the part or ask a scripting question",
                lines=4,
                placeholder="Create a parametric flange with 4 M6 bolt holes on a 60mm PCD…",
            )

            with gr.Accordion("Retrieval settings", open=False):
                use_bm25   = gr.Checkbox(value=True,  label="Enable BM25 (keyword retrieval)")
                use_dense  = gr.Checkbox(value=True,  label="Enable dense retrieval (semantic)")
                use_rerank = gr.Checkbox(value=True,  label="Enable cross-encoder reranking")
                top_n      = gr.Slider(minimum=3, maximum=10, value=5, step=1,
                                       label="Final chunks passed to LLM (top-N)")
                model      = gr.Dropdown(
                    choices=HF_MODELS, value=DEFAULT_MODEL, label="HuggingFace model"
                )

            run_btn = gr.Button("Generate", variant="primary")

            gr.Examples(examples=_EXAMPLES, inputs=[prompt], label="Example queries", cache_examples=False)

        # ── right column: outputs ─────────────────────────────────────────────
        with gr.Column(scale=3):
            code_out    = gr.Code(label="Generated FreeCAD Python", language="python")
            explain_out = gr.Markdown(label="Explanation & citations")

            with gr.Accordion("Retrieved chunks", open=False):
                chunks_out = gr.Dataframe(
                    headers=["#", "Page", "Section", "URL", "Score"],
                    wrap=True,
                    label="Top retrieved chunks",
                )

    run_btn.click(
        fn=run,
        inputs=[prompt, use_bm25, use_dense, use_rerank, top_n, model],
        outputs=[code_out, explain_out, chunks_out],
    )


if __name__ == "__main__":
    demo.launch()