Production Readiness Checklist

import logging
import os
import random
import time
import uuid
from dataclasses import dataclass
from typing import Any, Optional

import requests
from dotenv import load_dotenv

load_dotenv()

API_KEY = os.getenv("SWARMS_API_KEY")
BASE_URL = "https://api.swarms.world"

HEADERS = {"x-api-key": API_KEY, "Content-Type": "application/json"}

log = logging.getLogger("swarms.prod")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


@dataclass
class RunResult:
    request_id: str
    ok: bool
    status_code: int
    cost: float
    elapsed: float
    body: Any
    rate_limit: dict


class BudgetExceeded(RuntimeError):
    pass


class ProductionClient:
    """Production wrapper around the Swarms /v1/swarm/completions endpoint.

    Features:
      - Retry-with-backoff that honors Retry-After on 429s.
      - Hard per-call and cumulative budget caps.
      - Pre-flight credit check.
      - Per-request UUID for correlation against /v1/account/logs.
      - Rate-limit header capture on every response.
    """

    def __init__(
        self,
        per_call_budget: float = 5.00,
        cumulative_budget: float = 50.00,
        max_retries: int = 4,
        base_backoff: float = 1.5,
        request_timeout: float = 600.0,
    ):
        self.per_call_budget = per_call_budget
        self.cumulative_budget = cumulative_budget
        self.cumulative_spend = 0.0
        self.max_retries = max_retries
        self.base_backoff = base_backoff
        self.request_timeout = request_timeout

    # ---- pre-flight ----

    def credits_remaining(self) -> float:
        r = requests.get(f"{BASE_URL}/v1/account/credits", headers=HEADERS, timeout=10)
        r.raise_for_status()
        return float(r.json().get("total_credits", 0))

    def assert_credits(self, required: float):
        available = self.credits_remaining()
        if available < required:
            raise BudgetExceeded(
                f"Insufficient credits: have ${available:.2f}, need ${required:.2f}"
            )

    # ---- core call ----

    def run(self, payload: dict, request_id: Optional[str] = None) -> RunResult:
        rid = request_id or str(uuid.uuid4())

        # Pre-flight: refuse if we've already blown the cumulative budget.
        if self.cumulative_spend >= self.cumulative_budget:
            raise BudgetExceeded(
                f"Cumulative budget ${self.cumulative_budget:.2f} exceeded "
                f"(spent ${self.cumulative_spend:.2f})."
            )

        # Tag the payload so we can find it in logs later.
        payload = {**payload, "name": payload.get("name", "untitled") + f" [{rid[:8]}]"}

        last_error = None
        for attempt in range(self.max_retries + 1):
            t0 = time.monotonic()
            try:
                resp = requests.post(
                    f"{BASE_URL}/v1/swarm/completions",
                    headers=HEADERS,
                    json=payload,
                    timeout=self.request_timeout,
                )
            except (requests.ConnectionError, requests.Timeout) as e:
                last_error = e
                self._sleep_backoff(attempt)
                log.warning("network error, attempt=%s rid=%s err=%s", attempt, rid, e)
                continue

            elapsed = time.monotonic() - t0
            rl = {k: v for k, v in resp.headers.items() if k.lower().startswith("x-ratelimit")}

            # 429: honor Retry-After exactly.
            if resp.status_code == 429:
                wait = int(resp.headers.get("Retry-After", "30"))
                log.warning("429 rid=%s retry-after=%ss", rid, wait)
                time.sleep(wait)
                continue

            # 5xx: retry with backoff.
            if 500 <= resp.status_code < 600:
                log.warning("%s rid=%s attempt=%s body=%s",
                            resp.status_code, rid, attempt, resp.text[:200])
                self._sleep_backoff(attempt)
                last_error = RuntimeError(f"{resp.status_code}: {resp.text[:200]}")
                continue

            # 4xx (other than 429): permanent. Do not retry.
            if 400 <= resp.status_code < 500:
                log.error("4xx rid=%s status=%s body=%s",
                          rid, resp.status_code, resp.text[:500])
                return RunResult(
                    request_id=rid, ok=False, status_code=resp.status_code,
                    cost=0.0, elapsed=elapsed, body=resp.text, rate_limit=rl,
                )

            # 2xx: parse, enforce per-call budget, return.
            body = resp.json()
            cost = float(body.get("usage", {}).get("billing_info", {}).get("total_cost", 0))
            if cost > self.per_call_budget:
                log.error("rid=%s exceeded per-call budget: $%.4f > $%.2f",
                          rid, cost, self.per_call_budget)
            self.cumulative_spend += cost

            log.info(
                "OK rid=%s cost=$%.4f elapsed=%.1fs min=%s/%s",
                rid, cost, elapsed,
                rl.get("X-RateLimit-Remaining-Minute"), rl.get("X-RateLimit-Limit-Minute"),
            )

            return RunResult(
                request_id=rid, ok=True, status_code=resp.status_code,
                cost=cost, elapsed=elapsed, body=body, rate_limit=rl,
            )

        raise RuntimeError(f"rid={rid} exhausted {self.max_retries} retries: {last_error}")

    def _sleep_backoff(self, attempt: int):
        # Exponential backoff with full jitter.
        sleep = random.uniform(0, self.base_backoff * (2 ** attempt))
        time.sleep(sleep)

    # ---- post-flight ----

    def verify_in_logs(self, request_id: str) -> bool:
        """Confirm the request landed in /v1/account/logs (audit-trail check).

        Each log row's ``data`` field is the full SwarmCompletion response —
        the swarm's name comes back as ``swarm_name``, not ``name``. Only
        successful, billable runs are ever logged here, so this only confirms
        success; it can't tell you a request failed.
        """
        r = requests.get(f"{BASE_URL}/v1/account/logs", headers=HEADERS, timeout=30)
        r.raise_for_status()
        marker = request_id[:8]
        for entry in r.json().get("logs", []):
            data = entry.get("data") or {}
            swarm_name = data.get("swarm_name") or ""
            if marker in swarm_name:
                return True
        return False

client = ProductionClient(per_call_budget=2.00, cumulative_budget=20.00)

# Pre-flight: refuse to start if we don't have enough credit headroom.
client.assert_credits(required=client.cumulative_budget * 1.5)

payload = {
    "name": "Nightly Research Sweep",
    "swarm_type": "SequentialWorkflow",
    "max_loops": 1,
    "task": "Summarize today's macro-market developments.",
    "agents": [
        {
            "agent_name": "Macro Analyst",
            "system_prompt": "You are a senior macro analyst...",
            "model_name": "gpt-4.1",
            "role": "worker",
            "max_loops": 1,
            "max_tokens": 4096,
            "temperature": 0.2,
        },
    ],
}

result = client.run(payload)
print(f"request_id={result.request_id} cost=${result.cost:.4f} ok={result.ok}")

# Post-flight: confirm the audit trail is intact.
if not client.verify_in_logs(result.request_id):
    log.warning("rid=%s did not appear in /v1/account/logs", result.request_id)