tabs/helpers.py

"""
Cross-Entity Helpers tab — pivot between related entities.
"""

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

# All helper commands: (source_entity, helper_verb, id_label, id_placeholder, description)
HELPERS = [
    ("variant", "trials", "Variant", "BRAF V600E", "Find trials for a variant"),
    ("variant", "articles", "Variant", "BRAF V600E", "Find articles about a variant"),
    ("variant", "oncokb", "Variant", "BRAF V600E", "OncoKB therapy evidence (needs ONCOKB_TOKEN)"),
    ("drug", "adverse-events", "Drug", "pembrolizumab", "Adverse events for a drug"),
    ("drug", "trials", "Drug", "pembrolizumab", "Trials for a drug"),
    ("disease", "trials", "Disease", "melanoma", "Trials for a disease"),
    ("disease", "drugs", "Disease", "melanoma", "Drugs for a disease"),
    ("disease", "articles", "Disease", "Lynch syndrome", "Articles about a disease"),
    ("gene", "trials", "Gene", "BRAF", "Trials involving a gene"),
    ("gene", "drugs", "Gene", "BRAF", "Drugs targeting a gene"),
    ("gene", "articles", "Gene", "BRCA1", "Articles about a gene"),
    ("gene", "pathways", "Gene", "BRAF", "Pathways involving a gene"),
    ("pathway", "drugs", "Pathway ID", "R-HSA-5673001", "Drugs related to a pathway"),
    ("pathway", "articles", "Pathway ID", "R-HSA-5673001", "Articles about a pathway"),
    ("pathway", "trials", "Pathway ID", "R-HSA-5673001", "Trials related to a pathway"),
    ("protein", "structures", "UniProt ID", "P15056", "Protein structures"),
    ("article", "entities", "PMID", "22663011", "Named entities from an article"),
    ("article", "citations", "PMID", "22663011", "Papers that cite this article"),
    ("article", "references", "PMID", "22663011", "References cited by this article"),
    ("article", "recommendations", "PMID", "22663011", "Recommended similar articles"),
]

HELPER_LABELS = [f"{h[0]} {h[1]} — {h[4]}" for h in HELPERS]


def create_helpers_tab(session_keys):
    """Build the Cross-Entity Helpers tab."""

    with gr.Tab("🔗 Cross-Entity Helpers"):
        gr.Markdown(
            "## Cross-Entity Pivots\n"
            "Pivot from one entity to another without rebuilding filters.\n"
            "Select a helper command, enter the ID, and go."
        )

        helper_choice = gr.Dropdown(
            choices=HELPER_LABELS,
            value=HELPER_LABELS[0],
            label="Helper Command",
        )

        with gr.Row():
            helper_id = gr.Textbox(
                label="ID / Name",
                placeholder="BRAF V600E",
                scale=3,
            )
            helper_limit = gr.Number(label="Limit", value=5, minimum=1, maximum=50, scale=1)

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

        # Update placeholder when helper changes
        def update_placeholder(choice):
            idx = HELPER_LABELS.index(choice) if choice in HELPER_LABELS else 0
            h = HELPERS[idx]
            return gr.Textbox(placeholder=f"e.g., {h[3]}", label=h[2])

        helper_choice.change(fn=update_placeholder, inputs=[helper_choice], outputs=[helper_id])

        # Quick examples
        gr.Markdown("**Quick examples:**")
        with gr.Row():
            ex1 = gr.Button('variant trials "BRAF V600E"', size="sm", variant="secondary")
            ex2 = gr.Button("gene drugs BRAF", size="sm", variant="secondary")
            ex3 = gr.Button("drug adverse-events pembrolizumab", size="sm", variant="secondary")
            ex4 = gr.Button("article citations 22663011", size="sm", variant="secondary")

        def set_ex(label, eid):
            return label, eid

        ex1.click(fn=lambda: set_ex(HELPER_LABELS[0], "BRAF V600E"), outputs=[helper_choice, helper_id])
        ex2.click(fn=lambda: set_ex(HELPER_LABELS[9], "BRAF"), outputs=[helper_choice, helper_id])
        ex3.click(fn=lambda: set_ex(HELPER_LABELS[3], "pembrolizumab"), outputs=[helper_choice, helper_id])
        ex4.click(fn=lambda: set_ex(HELPER_LABELS[17], "22663011"), outputs=[helper_choice, helper_id])

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

        def run_helper(choice, eid, lim, skip_cache, keys):
            if not eid.strip():
                raise gr.Error("Please enter an ID or name (e.g., BRAF V600E, pembrolizumab, 22663011).")

            idx = HELPER_LABELS.index(choice) if choice in HELPER_LABELS else 0
            h = HELPERS[idx]
            entity, verb = h[0], h[1]

            args = [entity, verb, eid.strip()]
            lim = int(lim) if lim else 5
            args.extend(["--limit", str(lim)])

            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_helper,
            inputs=[helper_choice, helper_id, helper_limit, no_cache, session_keys],
            outputs=[output_md, output_json],
        )