tabs/batch.py

"""
Batch tab — parallel get calls for up to 10 IDs.
"""

import gradio as gr
from core import config, runner
from core.formatter import format_result

BATCH_ENTITIES = [
    "gene", "variant", "article", "trial", "drug",
    "disease", "pathway", "protein", "adverse-event", "pgx",
]


def create_batch_tab(session_keys):
    """Build the Batch tab."""

    with gr.Tab("📦 Batch"):
        gr.Markdown(
            "## Batch Mode\n"
            "Run parallel `get` calls for up to 10 entity IDs in one command.\n"
            "Enter comma-separated IDs."
        )

        with gr.Row():
            entity = gr.Dropdown(
                choices=BATCH_ENTITIES,
                value="gene",
                label="Entity Type",
                scale=1,
            )
            ids = gr.Textbox(
                label="IDs (comma-separated, max 10)",
                placeholder="e.g., BRAF,TP53 or NCT02576665,NCT03715933",
                scale=3,
            )

        with gr.Row():
            sections = gr.Textbox(
                label="Sections (comma-separated, optional)",
                placeholder="e.g., pathways,interactions",
            )
            source = gr.Textbox(
                label="Source (optional)",
                placeholder="e.g., nci",
            )

        with gr.Row():
            no_cache = gr.Checkbox(label="Bypass cache", value=False)

        # Quick examples
        gr.Markdown("**Quick examples:**")
        with gr.Row():
            ex1 = gr.Button("Genes: BRAF,TP53", size="sm", variant="secondary")
            ex2 = gr.Button("Trials: NCT02576665,NCT03715933", size="sm", variant="secondary")
            ex3 = gr.Button("Variants: BRAF V600E, KRAS G12D", size="sm", variant="secondary")

        ex1.click(fn=lambda: ("gene", "BRAF,TP53"), outputs=[entity, ids])
        ex2.click(fn=lambda: ("trial", "NCT02576665,NCT03715933"), outputs=[entity, ids])
        ex3.click(fn=lambda: ("variant", "BRAF V600E,KRAS G12D"), outputs=[entity, ids])

        run_btn = gr.Button("📦 Run Batch", variant="primary")
        output_md = gr.Markdown(label="Results")
        with gr.Accordion("Raw JSON", open=False):
            output_json = gr.Code(language="json")

        def run_batch(ent, id_str, secs, src, skip_cache, keys):
            if not id_str.strip():
                raise gr.Error("Please enter at least one ID (comma-separated, max 10).")

            args = ["batch", ent, id_str.strip()]
            if secs.strip():
                args.extend(["--sections", secs.strip()])
            if src.strip():
                args.extend(["--source", src.strip()])

            env = config.build_env_overrides(keys)
            result = runner.run(args, json_mode=True, no_cache=skip_cache, env_overrides=env)
            if not result["success"]:
                raise gr.Error(f"BioMCP error: {result['error']}")
            md, js = format_result(result)
            return md, js

        run_btn.click(
            fn=run_batch,
            inputs=[entity, ids, sections, source, no_cache, session_keys],
            outputs=[output_md, output_json],
        )