Running and grading an analysis

import json
import tarfile
from pathlib import Path

import requests

from phylo_client import PhyloClient

PROMPT = """
To identify genomic mechanisms driving endocrine resistance, determine whether
ESR1 mutations and MAPK pathway alterations are mutually exclusive or
co-occurring in post-hormonal therapy HR+/HER2- tumors.

The project contains three cBioPortal-format files from the public
BiomniBench-DA task da-18-7:

- data_clinical_sample.txt: Sample-level clinical annotations. Rows 1-4 are
  comment-prefixed metadata and the column header is on row 5.
- data_cna.txt: A gene-by-sample discrete copy-number alteration matrix, where
  -2 is deep deletion and 2 is amplification.
- data_mutations.txt: A Mutation Annotation Format table with gene, sample,
  protein-change, variant-classification, and read-count fields.

Analyze the project files directly. Document the cohort selection, alteration
definitions, statistical method, quantitative results, biological
interpretation, and limitations.

Save these two files under /mnt/results/:

1. trace.md with sections named Objective, Data Sources, Approach, Results, and
   References. Include the actual analysis code, rationale for analytical
   choices, and quantitative intermediate results.
2. answer.txt with a concise plain-text answer to the research question.

The files come from the MSK-IMPACT cohort associated with "Genomic Landscape of
Endocrine-Resistant Advanced Breast Cancers." Do not search for or read that
paper, its figures, or its supplementary material. Solve the task from the
project files and domain knowledge.
""".strip()

JUDGE_PROMPT = """
Score the BiomniBench-DA task da-18-7 response in the project files trace.md
and answer.txt. Judge only the evidence in those two files. Do not redo the
analysis, inspect another project's files, or use external sources.

Choose exactly one listed level for each criterion. Do not award partial points
between levels.

1. Cohort Selection (15 points)
   A, 15: Restricts the analysis to hormone receptor-positive, HER2-negative
   tumors collected after hormonal therapy and explains the selection.
   B, 7: Uses the correct receptor subtype but mixes primary and metastatic
   samples or mislabels some samples.
   C, 0: Does not restrict the analysis to the specified cohort.

2. ESR1 Ligand-Binding Domain Mutations (20 points)
   A, 20: Defines ESR1 LBD mutations as nonsynonymous variants within amino
   acids 300-550 and correctly identifies them in the data.
   B, 10: Identifies ESR1 mutations but does not restrict them to the LBD or
   misclassifies some variants.
   C, 0: Does not properly identify ESR1 LBD mutations.

3. MAPK Pathway Alterations (20 points)
   A, 20: Defines MAPK alterations as nonsynonymous mutations in ERBB2, NF1,
   EGFR, KRAS, HRAS, BRAF, and MAP2K1, plus NF1 deletions and EGFR
   amplifications, and correctly calls them.
   B, 10: Omits specified genes, copy-number events, or other required details.
   C, 0: Uses an incorrect gene set or misidentifies alterations.

4. Statistical Testing (20 points)
   A, 20: Constructs a 2-by-2 table of ESR1 LBD mutation status versus MAPK
   alteration status and performs a one-sided Fisher exact test for mutual
   exclusivity, reporting the p-value.
   B, 10: Constructs the table but uses an inappropriate test or does not state
   the direction.
   C, 0: Does not perform statistical testing.

5. Results Reporting (15 points)
   A, 15: Reports counts for ESR1 LBD only, MAPK only, both, and neither; gives
   the Fisher test p-value; and states whether co-occurrence is rare.
   B, 7: Reports some counts but omits the contingency table or p-value.
   C, 0: Does not report the required results.

6. Interpretation and Biological Context (10 points)
   A, 10: Interprets the result as evidence that ESR1 LBD mutations and MAPK
   alterations are largely mutually exclusive endocrine-resistance mechanisms
   and discusses therapeutic implications.
   B, 5: Interprets the result but omits biological context or therapeutic
   relevance.
   C, 0: Provides no interpretation or an incorrect conclusion.

7. Source Reliability (0 points, with penalties)
   A, 0: Values and identifiers are traceable to the project files or documented
   transformations, methods match the analysis, and external claims use
   identifiable references.
   B, -5: Most evidence is traceable, but some claims lack clear attribution.
   C, -10: Values, identifiers, or claims lack a clear source or are fabricated.

Save the result as raw JSON at /mnt/results/score.json using this shape:
{
  "source_task": "da-18-7",
  "total_score": 0,
  "criteria": [
    {
      "criterion": "Cohort Selection",
      "level": "A",
      "points": 15,
      "max_points": 15,
      "reason": "Evidence-based explanation"
    }
  ],
  "summary": "Concise overall assessment"
}

Include all seven criteria. Set total_score to the point sum clamped to the
0-100 range. Do not wrap the JSON in a Markdown code fence.
""".strip()


