agentpause 0.2.0__py3-none-any.whl

agentpause/errors.py

@@ -0,0 +1,66 @@
+"""Typed error hierarchy.
+
+Everything agentpause raises derives from :class:`AgentPauseError`, so callers
+can catch one base class — or handle each failure mode precisely:
+
+* :class:`RateLimitHit` — the provider returned 429 despite the prediction
+  (estimates are statistical, not oracular). Carries ``retry_after`` when the
+  provider says how long to wait. The scheduler's retry policy handles these
+  automatically; you only see one if all retries are exhausted.
+* :class:`TelemetryError` — the budget could not be read (missing headers).
+* :class:`CheckpointError` — the checkpoint could not be written or read
+  (disk full, permissions, corrupted file).
+* :class:`BackendError` — the LLM call failed for a non-rate-limit reason.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+__all__ = [
+    "AgentPauseError",
+    "RateLimitHit",
+    "TelemetryError",
+    "CheckpointError",
+    "BackendError",
+]
+
+
+class AgentPauseError(Exception):
+    """Base class for every error agentpause raises."""
+
+
+class RateLimitHit(AgentPauseError):
+    """The provider rate-limited a call (HTTP 429).
+
+    Args:
+        retry_after: seconds the provider asks to wait, when reported.
+    """
+
+    def __init__(self, message: str = "provider returned 429 (rate limit)",
+                 retry_after: Optional[float] = None) -> None:
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+class TelemetryError(AgentPauseError):
+    """The rate-limit budget could not be read from the provider response."""
+
+
+class CheckpointError(AgentPauseError):
+    """A checkpoint could not be persisted or restored."""
+
+
+class BackendError(AgentPauseError):
+    """The LLM call failed for a reason other than rate limiting.
+
+    Args:
+        retriable: True for transient infrastructure failures (HTTP 5xx,
+            timeouts) that deserve a retry; False for request problems
+            (4xx, validation) where retrying only wastes budget.
+    """
+
+    def __init__(self, message: str = "LLM call failed",
+                 retriable: bool = False) -> None:
+        super().__init__(message)
+        self.retriable = retriable

agentpause/estimator.py

@@ -0,0 +1,106 @@
+"""Cost estimation for the next agent step.
+
+The estimator predicts how many tokens the next LLM call will consume, so the
+scheduler can decide whether it fits within the remaining rate-limit budget.
+
+The estimate has two parts:
+
+* a *base* estimate from observable quantities (current input size, the
+  generation cap, and a fixed tool overhead), and
+* a data-driven correction ``epsilon`` — the moving average of recent
+  estimation errors — that adapts the base estimate to the actual behavior of
+  the model and workload.
+
+It also tracks ``sigma``, the standard deviation of realized consumption, which
+the decision rule uses as a safety margin.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Deque, Optional
+
+
+class Estimator:
+    """Predicts next-step token cost and tracks consumption statistics.
+
+    Args:
+        max_tokens: the generation cap passed to the LLM (upper bound on output).
+        tool_overhead: fixed token overhead attributed to tool calls.
+        error_window: how many recent estimation errors feed ``epsilon``.
+        min_estimate: a floor so the estimate is never absurdly small.
+    """
+
+    def __init__(
+        self,
+        max_tokens: int = 650,
+        tool_overhead: int = 280,
+        error_window: int = 6,
+        min_estimate: int = 200,
+    ) -> None:
+        self.max_tokens = max_tokens
+        self.tool_overhead = tool_overhead
+        self.min_estimate = min_estimate
+        self._errors: Deque[int] = deque(maxlen=error_window)
+        self._all_errors: list[int] = []   # full history, for quantile margins
+        self._realized: list[int] = []
+
+    # -- estimation ---------------------------------------------------------
+
+    def base_estimate(self, input_tokens: int) -> int:
+        """Estimate from observable quantities only, without the correction."""
+        return input_tokens + self.max_tokens + self.tool_overhead
+
+    def epsilon(self) -> int:
+        """Moving average of recent estimation errors (0 until data exists)."""
+        if not self._errors:
+            return 0
+        return int(sum(self._errors) / len(self._errors))
+
+    def estimate(self, input_tokens: int) -> int:
+        """Full next-step estimate: base plus the learned correction."""
+        return max(self.min_estimate, self.base_estimate(input_tokens) + self.epsilon())
+
+    def sigma(self, fallback_estimate: int) -> float:
+        """Std. dev. of realized consumption.
+
+        Falls back to a fraction of the current estimate until enough samples
+        exist (< 4), so early decisions stay prudent.
+        """
+        h = self._realized
+        if len(h) >= 4:
+            mean = sum(h) / len(h)
+            return (sum((x - mean) ** 2 for x in h) / len(h)) ** 0.5
+        return fallback_estimate * 0.27
+
+    # -- learning -----------------------------------------------------------
+
+    def record(self, input_tokens: int, realized: int) -> None:
+        """Update history with the actual consumption of a completed step."""
+        error = realized - self.base_estimate(input_tokens)
+        self._errors.append(error)
+        self._all_errors.append(error)
+        self._realized.append(realized)
+
+    def upper_estimate(self, input_tokens: int, q: float = 0.95,
+                       min_samples: int = 8) -> Optional[int]:
+        """Quantile-based upper bound for the next step's cost.
+
+        Consumption distributions are heavy-tailed: the mean-plus-``k·sigma``
+        margin under-covers the tail. When enough history exists, the
+        empirical ``q``-quantile of past estimation errors gives a
+        statistically honest bound: base estimate + q-th worst error seen.
+        Returns ``None`` until ``min_samples`` steps are recorded (callers
+        fall back to the ``k·sigma`` rule).
+        """
+        n = len(self._all_errors)
+        if n < min_samples:
+            return None
+        ordered = sorted(self._all_errors)
+        idx = min(n - 1, max(0, int(q * n + 0.999999) - 1))  # ceil(q·n)-1, clamped
+        return max(self.min_estimate, self.base_estimate(input_tokens) + ordered[idx])
+
+    @property
+    def samples(self) -> int:
+        """Number of completed steps recorded so far."""
+        return len(self._realized)

