| """Versioned REST API for the TP53 discovery workflow.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
| from uuid import uuid4 |
|
|
| import pandas as pd |
| from fastapi import BackgroundTasks, FastAPI, HTTPException, Query |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.staticfiles import StaticFiles |
| from pydantic import BaseModel, Field |
|
|
| from utils.calculations import ( |
| DEFAULT_WEIGHTS, |
| SCORING_VERSION, |
| apply_weights, |
| bmut_from_score, |
| consensus_score, |
| empirical_docking_score, |
| ingest_variant_frame, |
| md_metrics, |
| pocket_hydrophobicity, |
| pocket_records, |
| selectivity_from_scores, |
| structure_record, |
| ) |
| from api import store |
|
|
| DATA_DIR = Path(__file__).resolve().parent.parent / "data" |
|
|
| app = FastAPI( |
| title="TP53 Mutant Discovery API", |
| version="1.0.0", |
| description="Traceable variant-to-compound analysis with asynchronous jobs and versioned artifacts.", |
| ) |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["GET", "POST"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| class ProjectCreate(BaseModel): |
| name: str = Field(min_length=1) |
| cancer_type: str = "Lung Cancer" |
| metadata: dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class VariantIngest(BaseModel): |
| project_id: str |
| records: list[dict[str, Any]] |
|
|
|
|
| class VariantAnalyze(BaseModel): |
| project_id: str |
|
|
|
|
| class StructurePrepare(BaseModel): |
| project_id: str |
| variant_ids: list[str] |
| source_preference: str = "experimental_then_template" |
|
|
|
|
| class PocketDetect(BaseModel): |
| project_id: str |
| structure_ids: list[str] |
|
|
|
|
| class CompoundScreen(BaseModel): |
| project_id: str |
| variant_id: str |
| library_id: str = "platform-library-v1" |
| model_id: str = "complementarity-v1" |
| limit: int = Field(default=500, ge=1, le=2000) |
|
|
|
|
| class DockingRun(BaseModel): |
| project_id: str |
| variant_id: str |
| compounds: list[dict[str, Any]] |
| protocol_id: str = "matched-wt-mutant-v1" |
|
|
|
|
| class MDRun(BaseModel): |
| project_id: str |
| complexes: list[dict[str, Any]] |
| protocol_id: str = "explicit-solvent-3x100ns-v1" |
|
|
|
|
| class RankingCompute(BaseModel): |
| project_id: str |
| variant_id: str |
| candidates: list[dict[str, Any]] |
| weights: dict[str, float] = Field(default_factory=lambda: DEFAULT_WEIGHTS.copy()) |
|
|
|
|
| @app.on_event("startup") |
| def startup() -> None: |
| store.init_db() |
|
|
|
|
| def require_project(project_id: str) -> dict: |
| project = store.get_project(project_id) |
| if not project: |
| raise HTTPException(404, f"Project {project_id} not found") |
| return project |
|
|
|
|
| def submit(background: BackgroundTasks, project_id: str, operation: str, payload: dict) -> dict: |
| require_project(project_id) |
| job = store.create_job(project_id, operation, payload) |
| background.add_task(process_job, job["job_id"]) |
| return {"job_id": job["job_id"], "status": "queued", "status_url": f"/jobs/{job['job_id']}"} |
|
|
|
|
| @app.get("/health") |
| def health() -> dict: |
| return {"status": "ok", "version": app.version, "scoring_version": SCORING_VERSION} |
|
|
|
|
| @app.post("/projects", status_code=201) |
| def create_project(body: ProjectCreate) -> dict: |
| return store.create_project(body.name, body.cancer_type, body.metadata) |
|
|
|
|
| @app.get("/projects/{project_id}") |
| def get_project(project_id: str) -> dict: |
| return require_project(project_id) |
|
|
|
|
| @app.post("/variants:ingest", status_code=202) |
| def variants_ingest(body: VariantIngest, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "variants:ingest", body.model_dump()) |
|
|
|
|
| @app.post("/variants/{variant_id}/analyze", status_code=202) |
| def variant_analyze(variant_id: str, body: VariantAnalyze, background: BackgroundTasks) -> dict: |
| return submit( |
| background, |
| body.project_id, |
| "variant:analyze", |
| {**body.model_dump(), "variant_id": variant_id}, |
| ) |
|
|
|
|
| @app.post("/structures:prepare", status_code=202) |
| def structures_prepare(body: StructurePrepare, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "structures:prepare", body.model_dump()) |
|
|
|
|
| @app.post("/pockets:detect", status_code=202) |
| def pockets_detect(body: PocketDetect, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "pockets:detect", body.model_dump()) |
|
|
|
|
| @app.post("/compounds:screen", status_code=202) |
| def compounds_screen(body: CompoundScreen, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "compounds:screen", body.model_dump()) |
|
|
|
|
| @app.post("/docking:run", status_code=202) |
| def docking_run(body: DockingRun, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "docking:run", body.model_dump()) |
|
|
|
|
| @app.post("/md:run", status_code=202) |
| def md_run(body: MDRun, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "md:run", body.model_dump()) |
|
|
|
|
| @app.post("/rankings:compute", status_code=202) |
| def rankings_compute(body: RankingCompute, background: BackgroundTasks) -> dict: |
| return submit(background, body.project_id, "rankings:compute", body.model_dump()) |
|
|
|
|
| @app.get("/jobs/{job_id}") |
| def get_job(job_id: str) -> dict: |
| job = store.get_job(job_id) |
| if not job: |
| raise HTTPException(404, f"Job {job_id} not found") |
| return job |
|
|
|
|
| @app.get("/projects/{project_id}/jobs") |
| def project_jobs(project_id: str) -> list[dict]: |
| require_project(project_id) |
| return store.list_jobs(project_id) |
|
|
|
|
| @app.get("/projects/{project_id}/artifacts") |
| def project_artifacts(project_id: str) -> list[dict]: |
| require_project(project_id) |
| return store.list_artifacts(project_id) |
|
|
|
|
| @app.get("/artifacts/{artifact_id}") |
| def artifact(artifact_id: str) -> dict: |
| row = store.get_artifact(artifact_id) |
| if not row: |
| raise HTTPException(404, f"Artifact {artifact_id} not found") |
| return row |
|
|
|
|
| @app.get("/projects/{project_id}/report") |
| def project_report(project_id: str) -> dict: |
| project = require_project(project_id) |
| variants = store.list_records(project_id, "variant") |
| structures = store.list_records(project_id, "structure") |
| pockets = store.list_records(project_id, "pocket") |
| rankings = store.list_records(project_id, "ranking") |
| return { |
| "project": project, |
| "summary": { |
| "variants": len(variants), |
| "structures": len(structures), |
| "pockets": len(pockets), |
| "ranked_candidates": len(rankings), |
| }, |
| "top_candidates": sorted(rankings, key=lambda x: x.get("rescue_score", 0), reverse=True)[:10], |
| "jobs": store.list_jobs(project_id), |
| "artifacts": store.list_artifacts(project_id), |
| "versions": {"api": app.version, "scoring": SCORING_VERSION}, |
| } |
|
|
|
|
| def process_job(job_id: str) -> None: |
| job = store.get_job(job_id) |
| if not job: |
| return |
| try: |
| store.update_job(job_id, "running", log=f"Started {job['operation']}") |
| handler = HANDLERS[job["operation"]] |
| result = handler(job["project_id"], job["input"], job_id) |
| store.update_job(job_id, "succeeded", output=result, log="Completed successfully") |
| except Exception as exc: |
| store.update_job(job_id, "failed", error=str(exc), log=f"Failed: {exc}") |
|
|
|
|
| def handle_ingest(project_id: str, payload: dict, job_id: str) -> dict: |
| frame = pd.DataFrame(payload["records"]) |
| normalized = ingest_variant_frame(frame, require_project(project_id)["cancer_type"]) |
| records = [] |
| for _, row in normalized.iterrows(): |
| item = _jsonable(row.to_dict()) |
| item["variant_id"] = f"VAR-{uuid4().hex[:12]}" |
| records.append(item) |
| ids = store.add_records(project_id, "variant", records) |
| artifact = store.add_artifact( |
| project_id, "qc-report", "variant-qc.json", records, "application/json", job_id |
| ) |
| return {"variant_ids": ids, "normalized_records": records, "qc_artifact_id": artifact["artifact_id"]} |
|
|
|
|
| def handle_analyze(project_id: str, payload: dict, job_id: str) -> dict: |
| record = store.get_record(project_id, "variant", payload["variant_id"]) |
| if not record: |
| raise ValueError(f"Variant {payload['variant_id']} not found") |
| return { |
| "variant_id": payload["variant_id"], |
| "hgvs_p": record["hgvs_p"], |
| "priority": record["priority_score"], |
| "class": record["functional_class"], |
| "domain": record["domain"], |
| "route": record["route"], |
| "confidence": record["confidence"], |
| "qc_flags": record["qc_flags"], |
| } |
|
|
|
|
| def handle_structures(project_id: str, payload: dict, job_id: str) -> dict: |
| rows = [] |
| for variant_id in payload["variant_ids"]: |
| variant = store.get_record(project_id, "variant", variant_id) |
| if not variant: |
| raise ValueError(f"Variant {variant_id} not found") |
| rec = _jsonable(structure_record(pd.Series(variant))) |
| rec["structure_id"] = f"STR-{uuid4().hex[:12]}" |
| rec["variant_id"] = variant_id |
| rec["source_preference"] = payload["source_preference"] |
| rows.append(rec) |
| ids = store.add_records(project_id, "structure", rows) |
| artifact = store.add_artifact(project_id, "structures", "structures.json", rows, "application/json", job_id) |
| return {"structure_ids": ids, "structures": rows, "artifact_id": artifact["artifact_id"]} |
|
|
|
|
| def handle_pockets(project_id: str, payload: dict, job_id: str) -> dict: |
| rows = [] |
| for structure_id in payload["structure_ids"]: |
| structure = store.get_record(project_id, "structure", structure_id) |
| if not structure: |
| raise ValueError(f"Structure {structure_id} not found") |
| for rec in pocket_records(structure): |
| rec["pocket_id"] = f"POC-{uuid4().hex[:12]}" |
| rec["structure_id"] = structure_id |
| rows.append(_jsonable(rec)) |
| ids = store.add_records(project_id, "pocket", rows) |
| artifact = store.add_artifact(project_id, "pockets", "pockets.json", rows, "application/json", job_id) |
| return {"pocket_ids": ids, "pockets": rows, "artifact_id": artifact["artifact_id"]} |
|
|
|
|
| def handle_screen(project_id: str, payload: dict, job_id: str) -> dict: |
| variant = store.get_record(project_id, "variant", payload["variant_id"]) |
| if not variant: |
| raise ValueError(f"Variant {payload['variant_id']} not found") |
| source = pd.read_csv(DATA_DIR / "screening.csv") |
| rows = source[source["variant"] == variant["hgvs_p"]].nlargest(payload["limit"], "ai_score") |
| output = [_jsonable(row) for row in rows.to_dict("records")] |
| artifact = store.add_artifact(project_id, "screening", "screening.json", output, "application/json", job_id) |
| return { |
| "variant_id": payload["variant_id"], |
| "model_id": payload["model_id"], |
| "library_id": payload["library_id"], |
| "candidates": output, |
| "artifact_id": artifact["artifact_id"], |
| } |
|
|
|
|
| def handle_docking(project_id: str, payload: dict, job_id: str) -> dict: |
| variant = store.get_record(project_id, "variant", payload["variant_id"]) |
| if not variant: |
| raise ValueError(f"Variant {payload['variant_id']} not found") |
| structure = next( |
| (s for s in store.list_records(project_id, "structure") if s["variant_id"] == payload["variant_id"]), |
| None, |
| ) |
| if not structure: |
| structure = _jsonable(structure_record(pd.Series(variant))) |
| hyd = pocket_hydrophobicity(int(variant["position"]), variant["hgvs_p"]) |
| rows = [] |
| for compound in payload["compounds"]: |
| required = ("compound_id", "mw", "logp", "tpsa", "rotbonds") |
| missing = [key for key in required if key not in compound] |
| if missing: |
| raise ValueError(f"Compound missing {missing}") |
| dmut = empirical_docking_score( |
| compound["mw"], compound["logp"], compound["tpsa"], compound["rotbonds"], |
| structure.get("pocket_vol_mut") or 165, hyd, |
| ) |
| dwt = empirical_docking_score( |
| compound["mw"], compound["logp"], compound["tpsa"], compound["rotbonds"], |
| structure.get("pocket_vol_wt") or 165, 0.30, wt_pocket=True, |
| ) |
| rows.append({ |
| **compound, |
| "docking_id": f"DOCK-{uuid4().hex[:12]}", |
| "variant_id": payload["variant_id"], |
| "dock_mut": round(dmut, 3), |
| "dock_wt": round(dwt, 3), |
| "delta_dock": round(dmut - dwt, 3), |
| "mutant_preference": round(selectivity_from_scores(dmut, dwt), 3), |
| "protocol_version": payload["protocol_id"], |
| "units": "kcal/mol", |
| }) |
| ids = store.add_records(project_id, "docking", rows) |
| artifact = store.add_artifact(project_id, "docking", "docking.json", rows, "application/json", job_id) |
| return {"docking_ids": ids, "poses": rows, "artifact_id": artifact["artifact_id"]} |
|
|
|
|
| def handle_md(project_id: str, payload: dict, job_id: str) -> dict: |
| rows = [] |
| for complex_data in payload["complexes"]: |
| for key in ("complex_id", "dock_mut"): |
| if key not in complex_data: |
| raise ValueError(f"Complex missing {key}") |
| bmut = bmut_from_score(complex_data["dock_mut"]) |
| selectivity = float(complex_data.get("mutant_preference", 0)) |
| metrics = md_metrics( |
| complex_data["dock_mut"], complex_data.get("rotbonds", 4), bmut, selectivity |
| ) |
| rows.append({ |
| "md_id": f"MD-{uuid4().hex[:12]}", |
| **complex_data, |
| **metrics, |
| "replicas": 3, |
| "duration_ns": 100, |
| "protocol_version": payload["protocol_id"], |
| }) |
| ids = store.add_records(project_id, "md", rows) |
| artifact = store.add_artifact(project_id, "md", "md-summary.json", rows, "application/json", job_id) |
| return {"md_ids": ids, "results": rows, "artifact_id": artifact["artifact_id"]} |
|
|
|
|
| def handle_rankings(project_id: str, payload: dict, job_id: str) -> dict: |
| candidates = pd.DataFrame(payload["candidates"]) |
| required = list(DEFAULT_WEIGHTS) |
| missing = [col for col in required if col not in candidates.columns] |
| if missing: |
| raise ValueError(f"Candidates missing score components: {missing}") |
| if "variant" not in candidates: |
| candidates["variant"] = payload["variant_id"] |
| if "confidence" not in candidates: |
| candidates["confidence"] = 0.5 |
| if "admet_flag" not in candidates: |
| candidates["admet_flag"] = "Watch" |
| ranked = apply_weights(candidates, payload["weights"]) |
| rows = [] |
| for rec in ranked.to_dict("records"): |
| rec = _jsonable(rec) |
| rec["ranking_id"] = f"RNK-{uuid4().hex[:12]}" |
| rec["variant_id"] = payload["variant_id"] |
| rec["scoring_version"] = SCORING_VERSION |
| rows.append(rec) |
| ids = store.add_records(project_id, "ranking", rows) |
| artifact = store.add_artifact(project_id, "ranking", "ranking.json", rows, "application/json", job_id) |
| return {"ranking_ids": ids, "shortlist": rows[:10], "artifact_id": artifact["artifact_id"]} |
|
|
|
|
| HANDLERS = { |
| "variants:ingest": handle_ingest, |
| "variant:analyze": handle_analyze, |
| "structures:prepare": handle_structures, |
| "pockets:detect": handle_pockets, |
| "compounds:screen": handle_screen, |
| "docking:run": handle_docking, |
| "md:run": handle_md, |
| "rankings:compute": handle_rankings, |
| } |
|
|
|
|
| def _jsonable(value: Any) -> Any: |
| if isinstance(value, dict): |
| return {k: _jsonable(v) for k, v in value.items()} |
| if isinstance(value, list): |
| return [_jsonable(v) for v in value] |
| if pd.isna(value): |
| return None |
| if hasattr(value, "item"): |
| return value.item() |
| return value |
|
|
|
|
| |
| app.mount("/", StaticFiles(directory=Path(__file__).resolve().parent.parent, html=True), name="dashboard") |
|
|
