app.py

# app.py
import gradio as gr
from openai import OpenAI

# ----------------------- CONSTANTS ----------------------- #
SYSTEM_PROMPT = """
Given the research context, design an ablation study for the specified module or process.
Begin the design with a clear statement of the research objective, followed by a detailed description of the experiment setup.
Do not include the discussion of results or conclusions in the response, as the focus is solely on the experimental design.
The response should be within 300 words. Present the response in **Markdown** format (use headings, bold text, and bullet or numbered lists where appropriate).
""".strip()

# ----------------------- HELPERS ------------------------- #

def prepare_user_prompt(
    research_background: str,
    method: str,
    experiment_setup: str,
    experiment_results: str,
    module_name: str,
) -> str:
    """Craft the ‘user’ portion of the OpenAI chat based on form inputs."""
    research_background_block = f"### Research Background\n{research_background}\n"
    method_block = f"### Method Section\n{method}\n"
    experiment_block = (
        "### Main Experiment Setup\n"
        f"{experiment_setup}\n\n"
        "### Main Experiment Results\n"
        f"{experiment_results}\n"
    )

    return (
        "## Research Context\n"
        f"{research_background_block}{method_block}{experiment_block}\n\n"
        f"Design an **ablation study** about **{module_name}** based on the research context above."
    )


def generate_ablation_design(
    research_background,
    method,
    experiment_setup,
    experiment_results,
    module_name,
    api_key,
):
    """Combine inputs ➜ call OpenAI ➜ return the ablation‑study design text (Markdown)."""
    # 1) Validate the API key format.
    if not api_key or not api_key.startswith("sk-"):
        return "❌ **Please enter a valid OpenAI API key in the textbox above.**"

    # 2) Build the chat conversation payload.
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": prepare_user_prompt(
                research_background,
                method,
                experiment_setup,
                experiment_results,
                module_name,
            ),
        },
    ]

    # 3) Call the model and return the assistant response.
    client = OpenAI(api_key=api_key)
    try:
        response = client.chat.completions.create(
            model="gpt-4.1",
            messages=messages,
            max_tokens=2048,
            temperature=1,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"⚠️  **OpenAI error:** {e}"

# ----------------------- UI LAYOUT ----------------------- #
with gr.Blocks(title="Ablation Study Designer") as demo:
    # Main two‑column layout.
    with gr.Row():
        # ---------- LEFT COLUMN: INPUTS ----------
        with gr.Column(scale=1):
            gr.Markdown(
                """
                # 🧪 Ablation Study Designer  
                Fill in the study details below, then click **Generate** to receive a tailored ablation‑study design rendered in Markdown.
                """,
                elem_id="header",
            )

            # API‑key field (required)
            api_key = gr.Textbox(
                label="🔑 OpenAI API Key (required)",
                type="password",
                placeholder="sk-...",
            )

            research_background = gr.Textbox(
                label="Research Background", lines=6, placeholder="Describe the broader research context…"
            )
            method = gr.Textbox(
                label="Method Description", lines=6, placeholder="Summarize the method section…"
            )
            experiment_setup = gr.Textbox(
                label="Main Experiment – Setup", lines=6, placeholder="Datasets, hyper‑parameters, etc."
            )
            experiment_results = gr.Textbox(
                label="Main Experiment – Results", lines=6, placeholder="Key quantitative or qualitative findings…"
            )
            module_name = gr.Textbox(
                label="Module / Process for Ablation", placeholder="e.g., Attention mechanism"
            )

            generate_btn = gr.Button("Generate Ablation Study Design")

        # ---------- RIGHT COLUMN: OUTPUT ----------
        with gr.Column(scale=1):
            output_md = gr.Markdown(value="", label="Ablation Study Design")

    # Button click: trigger generation with a loading indicator.
    generate_btn.click(
        fn=generate_ablation_design,
        inputs=[
            research_background,
            method,
            experiment_setup,
            experiment_results,
            module_name,
            api_key,
        ],
        outputs=output_md,
        show_progress="full",  # Display a full‑screen progress bar.
    )

# ----------------------- LAUNCH -------------------------- #
if __name__ == "__main__":
    demo.launch()