agentpause/fallback.py

@@ -0,0 +1,74 @@
+"""Model fallback chain: route around a rate-limited or failing model.
+
+The escalation ladder of the accompanying research (§5.3, "model switch")
+as a composable piece: wrap any ordered list of backends and the chain tries
+them in sequence when the current one is rate-limited or fails transiently.
+
+    from agentpause import FallbackBackend, PredictiveScheduler
+
+    chain = FallbackBackend(
+        adapter_opus.backend,      # primary
+        adapter_haiku.backend,     # cheaper fallback
+        on_fallback=lambda i, exc: log.warning("switched to backend %d: %s", i, exc),
+    )
+    sched = PredictiveScheduler(backend=chain, telemetry=adapter_opus.budget)
+
+Only rate limits and *retriable* failures trigger the switch — a 400-style
+error means the request itself is broken, and no other model will fix it.
+
+Behavioral-consistency caveat (the "Sierra lesson"): different models produce
+differently-shaped outputs. The chain exposes ``last_index`` and the
+``on_fallback`` callback so downstream code can validate outputs that came
+from a fallback model instead of trusting them blindly.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from .errors import BackendError, RateLimitHit
+
+__all__ = ["FallbackBackend"]
+
+Backend = Callable[[List[Dict[str, str]]], Tuple[str, int]]
+
+
+class FallbackBackend:
+    """An ordered chain of backends, tried in sequence on transient failure.
+
+    Args:
+        *backends: two or more ``messages -> (reply, tokens_used)`` callables,
+            primary first, cheapest last.
+        on_fallback: optional callback ``(index, exception) -> None`` invoked
+            when the chain moves past a failing backend — use it for logging
+            or to flag outputs for extra validation.
+    """
+
+    def __init__(
+        self,
+        *backends: Backend,
+        on_fallback: Optional[Callable[[int, Exception], None]] = None,
+    ) -> None:
+        if len(backends) < 1:
+            raise ValueError("FallbackBackend needs at least one backend")
+        self.backends = list(backends)
+        self.on_fallback = on_fallback
+        self.last_index: Optional[int] = None   # which backend served the last call
+
+    def __call__(self, messages: List[Dict[str, str]]) -> Tuple[str, int]:
+        last_exc: Optional[Exception] = None
+        for i, backend in enumerate(self.backends):
+            try:
+                result = backend(messages)
+                self.last_index = i
+                return result
+            except RateLimitHit as exc:
+                last_exc = exc
+            except BackendError as exc:
+                if not exc.retriable:
+                    raise               # request problem: no model will fix it
+                last_exc = exc
+            if self.on_fallback is not None and i + 1 < len(self.backends):
+                self.on_fallback(i + 1, last_exc)
+        assert last_exc is not None
+        raise last_exc

agentpause/retry.py

@@ -0,0 +1,57 @@
+"""Retry policy for rate-limit hits the prediction did not prevent.
+
+Estimates are statistical: a 429 can still slip through. When it does, the
+scheduler waits and retries instead of crashing — honoring the provider's
+``retry-after`` when given, otherwise backing off exponentially.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+import time
+from dataclasses import dataclass, field
+from typing import Awaitable, Callable, Optional
+
+__all__ = ["RetryPolicy"]
+
+
+@dataclass
+class RetryPolicy:
+    """How to wait-and-retry after an unexpected 429.
+
+    Args:
+        max_retries: attempts after the first failure before giving up
+            (the original :class:`~agentpause.errors.RateLimitHit` is re-raised).
+        base_delay_s: first backoff delay.
+        factor: multiplier per attempt (exponential backoff).
+        max_delay_s: ceiling for any single delay.
+        sleep_fn: how to wait (defaults to ``time.sleep``; tests inject a fake).
+        async_sleep_fn: how to wait on the async path (defaults to
+            ``asyncio.sleep``; tests inject a fake).
+        jitter: randomize each delay by ±this fraction (default ±25%), so many
+            agents hitting the same window don't all retry in the same instant
+            (the "thundering herd" problem). Set 0.0 for deterministic delays.
+    """
+
+    max_retries: int = 3
+    base_delay_s: float = 1.0
+    factor: float = 2.0
+    max_delay_s: float = 60.0
+    sleep_fn: Callable[[float], None] = field(default=time.sleep)
+    async_sleep_fn: Optional[Callable[[float], Awaitable[None]]] = None
+    jitter: float = 0.25
+
+    def delay(self, attempt: int) -> float:
+        """Backoff delay for the given attempt number (0-based), with jitter."""
+        base = min(self.base_delay_s * self.factor ** attempt, self.max_delay_s)
+        if self.jitter <= 0:
+            return base
+        return base * (1 + self.jitter * (2 * random.random() - 1))
+
+    async def asleep(self, seconds: float) -> None:
+        """Async wait, honoring an injected ``async_sleep_fn``."""
+        if self.async_sleep_fn is not None:
+            await self.async_sleep_fn(seconds)
+        else:
+            await asyncio.sleep(seconds)

