| """Offline-first NVIDIA BioNeMo scientific-tool adapter. |
| |
| Only implements *preparation* and *response checking* for Boltz-2 NIM. |
| NO HTTP client, API key handling, Docker launch, or GPU execution here. |
| Never mistake schema validation for biological confirmation. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import math |
| import re |
| from dataclasses import dataclass |
| from typing import Any |
|
|
| from .models import Run |
|
|
| BIONEMO_REPO = "https://github.com/NVIDIA-BioNeMo/bionemo-agent-toolkit" |
| BOLTZ2_URL = "https://health.api.nvidia.com/v1/biology/mit/boltz2/predict" |
| SCHEMA = "frontierlab.bionemo.preflight.v1" |
| AA = frozenset("ACDEFGHIKLMNPQRSTVWY") |
|
|
| SKILLS = ( |
| {"name": "boltz2-nim", "task": "biomolecular structure prediction", |
| "stage": "OFFLINE_PREFLIGHT_ONLY", |
| "source": BIONEMO_REPO + "/tree/main/nim-skills/boltz2-nim"}, |
| {"name": "proteinmpnn-nim", "task": "sequence design for a known backbone", |
| "stage": "NOT_CONNECTED", |
| "source": BIONEMO_REPO + "/tree/main/nim-skills/proteinmpnn-nim"}, |
| {"name": "rfdiffusion-nim", "task": "protein backbone generation", |
| "stage": "NOT_CONNECTED", |
| "source": BIONEMO_REPO + "/tree/main/nim-skills/rfdiffusion-nim"}, |
| {"name": "openfold3-nim", "task": "biomolecular structure prediction", |
| "stage": "NOT_CONNECTED", |
| "source": BIONEMO_REPO + "/tree/main/nim-skills/openfold3-nim"}, |
| ) |
|
|
| def registry() -> list[dict[str, str]]: |
| """Tool availability is explicit; no endpoint is called or probed.""" |
| return [dict(tool) for tool in SKILLS] |
|
|
|
|
| def _canonical(payload: Any) -> bytes: |
| return json.dumps(payload, ensure_ascii=True, sort_keys=True, |
| separators=(",", ":"), allow_nan=False).encode("utf-8") |
|
|
|
|
| def _sha256(payload: bytes) -> str: |
| return hashlib.sha256(payload).hexdigest() |
|
|
|
|
| def _validate_sequence(sequence: str) -> str: |
| if not isinstance(sequence, str) or not 8 <= len(sequence) <= 4096: |
| raise ValueError("Protein sequence must have 8–4096 residues") |
| if any(residue not in AA for residue in sequence): |
| raise ValueError("Only uppercase canonical 20-residue protein sequences are supported") |
| return sequence |
|
|
|
|
| def _validate_id(source_id: str) -> str: |
| if not isinstance(source_id, str) or not re.fullmatch(r"[A-Za-z0-9:._/-]{3,160}", source_id): |
| raise ValueError("Source ID must be a nonempty plain identifier") |
| return source_id |
|
|
|
|
| @dataclass(frozen=True) |
| class Boltz2Preflight: |
| """Internal request contains the sequence. Publish .manifest only after review.""" |
| request: dict[str, Any] |
| manifest: dict[str, Any] |
|
|
|
|
| def _approved_research(run: Run) -> list[str]: |
| from .review import approved_claims, review_packet |
| if run.mode != "live": |
| raise PermissionError("Only live literature runs can authorize research preflight") |
| if not isinstance(run.review, dict) or run.review.get("dataset_id") != review_packet(run)["dataset_id"]: |
| raise PermissionError("A valid, source-bound human review is required") |
| if run.review.get("status") != "READY_FOR_HUMAN_GATED_SYNTHESIS": |
| raise PermissionError("All claims need human review and two distinct supporting sources") |
| sources = sorted({c.paper_id for c in approved_claims(run)}) |
| if len(sources) < 2: |
| raise PermissionError("At least two distinct human-supported evidence sources required") |
| return sources |
|
|
|
|
| def prepare_boltz2( |
| sequence: str, |
| *, |
| source_id: str, |
| mode: str = "offline_fixture", |
| run: Run | None = None, |
| researcher_authorized: bool = False, |
| recycling_steps: int = 3, |
| sampling_steps: int = 50, |
| diffusion_samples: int = 1, |
| ) -> Boltz2Preflight: |
| """Build a standard protein-only Boltz2 request, without executing it. |
| |
| The offline fixture mode is for inert code tests and exploration. A research |
| preflight requires a completed human review and separate explicit consent. |
| Neither mode authorizes remote submission of the sequence. |
| """ |
| seq = _validate_sequence(sequence) |
| sid = _validate_id(source_id) |
| if type(recycling_steps) is not int or not 1 <= recycling_steps <= 10: |
| raise ValueError("recycling_steps must be an integer from 1 to 10") |
| if type(sampling_steps) is not int or not 10 <= sampling_steps <= 1000: |
| raise ValueError("sampling_steps must be an integer from 10 to 1000") |
| if type(diffusion_samples) is not int or not 1 <= diffusion_samples <= 25: |
| raise ValueError("diffusion_samples must be an integer from 1 to 25") |
|
|
| if mode == "offline_fixture": |
| if run is not None or researcher_authorized: |
| raise ValueError("Offline fixtures cannot carry research approvals or runs") |
| approved_sources: list[str] = [] |
| elif mode == "research": |
| if run is None or not researcher_authorized: |
| raise PermissionError("Real research preflight requires a reviewed run and explicit authorization") |
| approved_sources = _approved_research(run) |
| else: |
| raise ValueError("Unsupported BioNeMo preflight mode") |
|
|
| request = { |
| "polymers": [{"id": "A", "molecule_type": "protein", "sequence": seq}], |
| "recycling_steps": recycling_steps, |
| "sampling_steps": sampling_steps, |
| "diffusion_samples": diffusion_samples, |
| "step_scale": 1.638, |
| "output_format": "mmcif", |
| } |
| req_hash = _sha256(_canonical(request)) |
| manifest = { |
| "schema": SCHEMA, |
| "status": "PREPARED_NOT_EXECUTED", |
| "tool": "boltz2-nim", |
| "model": "NVIDIA-hosted Boltz-2 NIM (not called)", |
| "endpoint_reference": BOLTZ2_URL, |
| "input_source_id": sid, |
| "input_sequence_sha256": _sha256(seq.encode("ascii")), |
| "input_sequence_length": len(seq), |
| "request_sha256": req_hash, |
| "request_settings": { |
| "recycling_steps": recycling_steps, |
| "sampling_steps": sampling_steps, |
| "diffusion_samples": diffusion_samples, |
| "output_format": "mmcif", |
| }, |
| "mode": mode, |
| "researcher_preflight_authorized": researcher_authorized, |
| "human_supported_evidence_ids": approved_sources, |
| "network_invocation_occurred": False, |
| "biological_activity_confirmed": False, |
| "scientific_interpretation": "NO MODEL RESULT. This is a schema preflight, not a prediction.", |
| "source_skill_url": SKILLS[0]["source"], |
| } |
| return Boltz2Preflight(request=request, manifest=manifest) |
|
|
|
|
| def validate_boltz2_response(prepared: Boltz2Preflight, response: dict[str, Any]) -> dict[str, Any]: |
| """Inspect actual response data, never synthesize or fill missing values. |
| |
| A validated response is only a correctly shaped computational prediction. |
| Provenance requires an external signed/runtime receipt before 'live' |
| attribution; caller must not claim HTTP success from this function. |
| """ |
| if not isinstance(prepared, Boltz2Preflight) or not isinstance(response, dict): |
| raise TypeError("Boltz2 response and prepared request required") |
| structures = response.get("structures") |
| confidence = response.get("confidence_scores") |
| expected = prepared.request["diffusion_samples"] |
| if not isinstance(structures, list) or len(structures) != expected: |
| raise ValueError("Missing or mismatched returned structures") |
| if not isinstance(confidence, list) or len(confidence) != expected: |
| raise ValueError("Missing or mismatched confidence_scores") |
| content_hashes = [] |
| for structure in structures: |
| if not isinstance(structure, dict) or structure.get("format") != "mmcif": |
| raise ValueError("Every structure must declare format=mmcif") |
| content = structure.get("structure") |
| if not isinstance(content, str) or not content.strip(): |
| raise ValueError("Empty or malformed structure") |
| content_hashes.append(_sha256(content.encode("utf-8"))) |
| for score in confidence: |
| if type(score) not in (float, int) or not math.isfinite(score): |
| raise ValueError("Nonfinite or nonnumeric model confidence") |
| |
| return { |
| **prepared.manifest, |
| "status": "RESPONSE_SHAPE_CHECKED_ORIGIN_UNVERIFIED", |
| "structure_count": expected, |
| "structure_sha256": content_hashes, |
| "confidence_scores": confidence.copy(), |
| "response_sha256": _sha256(_canonical(response)), |
| "source_response_validated_as_live": False, |
| "scientific_interpretation": ( |
| "Computational model outputs only; confidence_scores are not automatically " |
| "pLDDT or proof of experimental binding, function, or clinical efficacy." |
| ), |
| } |
|
|
