Oracles

def oracles():
    return []

def oracles() -> list[OracleSpec]:
    """Return one entry per oracle this file provides."""
def prepare_oracle() -> None:
    """Optional. Runs once per worker process, lazily, before the first
    fold, never at container startup. This is where a co-resident model
    loads its weights, or where a client to an externally hosted model opens
    its connection."""

FoldRequest(
    example_id,             # names the CIF and confidence JSON written for this design
    cif_text,               # the design, already serialised
    atom_array,             # the same design as a biotite AtomArray
    chain_of_interest,      # the binder/designed chain, when the request names one
    msa_paths_by_chain_id,  # per-chain MSA paths, when any were supplied
)
FoldResult(
    cif_text,               # the refolded structure
    confidences,            # {key: float}, covered by published_keys
)

def oracles():
    return [
        OracleSpec(
            name="oracle_boltz2",
            supported_tags={"unconditional_refolding"},
            published_keys={"overall_plddt", "overall_ptm", "overall_iptm"},
            fold=_fold_general,
        ),
        OracleSpec(
            name="oracle_boltz2_ppi",
            supported_tags={"ppi_refolding"},
            published_keys={"overall_plddt", "overall_ptm", "overall_iptm", "ptm_coi", "iptm_coi"},
            fold=_fold_ppi,
        ),
    ]

"""reward.py: refold through the NVCF-hosted OpenFold3 NIM."""
from __future__ import annotations
import os
from api.oracles.contract import FoldRequest, FoldResult, OracleSpec, UnsupportedDesign
#: An NVCF invoke key. Required -- there is no anonymous access to a hosted
#: function. See the module docstring for why this is a different credential
#: from an NGC_API_KEY that only pulls container images.
_API_KEY_ENV_VAR = "OF3_NIM_API_KEY"
#: Fixed -- this is the one public endpoint, not a per-deployment address the
#: way a self-hosted NIM's is. Override only if NVCF's own base URL changes.
_URL_ENV_VAR = "OF3_NIM_URL"
_DEFAULT_URL = "https://health.api.nvidia.com/v1"
#: Sent alongside the API key, transcribed from curl.sh -- NVCF's own knob
#: for how long the caller is willing to wait inline before it should be
#: expected to poll separately instead. 300s covers this oracle's sampling
#: schedule with room to spare.
_POLL_SECONDS_ENV_VAR = "OF3_NIM_POLL_SECONDS"
_DEFAULT_POLL_SECONDS = "300"
#: Sampling settings sent with every fold. Deliberately light -- this oracle
#: is refolding-for-scoring, not a from-scratch structure prediction, and a
#: heavier schedule only costs latency (and, here, NVCF billing) the
#: search/metrics path is not asking for. Override via env for more samples.
_RECYCLING_STEPS = int(os.environ.get("OF3_NIM_RECYCLING_STEPS", "3"))
_DIFFUSION_SAMPLES = int(os.environ.get("OF3_NIM_DIFFUSION_SAMPLES", "1"))
_SAMPLING_STEPS = int(os.environ.get("OF3_NIM_SAMPLING_STEPS", "20"))
#: HTTP timeout for one fold call. Generous, and independent of
#: NVCF-POLL-SECONDS: that header tells the *server* how long to hold the
#: connection open before switching to async polling; this is this client's
#: own socket timeout and has to be at least as long, so it is kept a little
#: longer rather than assumed equal.
_TIMEOUT_SECONDS = float(os.environ.get("OF3_NIM_TIMEOUT_SECONDS", "330"))
_SUPPORTED_TAGS = frozenset({"unconditional_refolding"})
#: See the module docstring for why this stops at complex-level confidence.
_PUBLISHED_KEYS = frozenset({"overall_plddt", "overall_ptm", "overall_iptm"})
def oracles() -> list[OracleSpec]:
    return [
        OracleSpec(
            name="oracle_of3_nvcf",
            supported_tags=_SUPPORTED_TAGS,
            published_keys=_PUBLISHED_KEYS,
            fold=_fold,
        )
    ]
def prepare_oracle() -> None:
    """Fail at startup on a missing credential, not on the first design's HTTP 401."""
    if not os.environ.get(_API_KEY_ENV_VAR):
        raise RuntimeError(
            f"{_API_KEY_ENV_VAR} is not set. This oracle calls the NVCF-hosted OpenFold3 NIM, "
            "which requires an invoke key (Authorization: Bearer) on every request -- there is "
            "no anonymous or health-checkable route to fall back to."
        )