agentpause/risk.py

@@ -0,0 +1,138 @@
+"""Risk model and the suspension decision rules.
+
+The pieces:
+
+* :func:`should_checkpoint` — the hard token rule. Suspend when the remaining
+  token budget cannot cover the estimated next step plus a safety margin
+  ``k * sigma``.
+
+* :class:`Budget` — what telemetry can report beyond a bare token count:
+  remaining requests (RPM limits exist alongside TPM limits) and the time
+  until the window resets.
+
+* :func:`decide` — the three-valued rule built on top: ``continue`` when the
+  next step fits (both tokens AND requests), ``wait`` when it does not fit but
+  the window resets imminently (a short pause beats a full suspend/resume
+  cycle), ``checkpoint`` otherwise.
+
+* :class:`RiskModel` — a soft, multi-dimensional risk score used for monitoring
+  and for richer policies. It aggregates rate-limit pressure, context-window
+  pressure, and monetary-budget pressure into a single number.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+def should_checkpoint(remaining: int, estimated: int, sigma: float, k: float = 2.0) -> bool:
+    """Return True when the next step does not safely fit in the remaining budget.
+
+    Rule:  ``remaining < estimated + k * sigma``.
+
+    Args:
+        remaining: tokens left in the current rate-limit window.
+        estimated: estimated cost of the next step.
+        sigma: observed std. dev. of consumption (the margin's scale).
+        k: safety factor; higher ``k`` suspends earlier (more cautious).
+    """
+    return remaining < estimated + k * sigma
+
+
+@dataclass
+class Budget:
+    """A telemetry reading richer than a bare token count.
+
+    Args:
+        remaining_tokens: tokens left in the current rate-limit window (TPM).
+        remaining_requests: requests left in the window (RPM), if known.
+        reset_seconds: seconds until the window resets, if known.
+        remaining_input_tokens: input-token budget, if the provider limits
+            input and output separately (Anthropic does).
+        remaining_output_tokens: output-token budget, if reported separately —
+            a reasoning-heavy agent can drain this while input budget abounds.
+    """
+
+    remaining_tokens: int
+    remaining_requests: Optional[int] = None
+    reset_seconds: Optional[float] = None
+    remaining_input_tokens: Optional[int] = None
+    remaining_output_tokens: Optional[int] = None
+
+
+@dataclass
+class Decision:
+    """The outcome of :func:`decide`, with its inputs kept for diagnostics."""
+
+    action: str            # "continue" | "wait" | "checkpoint"
+    budget: Budget
+    estimated: int
+    sigma: float
+
+
+def decide(
+    budget: Budget,
+    estimated: int,
+    sigma: float,
+    k: float = 2.0,
+    wait_threshold_s: float = 15.0,
+    estimated_input: Optional[int] = None,
+    estimated_output: Optional[int] = None,
+) -> Decision:
+    """The three-valued decision rule.
+
+    ``continue``    the next step fits: tokens cover ``estimated + k*sigma``
+                    AND at least one request remains (when RPM is known).
+    ``wait``        it does not fit, but the window resets within
+                    ``wait_threshold_s`` — pausing in place is cheaper than a
+                    full checkpoint + relaunch.
+    ``checkpoint``  it does not fit and the reset is far away (or unknown).
+
+    Unknown dimensions never block: if the provider reports no RPM, reset,
+    or split input/output information, the rule degrades gracefully to the
+    combined token-only behavior.
+    """
+    tokens_fit = not should_checkpoint(budget.remaining_tokens, estimated, sigma, k)
+    requests_fit = budget.remaining_requests is None or budget.remaining_requests >= 1
+    input_fit = (budget.remaining_input_tokens is None or estimated_input is None
+                 or estimated_input <= budget.remaining_input_tokens)
+    output_fit = (budget.remaining_output_tokens is None or estimated_output is None
+                  or estimated_output <= budget.remaining_output_tokens)
+    if tokens_fit and requests_fit and input_fit and output_fit:
+        action = "continue"
+    elif budget.reset_seconds is not None and budget.reset_seconds <= wait_threshold_s:
+        action = "wait"
+    else:
+        action = "checkpoint"
+    return Decision(action=action, budget=budget, estimated=estimated, sigma=sigma)
+
+
+@dataclass
+class RiskModel:
+    """Weighted, multi-dimensional risk score (diagnostic / for rich policies).
+
+    ``risk = w1 * est/remaining + w2 * context/context_max + w3 * cost/budget``
+
+    The weights default to values validated in the accompanying research and can
+    be overridden or, in future, learned online.
+    """
+
+    w1: float = 0.55  # rate-limit pressure
+    w2: float = 0.20  # context-window pressure
+    w3: float = 0.25  # monetary-budget pressure
+
+    def score(
+        self,
+        estimated: int,
+        remaining: int,
+        context_tokens: int,
+        context_max: int,
+        cost_estimate: float = 0.0,
+        budget_remaining: float = 1.0,
+    ) -> float:
+        """Compute the aggregate risk score (>= 0; higher means riskier)."""
+        rate = estimated / max(remaining, 1)
+        ctx = context_tokens / max(context_max, 1)
+        econ = cost_estimate / max(budget_remaining, 1e-9)
+        return self.w1 * rate + self.w2 * ctx + self.w3 * econ