Handle Errors and Retries

import os
import random
import time

import requests

API = "https://api.zerodrift.ai"
HEADERS = {"x-api-key": os.environ["ZERODRIFT_API_KEY"]}
MAX_ATTEMPTS = 4
MAX_WAIT_S = 30
TERMINAL = {"completed", "failed"}


class ClientError(Exception):
    """A 4xx other than 429. Fix the request or key. Never fail open on this."""


class Unavailable(Exception):
    """Retries ran out, the API asked for a long wait, or the job failed."""


def request_with_retry(method: str, path: str, **kwargs) -> requests.Response:
    last_error = "no response"
    for attempt in range(MAX_ATTEMPTS):
        retry_after = None
        try:
            r = requests.request(
                method, f"{API}{path}", headers=HEADERS, timeout=30, **kwargs
            )
        except (requests.ConnectionError, requests.Timeout) as e:
            last_error = str(e)
        else:
            if r.ok:
                return r
            if r.status_code != 429 and r.status_code < 500:
                raise ClientError(f"HTTP {r.status_code}: {r.text[:200]}")
            last_error = f"HTTP {r.status_code}"
            retry_after = r.headers.get("Retry-After")
        delay = float(retry_after) if retry_after else min(2**attempt, 8)
        if delay > MAX_WAIT_S or attempt == MAX_ATTEMPTS - 1:
            break
        time.sleep(delay + random.random())
    raise Unavailable(last_error)


def check(content: str, fail_closed: bool = True) -> dict:
    try:
        r = request_with_retry(
            "POST",
            "/api/v3/content/validate",
            json={"content": content, "mode": "sync", "model_engine": "anchor_3_0"},
        )
        job = r.json()
        if job.get("mode") == "async":
            poll_url = job["poll"]["url"]
            deadline = time.monotonic() + 120
            while job["status"] not in TERMINAL:
                if time.monotonic() > deadline:
                    raise Unavailable(f"timeout polling {poll_url}")
                time.sleep(1)
                job = request_with_retry("GET", poll_url).json()
        if job["status"] == "failed":
            raise Unavailable(job.get("error", "enforcement failed"))
        return {"checked": True, "job": job}
    except Unavailable as e:
        print(f"ZeroDrift unchecked: {e}")  # log it, and alert on the count
        return {"checked": False, "deliver": not fail_closed, "reason": str(e)}


outcome = check("Our fund guarantees 20% returns with zero risk!", fail_closed=True)
if not outcome["checked"] and not outcome["deliver"]:
    print("Held: could not check this message.")

const API = "https://api.zerodrift.ai";
const HEADERS = {
  "x-api-key": process.env.ZERODRIFT_API_KEY!,
  "Content-Type": "application/json",
};
const MAX_ATTEMPTS = 4;
const MAX_WAIT_MS = 30_000;
const TERMINAL = new Set(["completed", "failed"]);

/** A 4xx other than 429. Fix the request or key. Never fail open on this. */
class ClientError extends Error {}
/** Retries ran out, the API asked for a long wait, or the job failed. */
class Unavailable extends Error {}

const sleep = (ms: number) => new Promise((res) => setTimeout(res, ms));

async function requestWithRetry(
  method: string,
  path: string,
  body?: unknown,
): Promise<Response> {
  let lastError = "no response";
  for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
    let retryAfter: string | null = null;
    try {
      const r = await fetch(`${API}${path}`, {
        method,
        headers: HEADERS,
        body: body === undefined ? undefined : JSON.stringify(body),
        signal: AbortSignal.timeout(30_000),
      });
      if (r.ok) return r;
      if (r.status !== 429 && r.status < 500) {
        throw new ClientError(`HTTP ${r.status}`);
      }
      lastError = `HTTP ${r.status}`;
      retryAfter = r.headers.get("Retry-After");
    } catch (e) {
      if (e instanceof ClientError) throw e;
      lastError = String(e);
    }
    const delay = retryAfter
      ? Number(retryAfter) * 1000
      : Math.min(2 ** attempt, 8) * 1000;
    if (delay > MAX_WAIT_MS || attempt === MAX_ATTEMPTS - 1) break;
    await sleep(delay + Math.random() * 1000);
  }
  throw new Unavailable(lastError);
}

type Outcome =
  | { checked: true; job: any }
  | { checked: false; deliver: boolean; reason: string };

export async function check(content: string, failClosed = true): Promise<Outcome> {
  try {
    const r = await requestWithRetry("POST", "/api/v3/content/validate", {
      content,
      mode: "sync",
      model_engine: "anchor_3_0",
    });
    let job = await r.json();
    if (job.mode === "async") {
      const pollUrl: string = job.poll.url;
      const deadline = Date.now() + 120_000;
      while (!TERMINAL.has(job.status)) {
        if (Date.now() > deadline) throw new Unavailable(`timeout polling ${pollUrl}`);
        await sleep(1000);
        job = await (await requestWithRetry("GET", pollUrl)).json();
      }
    }
    if (job.status === "failed") {
      throw new Unavailable(job.error ?? "enforcement failed");
    }
    return { checked: true, job };
  } catch (e) {
    if (!(e instanceof Unavailable)) throw e;
    console.error(`ZeroDrift unchecked: ${e.message}`); // log it, and alert on the count
    return { checked: false, deliver: !failClosed, reason: e.message };
  }
}

const outcome = await check("Our fund guarantees 20% returns with zero risk!", true);
if (!outcome.checked && !outcome.deliver) {
  console.log("Held: could not check this message.");
}