Quickstart

export PROTEINIQ_API_KEY="YOUR_API_KEY"
export PROTEINIQ_RUN_KEY="$(python3 -c 'import uuid; print(uuid.uuid4())')"

import json
import os
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from pathlib import Path
from uuid import UUID

BASE_URL = "https://proteiniq.io"
API_KEY = os.environ["PROTEINIQ_API_KEY"]
RUN_KEY = str(UUID(os.environ["PROTEINIQ_RUN_KEY"]))
TERMINAL = {"COMPLETED", "FAILED", "TIMEOUT", "CANCELLED", "BUDGET_EXCEEDED"}

def retry_delay(headers):
    value = headers.get("Retry-After", "5")
    try:
        return max(0, float(value))
    except ValueError:
        reset = parsedate_to_datetime(value)
        return max(0, (reset - datetime.now(timezone.utc)).total_seconds())

def request_json(method, path, body=None, idempotency_key=None):
    headers = {"Authorization": f"Bearer {API_KEY}"}
    if body is not None:
        headers["Content-Type"] = "application/json"
    if idempotency_key:
        headers["Idempotency-Key"] = idempotency_key
    data = None if body is None else json.dumps(body).encode()
    for attempt in range(6):
        request = urllib.request.Request(
            BASE_URL + path, data=data, headers=headers, method=method
        )
        try:
            with urllib.request.urlopen(request, timeout=60) as response:
                return json.load(response), response.headers
        except urllib.error.HTTPError as exc:
            payload = json.loads(exc.read())
            error = payload.get("error", {})
            retryable = exc.code == 429 or (
                method == "GET" and error.get("code") == "job_not_completed"
            )
            if not retryable or attempt == 5:
                raise RuntimeError(
                    f"{exc.code} {error.get('code')}: {error.get('message')}"
                ) from exc
            time.sleep(retry_delay(exc.headers))

request = {
    "tool": "esmfold",
    "input": {"inputs": [{
        "id": "protein_1",
        "slotId": "protein",
        "kind": "protein",
        "format": "fasta",
        "content": ">example\nACDEFGHIKLMNPQRSTVWY",
        "source": {"type": "text"},
    }]},
    "settings": {},
}

saved_request = Path("submissions") / f"{RUN_KEY}.json"
if saved_request.exists():
    # Recover the exact approved request, including its accepted price.
    submission = json.loads(saved_request.read_text(encoding="utf-8"))
    print("Reusing saved submission:", saved_request)
else:
    tool, _ = request_json("GET", "/api/v1/tools/esmfold")
    print("Tool:", tool["name"])
    quote, _ = request_json("POST", "/api/v1/jobs/quote", request)
    print(json.dumps(quote, indent=2))
    if quote.get("blocking_errors"):
        raise RuntimeError("Resolve the quote's blocking_errors before submitting.")
    if input("Submit this paid analysis? [y/N] ").strip().lower() != "y":
        raise SystemExit("No job submitted.")
    submission = {
        **request,
        "name": "ESMfold API example",
        "max_quoted_credits": quote["billable_credits"],
    }
    saved_request.parent.mkdir(parents=True, exist_ok=True)
    saved_request.write_text(json.dumps(submission, indent=2), encoding="utf-8")
job, _ = request_json("POST", "/api/v1/jobs", submission, RUN_KEY)
job_id = job["id"]
print("Job ID:", job_id, flush=True)
print("Results page:", BASE_URL + "/app/jobs/" + job_id)

delay = 5
deadline = time.monotonic() + 3600
while job["status"] not in TERMINAL:
    if time.monotonic() >= deadline:
        raise TimeoutError(f"Stopped waiting; job {job_id} may still be running.")
    time.sleep(delay)
    job, headers = request_json("GET", f"/api/v1/jobs/{job_id}/status")
    print("Status:", job["status"], flush=True)
    delay = max(min(delay * 2, 30), retry_delay(headers))

if job["status"] != "COMPLETED":
    raise RuntimeError(
        f"Job {job_id} ended as {job['status']}: {job.get('error')}. "
        "Check the Results reference for partial-output availability."
    )

result, _ = request_json("GET", f"/api/v1/results/{job_id}")
output_dir = Path("outputs") / job_id
output_dir.mkdir(parents=True, exist_ok=True)
(output_dir / "result.json").write_text(json.dumps(result, indent=2), encoding="utf-8")

downloaded = 0
for index, file in enumerate(result.get("files", []), start=1):
    if not file.get("url"):
        continue
    # Use a local filename, never an untrusted path from the response.
    name = Path(file.get("name") or "result-file").name
    destination = output_dir / f"{index}-{name}"
    # Signed URLs carry their own access: do not send the API key.
    with urllib.request.urlopen(file["url"], timeout=60) as response:
        with destination.open("wb") as output:
            while chunk := response.read(1024 * 1024):
                output.write(chunk)
    downloaded += 1
    print("Saved:", destination)

print(f"Saved result.json and {downloaded} files in {output_dir}")