app/gradio_app.py

"""Gradio chat UI — Hugging Face Spaces deploy target.

Locally:  python app/gradio_app.py
HF Spaces: set app_file: app/gradio_app.py in README front-matter.
"""
from __future__ import annotations

import sys
from pathlib import Path

import gradio as gr

sys.path.append(str(Path(__file__).resolve().parents[1]))

_assistant = None


def _get_assistant():
    global _assistant
    if _assistant is None:
        from src.rag.generator import CodeAssistant
        _assistant = CodeAssistant.from_config(with_index=True)
    return _assistant


def respond(message: str, history, mode: str) -> str:
    assistant = _get_assistant()

    result = assistant.generate(message, mode=mode, max_new_tokens=64, return_sources=True)
    if isinstance(result, tuple):
        code, sources = result
        src_lines = "\n".join(
            f"- `{s['score']:.2f}` {s['docstring']}" for s in sources
        )
        return f"```python\n{code}\n```\n\n**Retrieved examples**\n{src_lines}"
    return f"```python\n{result}\n```"


with gr.Blocks(theme=gr.themes.Soft(), title="Code Generation Assistant") as demo:
    gr.Markdown(
        "# Code Generation Assistant\n"
        "Writes Python from a natural-language request, grounded in retrieved "
        "CodeSearchNet examples (RAG)."
    )
    with gr.Row():
        mode = gr.Radio(
            ["rag", "baseline"],
            value="rag",
            label="Mode",
            info="RAG conditions the model on retrieved CodeSearchNet examples; baseline is zero-shot.",
        )
    gr.ChatInterface(
        fn=respond,
        additional_inputs=[mode],
        examples=[
            ["Write a function to check if a number is prime."],
            ["Parse a CSV file and return a list of dictionaries."],
            ["Reverse a singly linked list."],
            ["Compute the Levenshtein distance between two strings."],
        ],
        cache_examples=False,
    )

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