def _base_url() -> str:
    return os.environ.get(_URL_ENV_VAR, _DEFAULT_URL).rstrip("/")
def _headers() -> dict[str, str]:
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ[_API_KEY_ENV_VAR]}",
        "NVCF-POLL-SECONDS": os.environ.get(
            _POLL_SECONDS_ENV_VAR, _DEFAULT_POLL_SECONDS
        ),
    }
def _fold(request: FoldRequest) -> FoldResult:
    molecules = _build_molecules(request)
    payload = {
        "request_id": request.example_id,
        "inputs": [
            {
                "input_id": request.example_id,
                "molecules": molecules,
                "recycling_steps": _RECYCLING_STEPS,
                "diffusion_samples": _DIFFUSION_SAMPLES,
                "sampling_steps": _SAMPLING_STEPS,
                "output_format": "cif",
            }
        ],
    }
    response_body = _post(payload)
    output = response_body["outputs"][0]
    structure = output["structures_with_scores"][0]
    confidences: dict[str, float] = {
        "overall_plddt": structure["complex_plddt_score"],
        "overall_ptm": structure["ptm_score"],
    }
    # Monomers get iptm_score: null -- omitted rather than published as NaN,
    # so a single-chain design does not carry a meaningless interface score.
    if structure.get("iptm_score") is not None:
        confidences["overall_iptm"] = structure["iptm_score"]
    return FoldResult(cif_text=structure["structure"], confidences=confidences)
def _build_molecules(request: FoldRequest) -> list[dict]:
    """One protein-chain entry per polymer protein chain, MSA-free.
    MSA-free refolding is not a shortcut here -- AGENTS.md's "Refolding
    oracles" section is explicit that every upstream predictor runs this way
    in the refolding role, "database-free". A single-sequence a3m (the query
    against itself, no real hits) is how this API expresses that: its DTO
    requires a non-empty MSA per protein, verified live against the
    self-hosted sibling (identical DTO) -- an absent ``msa`` field 422s with
    "Each protein sequence ... must be accompanied by a non-empty MSA".
    Raises ``UnsupportedDesign`` for anything that is not a plain protein
    chain (ligand, nucleic acid) -- this file does not yet build those
    molecule types. See the module docstring for what extending this needs.
    """
    from atomworks.enums import ChainTypeInfo
    from atomworks.io.utils.chain_info import build_chain_info
    chain_info = build_chain_info(request.atom_array)
    molecules = []
    for chain_id in sorted(chain_info):
        info = chain_info[chain_id]
        if info["chain_type"] not in ChainTypeInfo.PROTEINS:
            raise UnsupportedDesign(
                f"oracle_of3_nvcf only builds plain protein chains today; chain {chain_id!r} is "
                f"{info['chain_type']!r} ({request.example_id})"
            )
        sequence = info["processed_entity_canonical_sequence"]
        molecules.append(
            {
                "type": "protein",
                "id": chain_id,
                "sequence": sequence,
                "msa": {
                    "query": {
                        "a3m": {
                            "alignment": f">query\n{sequence}\n",
                            "format": "a3m",
                            "rank": 0,
                        }
                    }
                },
            }
        )
    if not molecules:
        raise UnsupportedDesign(
            f"no protein chains found to refold ({request.example_id})"
        )
    return molecules
def _post(payload: dict) -> dict:
    from api.oracles.http import post_json
    return post_json(
        f"{_base_url()}/biology/openfold/openfold3/predict",
        payload,
        headers=_headers(),
        timeout_seconds=_TIMEOUT_SECONDS,
        poll_base_url=f"{_base_url()}/status",
    )