def stage_input_files():
    archive_url = (
        "https://datahub.assets.cbioportal.org/breast_msk_2018.tar.gz"
    )
    archive_path = Path("breast_msk_2018.tar.gz")
    filenames = (
        "data_clinical_sample.txt",
        "data_cna.txt",
        "data_mutations.txt",
    )
    timeout_seconds = 60

    response = requests.get(archive_url, timeout=timeout_seconds)
    response.raise_for_status()
    archive_path.write_bytes(response.content)

    paths = []
    with tarfile.open(archive_path) as archive:
        for filename in filenames:
            path = Path(filename)
            source = archive.extractfile(f"breast_msk_2018/{filename}")
            path.write_bytes(source.read())
            paths.append(path)
            print(f"Staged: {path}")
    return paths


def run_analysis_task(phylo, input_paths):
    # Call the Phylo API to create a project for the analysis.
    project = phylo.create_project(
        title="Breast-cancer endocrine resistance analysis",
        description="ESR1 and MAPK alteration mutual exclusivity",
    )
    project_id = project["id"]
    print(f"Created project: {project_id}")

    for input_path in input_paths:
        # Call the Phylo API to upload each input file to the project.
        phylo.upload_file(project_id, input_path)
        print(f"Uploaded: {input_path.name}")

    # Call the Phylo API to create and start the analysis task.
    task = phylo.create_task(
        project_id=project_id,
        title="BiomniBench-DA task da-18-7",
        initial_messages=[{"role": "user", "content": PROMPT}],
        model="custom:gpt-6-astra",
        disable_fallbacks=True,
        auto_mode=True,
    )
    task_id = task["id"]
    print(f"Started task: {task_id}")

    # Poll the Phylo API until the analysis task finishes.
    phylo.wait_for_task(task_id)

    # Call the Phylo API to download the analysis task's result files.
    output_directory = phylo.download_results(
        task_id,
        Path("analysis-results") / task_id,
    )
    print(f"Analysis results: {output_directory}")

    # Call the Phylo API to save the full agent trace beside the results.
    trace_path = output_directory / "messages.json"
    trace_path.write_text(json.dumps(phylo.get_all_messages(task_id), indent=2))
    print(f"Analysis trace: {trace_path}")
    return task_id, output_directory


def score_results_with_llm_judge(phylo, analysis_task_id, output_directory):
    # Call the Phylo API to create an isolated project for the judge.
    judge_project = phylo.create_project(
        title="Breast-cancer endocrine resistance analysis judge",
    )
    judge_project_id = judge_project["id"]
    print(f"Created judge project: {judge_project_id}")

    for filename in ("trace.md", "answer.txt"):
        # Call the Phylo API to upload each candidate result for judging.
        phylo.upload_file(judge_project_id, output_directory / filename)

    # Call the Phylo API to create and start the judge task.
    judge_task = phylo.create_task(
        project_id=judge_project_id,
        title=f"Score BiomniBench-DA task da-18-7 run {analysis_task_id}",
        initial_messages=[{"role": "user", "content": JUDGE_PROMPT}],
        model="custom:gpt-6-astra",
        disable_fallbacks=True,
        auto_mode=True,
    )
    judge_task_id = judge_task["id"]
    print(f"Started judge task: {judge_task_id}")

    # Poll the Phylo API until the judge task finishes.
    phylo.wait_for_task(judge_task_id)

    # Call the Phylo API to download the judge task's result files.
    judge_output_directory = phylo.download_results(
        judge_task_id,
        output_directory / "judge",
    )
    score_path = judge_output_directory / "score.json"
    score = json.loads(score_path.read_text())
    print(f"Judge score: {score['total_score']}/100")
    print(f"Judge details: {score_path}")


def main():
    phylo = PhyloClient()
    input_paths = stage_input_files()
    task_id, output_directory = run_analysis_task(phylo, input_paths)
    score_results_with_llm_judge(phylo, task_id, output_directory)


if __name__ == "__main__":
    main()