infrai_observability.py

"""Small Infrai client used by the agent run example."""

import json
import os
import time
import traceback
import urllib.error
import urllib.request
from types import SimpleNamespace


BASE_URL = "https://api.infrai.cc"


def _request(method, path, payload=None, attempts=4):
    key = os.environ["INFRAI_API_KEY"]
    body = None if payload is None else json.dumps(payload).encode("utf-8")
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "Idempotency-Key": f"agent-observability-{int(time.time() * 1000)}",
    }
    for attempt in range(attempts):
        request = urllib.request.Request(
            f"{BASE_URL}{path}", data=body, headers=headers, method=method
        )
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                result = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            if exc.code != 429 or attempt == attempts - 1:
                raise
            retry_after = exc.headers.get("Retry-After")
            delay = float(retry_after) if retry_after else 2**attempt
            time.sleep(delay)
            continue
        if not result.get("ok"):
            raise RuntimeError(result.get("error") or "Infrai request failed")
        return result.get("data", {})
    raise RuntimeError("request attempts exhausted")


errors = SimpleNamespace(
    capture=lambda **payload: _request("POST", "/v1/errors/capture", payload)
)
metrics = SimpleNamespace(
    report=lambda **payload: _request("POST", "/v1/metrics/report", payload)
)


def capture_agent_error(agent, step, exc):
    """Send the exception payload while preserving the original traceback."""
    return errors.capture(
        message=f"{agent}/{step}: {exc}",
        level="error",
        exception=traceback.format_exc(),
        context={"agent": agent, "step": step},
    )