"""reward.py: two roles of the NVCF-hosted Boltz-2 NIM, one file."""
from __future__ import annotations
import os
from api.oracles.contract import FoldRequest, FoldResult, OracleSpec, UnsupportedDesign
#: An NVCF invoke key. Required -- there is no anonymous access to a hosted
#: function. Same variable name as the OpenFold3 NVCF example, deliberately
#: -- both are NVCF invoke keys, not NGC_API_KEYs, and a deployment mounting
#: either oracle configures the same credential the same way.
_API_KEY_ENV_VAR = "OF3_NIM_API_KEY"
#: Fixed -- this is the one public endpoint, not a per-deployment address.
_URL_ENV_VAR = "BOLTZ2_NIM_URL"
_DEFAULT_URL = "https://health.api.nvidia.com/v1"
_POLL_SECONDS_ENV_VAR = "BOLTZ2_NIM_POLL_SECONDS"
_DEFAULT_POLL_SECONDS = "300"
#: Sampling settings sent with every fold. Deliberately light and fixed at
#: diffusion_samples=1 -- Boltz2's per-sample score fields
#: (chains_ptm_scores in particular) are ambiguous about how they flatten
#: across multiple diffusion samples in the schema itself, and this file
#: only verified the diffusion_samples=1 shape live. Raise it only after
#: re-verifying how chains_ptm_scores indexes with diffusion_samples > 1.
_RECYCLING_STEPS = int(os.environ.get("BOLTZ2_NIM_RECYCLING_STEPS", "1"))
_SAMPLING_STEPS = int(os.environ.get("BOLTZ2_NIM_SAMPLING_STEPS", "20"))
_DIFFUSION_SAMPLES = 1
_TIMEOUT_SECONDS = float(os.environ.get("BOLTZ2_NIM_TIMEOUT_SECONDS", "330"))
#: Confidence scale this NIM's pass-criterion thresholds are written for
#: (e.g. "min: 80.0"). Boltz2 returns 0-1; multiply by this before returning
#: anything through FoldResult.confidences.
_SCALE = 100.0
_GENERAL_TAGS = frozenset({"unconditional_refolding"})
_GENERAL_KEYS = frozenset({"overall_plddt", "overall_ptm", "overall_iptm"})
_PPI_TAGS = frozenset({"ppi_refolding"})
#: See the module docstring for why min_ipae_coi cannot be here.
_PPI_KEYS = frozenset(
    {"overall_plddt", "overall_ptm", "overall_iptm", "ptm_coi", "iptm_coi"}
)
def oracles() -> list[OracleSpec]:
    return [
        OracleSpec(
            name="oracle_boltz2",
            supported_tags=_GENERAL_TAGS,
            published_keys=_GENERAL_KEYS,
            fold=_fold_general,
        ),
        OracleSpec(
            name="oracle_boltz2_ppi",
            supported_tags=_PPI_TAGS,
            published_keys=_PPI_KEYS,
            fold=_fold_ppi,
        ),
    ]
def prepare_oracle() -> None:
    """Fail at startup on a missing credential, not on the first design's HTTP 401.
    Shared by both roles: one credential, checked once, regardless of which
    of the two oracles above a worker reaches first.
    """
    if not os.environ.get(_API_KEY_ENV_VAR):
        raise RuntimeError(
            f"{_API_KEY_ENV_VAR} is not set. Both Boltz2 oracles in this file call the "
            "NVCF-hosted Boltz2 NIM, which requires an invoke key (Authorization: Bearer) on "
            "every request -- there is no anonymous or health-checkable route to fall back to."
        )
def _fold_general(request: FoldRequest) -> FoldResult:
    """The general role: complex-level confidence only, no interface reading."""
    chain_ids, cif_text, response_body = _predict(request)
    confidences: dict[str, float] = {
        "overall_plddt": response_body["complex_plddt_scores"][0] * _SCALE,
        "overall_ptm": response_body["ptm_scores"][0] * _SCALE,
    }
    # Zero from a single-chain input is "no interface", not "measured zero" --
    # omitted rather than published as a real (bad) score.
    if len(chain_ids) > 1:
        confidences["overall_iptm"] = response_body["iptm_scores"][0] * _SCALE
    return FoldResult(cif_text=cif_text, confidences=confidences)
def _fold_ppi(request: FoldRequest) -> FoldResult:
    """The PPI role: same fold, plus chain-of-interest-specific confidence.
    Declines (``UnsupportedDesign``) a single-chain design outright -- there
    is no interface for a "PPI" oracle to score, so answering with zeros
    would misrepresent absence as a measured bad result.
    """
    chain_ids, cif_text, response_body = _predict(request)
    if len(chain_ids) < 2:
        raise UnsupportedDesign(
            f"oracle_boltz2_ppi needs at least two chains to score an interface; "
            f"got {len(chain_ids)} ({request.example_id})"
        )
    confidences: dict[str, float] = {
        "overall_plddt": response_body["complex_plddt_scores"][0] * _SCALE,
        "overall_ptm": response_body["ptm_scores"][0] * _SCALE,
        "overall_iptm": response_body["iptm_scores"][0] * _SCALE,
    }
    coi = request.chain_of_interest
    if coi is not None and coi in chain_ids:
        coi_index = chain_ids.index(coi)
        chains_ptm = response_body.get("chains_ptm_scores") or []
        if coi_index < len(chains_ptm):
            confidences["ptm_coi"] = chains_ptm[coi_index] * _SCALE
        pair_iptm_samples = response_body.get("pair_chains_iptm_scores") or []
        if pair_iptm_samples:
            # Not symmetric in practice (verified live: pair_iptm["0"]["1"] !=
            # pair_iptm["1"]["0"]), so both directions are averaged to match upstream's
            # chain-of-interest confidence convention.
            pair_iptm = pair_iptm_samples[0]
            interface_scores = []
            for j in range(len(chain_ids)):
                if j == coi_index:
                    continue
                values = [
                    pair_iptm.get(str(coi_index), {}).get(str(j)),
                    pair_iptm.get(str(j), {}).get(str(coi_index)),
                ]
                values = [value for value in values if value is not None]
                if values:
                    interface_scores.append(sum(values) / len(values))
            if interface_scores:
                confidences["iptm_coi"] = (
                    sum(interface_scores) / len(interface_scores)
                ) * _SCALE
    return FoldResult(cif_text=cif_text, confidences=confidences)
def _predict(request: FoldRequest) -> tuple[list[str], str, dict]:
    """Shared by both roles: build the request, call Boltz2 once, return the raw response.
    Returns ``(chain_ids, cif_text, response_body)`` -- ``chain_ids[i]`` names
    the chain the request's ``polymers[i]`` described, which is also the
    index Boltz2's own ``chains_ptm_scores``/``pair_chains_iptm_scores`` use
    to refer back to that chain. Building the index this way, once, is what
    keeps those two in agreement -- Boltz2's response carries no chain
    *names*, only positions.
    """
    chain_ids, polymers = _build_polymers(request)
    payload = {
        "polymers": polymers,
        "recycling_steps": _RECYCLING_STEPS,
        "sampling_steps": _SAMPLING_STEPS,
        "diffusion_samples": _DIFFUSION_SAMPLES,
        "step_scale": 1.2,
        "without_potentials": True,
        "output_format": "mmcif",
    }
    response_body = _post(payload)
    cif_text = response_body["structures"][0]["structure"]
    return chain_ids, cif_text, response_body
def _build_polymers(request: FoldRequest) -> tuple[list[str], list[dict]]:
    """One protein-chain entry per polymer protein chain, MSA-free, in a fixed order.
    MSA-free refolding is not a shortcut here -- AGENTS.md's "Refolding
    oracles" section is explicit that every upstream predictor runs this way
    in the refolding role, "database-free". A single-sequence a3m (the query
    against itself, no real hits) is the same convention the OpenFold3
    examples use.
    Raises ``UnsupportedDesign`` for anything that is not a plain protein
    chain (ligand, nucleic acid) -- this file does not yet build those
    molecule types. See the module docstring for what extending this needs.
    """
    from atomworks.enums import ChainTypeInfo
    from atomworks.io.utils.chain_info import build_chain_info
    chain_info = build_chain_info(request.atom_array)
    chain_ids: list[str] = []
    polymers: list[dict] = []
    for chain_id in sorted(chain_info):
        info = chain_info[chain_id]
        if info["chain_type"] not in ChainTypeInfo.PROTEINS:
            raise UnsupportedDesign(
                f"this file only builds plain protein chains today; chain {chain_id!r} is "
                f"{info['chain_type']!r} ({request.example_id})"
            )
        sequence = info["processed_entity_canonical_sequence"]
        chain_ids.append(chain_id)
        polymers.append(
            {
                "id": chain_id,
                "molecule_type": "protein",
                "sequence": sequence,
                "msa": {
                    "uniref90": {
                        "a3m": {"alignment": f">query\n{sequence}", "format": "a3m"}
                    }
                },
            }
        )
    if not polymers:
        raise UnsupportedDesign(
            f"no protein chains found to refold ({request.example_id})"
        )
    return chain_ids, polymers
def _base_url() -> str:
    return os.environ.get(_URL_ENV_VAR, _DEFAULT_URL).rstrip("/")
def _headers() -> dict[str, str]:
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ[_API_KEY_ENV_VAR]}",
        "NVCF-POLL-SECONDS": os.environ.get(
            _POLL_SECONDS_ENV_VAR, _DEFAULT_POLL_SECONDS
        ),
    }
def _post(payload: dict) -> dict:
    from api.oracles.http import post_json
    return post_json(
        f"{_base_url()}/biology/mit/boltz2/predict",
        payload,
        headers=_headers(),
        timeout_seconds=_TIMEOUT_SECONDS,
        poll_base_url=f"{_base_url()}/status",
    )