agentpause 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. {agentpause-0.2.0/src/agentpause.egg-info → agentpause-0.2.1}/PKG-INFO +16 -1
  2. {agentpause-0.2.0 → agentpause-0.2.1}/README.md +15 -0
  3. {agentpause-0.2.0 → agentpause-0.2.1}/pyproject.toml +1 -1
  4. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/__init__.py +3 -1
  5. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/langgraph.py +9 -3
  6. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/litellm.py +13 -0
  7. agentpause-0.2.1/src/agentpause/adapters/openai_compat.py +206 -0
  8. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/estimator.py +7 -2
  9. agentpause-0.2.1/src/agentpause/refill.py +75 -0
  10. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/risk.py +55 -6
  11. {agentpause-0.2.0 → agentpause-0.2.1/src/agentpause.egg-info}/PKG-INFO +16 -1
  12. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/SOURCES.txt +5 -0
  13. agentpause-0.2.1/tests/test_openai_compat.py +98 -0
  14. agentpause-0.2.1/tests/test_refill_wait.py +77 -0
  15. agentpause-0.2.1/tests/test_regime.py +98 -0
  16. {agentpause-0.2.0 → agentpause-0.2.1}/LICENSE +0 -0
  17. {agentpause-0.2.0 → agentpause-0.2.1}/setup.cfg +0 -0
  18. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/__init__.py +0 -0
  19. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/breaker.py +0 -0
  20. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/errors.py +0 -0
  21. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/fallback.py +0 -0
  22. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/retry.py +0 -0
  23. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/scheduler.py +0 -0
  24. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/state.py +0 -0
  25. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/dependency_links.txt +0 -0
  26. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/requires.txt +0 -0
  27. {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/top_level.txt +0 -0
  28. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_async.py +0 -0
  29. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_breaker.py +0 -0
  30. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_estimator.py +0 -0
  31. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_fallback.py +0 -0
  32. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_io_tpm.py +0 -0
  33. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_langgraph_adapter.py +0 -0
  34. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_litellm_adapter.py +0 -0
  35. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_money_and_hooks.py +0 -0
  36. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_quantile.py +0 -0
  37. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_risk.py +0 -0
  38. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_robustness.py +0 -0
  39. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_scheduler.py +0 -0
  40. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_state.py +0 -0
  41. {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_wait_and_rpm.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentpause
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  Summary: Predictive scheduling for autonomous LLM agents: suspend gracefully before rate limits, resume without losing work.
5
5
  Author: Andrea Lelio Ciampolillo
6
6
  License: MIT
@@ -40,6 +40,21 @@ plugin for self-hosted runtimes (llama.cpp, vLLM).
40
40
  | End-to-end A/B (thinking models, 4B/8B) | recovery **54×–93× faster**; total task time −19% |
41
41
 
42
42
  Reproduce them yourself with a free key: `python scripts/benchmark_groq.py`.
43
+ Live run (2026-07-08, Groq free tier, 12 steps, ~1k-token context,
44
+ refill-aware waiting enabled):
45
+
46
+ | | reactive baseline | agentpause |
47
+ |---|---|---|
48
+ | 429 errors suffered | 7 | **0** |
49
+ | steps redone | 7 | **0** |
50
+ | tokens re-sent (waste) | 12,520 | **0** |
51
+ | wall-clock | 79 s | 145 s* |
52
+
53
+ \* zero waste costs some waiting. Refill-aware math already cuts most waits
54
+ from ~50 s to ~10 s; the residue is a benchmark artifact (condition B starts
55
+ right after A has drained the shared TPM window) plus the safety cap when a
56
+ late-task call needs nearly the whole bucket. Waste is money on paid tiers —
57
+ waiting is free.
43
58
 
44
59
  > Status: **early alpha (0.1)**. Core components, the high-level
45
60
  > `PredictiveScheduler` API, the LiteLLM adapter (any provider), the LangGraph
@@ -17,6 +17,21 @@ plugin for self-hosted runtimes (llama.cpp, vLLM).
17
17
  | End-to-end A/B (thinking models, 4B/8B) | recovery **54×–93× faster**; total task time −19% |
18
18
 
19
19
  Reproduce them yourself with a free key: `python scripts/benchmark_groq.py`.
20
+ Live run (2026-07-08, Groq free tier, 12 steps, ~1k-token context,
21
+ refill-aware waiting enabled):
22
+
23
+ | | reactive baseline | agentpause |
24
+ |---|---|---|
25
+ | 429 errors suffered | 7 | **0** |
26
+ | steps redone | 7 | **0** |
27
+ | tokens re-sent (waste) | 12,520 | **0** |
28
+ | wall-clock | 79 s | 145 s* |
29
+
30
+ \* zero waste costs some waiting. Refill-aware math already cuts most waits
31
+ from ~50 s to ~10 s; the residue is a benchmark artifact (condition B starts
32
+ right after A has drained the shared TPM window) plus the safety cap when a
33
+ late-task call needs nearly the whole bucket. Waste is money on paid tiers —
34
+ waiting is free.
20
35
 
21
36
  > Status: **early alpha (0.1)**. Core components, the high-level
22
37
  > `PredictiveScheduler` API, the LiteLLM adapter (any provider), the LangGraph
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "agentpause"
7
- version = "0.2.0"
7
+ version = "0.2.1"
8
8
  description = "Predictive scheduling for autonomous LLM agents: suspend gracefully before rate limits, resume without losing work."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.9"
@@ -15,12 +15,13 @@ from .errors import (
15
15
  from .breaker import CircuitBreaker, CircuitOpenError
16
16
  from .estimator import Estimator
17
17
  from .fallback import FallbackBackend
18
+ from .refill import RegimeDetector
18
19
  from .retry import RetryPolicy
19
20
  from .risk import Budget, Decision, RiskModel, decide, should_checkpoint
20
21
  from .state import Checkpoint, StateStore
21
22
  from .scheduler import PredictiveScheduler, Session
22
23
 
23
- __version__ = "0.2.0"
24
+ __version__ = "0.2.1"
24
25
 
25
26
  __all__ = [
26
27
  "PredictiveScheduler",
@@ -37,6 +38,7 @@ __all__ = [
37
38
  "FallbackBackend",
38
39
  "CircuitBreaker",
39
40
  "CircuitOpenError",
41
+ "RegimeDetector",
40
42
  "AgentPauseError",
41
43
  "RateLimitHit",
42
44
  "TelemetryError",
@@ -97,6 +97,7 @@ class AgentPauseGuard:
97
97
  wait_threshold_s: float = 15.0,
98
98
  sleep_fn: Callable[[float], None] = time.sleep,
99
99
  async_sleep_fn: Optional[Callable[[float], Any]] = None,
100
+ chunk_s: float = 10.0,
100
101
  max_resumes: int = 100,
101
102
  ) -> None:
102
103
  self.telemetry = telemetry
@@ -107,6 +108,7 @@ class AgentPauseGuard:
107
108
  self.wait_threshold_s = wait_threshold_s
108
109
  self.sleep_fn = sleep_fn
109
110
  self.async_sleep_fn = async_sleep_fn
111
+ self.chunk_s = chunk_s
110
112
  self.max_resumes = max_resumes
111
113
 
112
114
  # -- internals -----------------------------------------------------------
@@ -144,8 +146,11 @@ class AgentPauseGuard:
144
146
  if d.action == "continue":
145
147
  return
146
148
  if d.action == "wait":
147
- # reset is imminent: sleeping beats a suspend/resume cycle
148
- self.sleep_fn((budget.reset_seconds or 1.0) + 0.5)
149
+ # reset is imminent: sleeping beats a suspend/resume cycle.
150
+ # Sleep in chunks and re-read: if the bucket refills sooner,
151
+ # resume sooner (and the samples feed regime detection).
152
+ wait_s = d.wait_seconds or budget.reset_seconds or 1.0
153
+ self.sleep_fn(min(wait_s, self.chunk_s) + 0.5)
149
154
  continue
150
155
  # interrupt() raises on the first pass (the node aborts here);
151
156
  # on a resume pass it RETURNS, and the loop re-reads telemetry.
@@ -177,7 +182,8 @@ class AgentPauseGuard:
177
182
  if d.action == "continue":
178
183
  return
179
184
  if d.action == "wait":
180
- delay = (budget.reset_seconds or 1.0) + 0.5
185
+ wait_s = d.wait_seconds or budget.reset_seconds or 1.0
186
+ delay = min(wait_s, self.chunk_s) + 0.5
181
187
  if self.async_sleep_fn is not None:
182
188
  await self.async_sleep_fn(delay)
183
189
  else:
@@ -71,11 +71,23 @@ _INSTALL_HINT = ("The LiteLLM adapter needs the 'litellm' package. "
71
71
  "Install it with: pip install agentpause[litellm]")
72
72
 
73
73
 
74
+ def _enable_headers() -> None:
75
+ """Ask litellm to surface provider response headers.
76
+
77
+ Without this flag litellm does NOT populate
78
+ ``_hidden_params["additional_headers"]`` — telemetry would be blind.
79
+ (Found in the field on the very first real-provider run.)
80
+ """
81
+ import litellm
82
+ litellm.return_response_headers = True
83
+
84
+
74
85
  def _default_completion() -> Callable[..., Any]:
75
86
  try:
76
87
  from litellm import completion
77
88
  except ImportError as exc: # pragma: no cover - depends on environment
78
89
  raise ImportError(_INSTALL_HINT) from exc
90
+ _enable_headers()
79
91
  return completion
80
92
 
81
93
 
@@ -84,6 +96,7 @@ def _default_acompletion() -> Callable[..., Any]:
84
96
  from litellm import acompletion
85
97
  except ImportError as exc: # pragma: no cover - depends on environment
86
98
  raise ImportError(_INSTALL_HINT) from exc
99
+ _enable_headers()
87
100
  return acompletion
88
101
 
89
102
 
@@ -0,0 +1,206 @@
1
+ """Direct adapter for OpenAI-compatible providers (Groq, OpenAI, ...).
2
+
3
+ Reads rate-limit headers straight from the HTTP response — no middle layer.
4
+ Born in the field: litellm currently drops provider response headers
5
+ (BerriAI/litellm#11749), which blinds telemetry. This adapter talks to the
6
+ provider's ``/chat/completions`` endpoint directly, so the headers the
7
+ scheduler needs are always there.
8
+
9
+ from agentpause.adapters.openai_compat import OpenAICompatAdapter
10
+
11
+ adapter = OpenAICompatAdapter.for_model("groq/llama-3.1-8b-instant")
12
+ sched = PredictiveScheduler(backend=adapter.backend,
13
+ telemetry=adapter.budget,
14
+ count_tokens=adapter.count_tokens)
15
+
16
+ Works with any provider exposing the OpenAI chat API: pass ``base_url`` and
17
+ the env var holding the key, or use :meth:`for_model` for known prefixes.
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ import os
23
+ import time
24
+ from typing import Any, Callable, Dict, List, Optional, Tuple
25
+
26
+ from ..errors import BackendError, RateLimitHit, TelemetryError
27
+ from ..refill import RegimeDetector
28
+ from ..risk import Budget
29
+ from .litellm import _parse_reset
30
+
31
+ __all__ = ["OpenAICompatAdapter"]
32
+
33
+ KNOWN_PROVIDERS = {
34
+ "groq": ("https://api.groq.com/openai/v1", "GROQ_API_KEY"),
35
+ "openai": ("https://api.openai.com/v1", "OPENAI_API_KEY"),
36
+ }
37
+
38
+
39
+ def _default_post(url: str, headers: Dict[str, str], payload: Dict[str, Any]):
40
+ """Perform the HTTP POST; returns (status, response_headers, json_body)."""
41
+ import httpx
42
+ r = httpx.post(url, headers=headers, json=payload, timeout=120.0)
43
+ try:
44
+ body = r.json()
45
+ except Exception:
46
+ body = {}
47
+ return r.status_code, dict(r.headers), body
48
+
49
+
50
+ class OpenAICompatAdapter:
51
+ """Backend + telemetry via the provider's raw HTTP API.
52
+
53
+ Args:
54
+ model: model name as the PROVIDER knows it (e.g. ``llama-3.1-8b-instant``).
55
+ base_url: API root, e.g. ``https://api.groq.com/openai/v1``.
56
+ api_key: literal key, or None to read it from ``api_key_env``.
57
+ api_key_env: environment variable holding the key.
58
+ post_fn: transport, injectable for tests:
59
+ ``(url, headers, payload) -> (status, headers, json)``.
60
+ fallback_remaining / max_age_s / clock: as in the LiteLLM adapter.
61
+ **request_kwargs: extra payload fields sent on every call.
62
+ """
63
+
64
+ def __init__(
65
+ self,
66
+ model: str,
67
+ base_url: str,
68
+ api_key: Optional[str] = None,
69
+ api_key_env: str = "OPENAI_API_KEY",
70
+ post_fn: Optional[Callable[..., Tuple[int, Dict[str, str], Dict[str, Any]]]] = None,
71
+ fallback_remaining: Optional[int] = None,
72
+ max_age_s: float = 10.0,
73
+ clock: Callable[[], float] = time.monotonic,
74
+ **request_kwargs: Any,
75
+ ) -> None:
76
+ self.model = model
77
+ self.base_url = base_url.rstrip("/")
78
+ key = api_key if api_key is not None else os.environ.get(api_key_env)
79
+ if not key:
80
+ raise TelemetryError(f"No API key: set {api_key_env} or pass api_key=...")
81
+ self._auth = {"Authorization": f"Bearer {key}",
82
+ "Content-Type": "application/json"}
83
+ self._post = post_fn if post_fn is not None else _default_post
84
+ self.fallback_remaining = fallback_remaining
85
+ self.max_age_s = max_age_s
86
+ self._clock = clock
87
+ self.request_kwargs = request_kwargs
88
+ self._budget: Optional[Budget] = None
89
+ self._read_at: Optional[float] = None
90
+ # D8: refill-regime detection from ping-to-ping samples
91
+ self.detector = RegimeDetector()
92
+ self._prev_ping: Optional[Tuple[float, int, float]] = None
93
+
94
+ @classmethod
95
+ def for_model(cls, model: str, **kwargs: Any) -> "OpenAICompatAdapter":
96
+ """Build from a prefixed model string: ``groq/...``, ``openai/...``,
97
+ or a bare OpenAI model name like ``gpt-4o-mini``."""
98
+ prefix, _, rest = model.partition("/")
99
+ if rest and prefix in KNOWN_PROVIDERS:
100
+ base_url, env = KNOWN_PROVIDERS[prefix]
101
+ return cls(rest, base_url=base_url, api_key_env=env, **kwargs)
102
+ base_url, env = KNOWN_PROVIDERS["openai"]
103
+ return cls(model, base_url=base_url, api_key_env=env, **kwargs)
104
+
105
+ # -- the callables the scheduler wants -------------------------------------
106
+
107
+ def backend(self, messages: List[Dict[str, str]]) -> Tuple[str, int]:
108
+ body = self._call(messages, **self.request_kwargs)
109
+ self._prev_ping = None # a real call breaks the sampling chain
110
+ reply = (body.get("choices") or [{}])[0].get("message", {}).get("content") or ""
111
+ used = int(body.get("usage", {}).get("total_tokens") or 0)
112
+ return reply, used
113
+
114
+ def telemetry(self) -> int:
115
+ return self.budget().remaining_tokens
116
+
117
+ def budget(self) -> Budget:
118
+ if not self._is_fresh():
119
+ self.ping()
120
+ if self._budget is not None:
121
+ regime = self.detector.regime
122
+ if regime != "unknown":
123
+ self._budget.refill_regime = regime
124
+ return self._budget
125
+ if self.fallback_remaining is not None:
126
+ return Budget(remaining_tokens=self.fallback_remaining)
127
+ raise TelemetryError(
128
+ "Provider sent no rate-limit headers; set fallback_remaining=... "
129
+ "to proceed with a fixed assumption."
130
+ )
131
+
132
+ def ping(self) -> None:
133
+ """1-token call whose only purpose is fresh rate-limit headers.
134
+
135
+ Consecutive pings with no real call in between are exactly the
136
+ no-traffic samples the regime detector needs: it observes whether
137
+ ``remaining`` rises between them (continuous bucket) or stays flat
138
+ (fixed window).
139
+ """
140
+ prev = self._prev_ping
141
+ self._call([{"role": "user", "content": "ping"}], max_tokens=1)
142
+ b, t = self._budget, self._read_at
143
+ if b is None or t is None:
144
+ return
145
+ implied = 0.0
146
+ if b.limit_tokens is not None and (b.reset_seconds or 0) > 0:
147
+ implied = (b.limit_tokens - b.remaining_tokens) / b.reset_seconds
148
+ if prev is not None:
149
+ self.detector.feed(dt=t - prev[0],
150
+ dr=b.remaining_tokens - prev[1],
151
+ implied_rate=prev[2])
152
+ self._prev_ping = (t, b.remaining_tokens, implied)
153
+
154
+ def count_tokens(self, text: str) -> int:
155
+ return max(1, len(text) // 4)
156
+
157
+ # -- internals ---------------------------------------------------------------
158
+
159
+ def _is_fresh(self) -> bool:
160
+ return (self._read_at is not None
161
+ and (self._clock() - self._read_at) <= self.max_age_s)
162
+
163
+ def _call(self, messages: List[Dict[str, str]], **extra: Any) -> Dict[str, Any]:
164
+ status, headers, body = self._post(
165
+ f"{self.base_url}/chat/completions",
166
+ self._auth,
167
+ {"model": self.model, "messages": messages, **extra},
168
+ )
169
+ self._absorb(headers)
170
+ if status == 429:
171
+ ra = headers.get("retry-after")
172
+ raise RateLimitHit(retry_after=float(ra) if ra else None)
173
+ if 500 <= status <= 599:
174
+ raise BackendError(f"provider returned {status}", retriable=True)
175
+ if status >= 400:
176
+ raise BackendError(f"provider returned {status}: {str(body)[:200]}")
177
+ return body
178
+
179
+ def _absorb(self, headers: Dict[str, str]) -> None:
180
+ h = {str(k).lower(): v for k, v in headers.items()}
181
+
182
+ def grab(*names: str) -> Optional[str]:
183
+ for n in names:
184
+ if n in h:
185
+ return h[n]
186
+ return None
187
+
188
+ tokens = grab("x-ratelimit-remaining-tokens",
189
+ "anthropic-ratelimit-tokens-remaining")
190
+ if tokens is None:
191
+ return # no telemetry in this response
192
+ limit = grab("x-ratelimit-limit-tokens", "anthropic-ratelimit-tokens-limit")
193
+ reset = grab("x-ratelimit-reset-tokens", "anthropic-ratelimit-tokens-reset")
194
+ requests_ = grab("x-ratelimit-remaining-requests",
195
+ "anthropic-ratelimit-requests-remaining")
196
+ in_ = grab("anthropic-ratelimit-input-tokens-remaining")
197
+ out = grab("anthropic-ratelimit-output-tokens-remaining")
198
+ self._budget = Budget(
199
+ remaining_tokens=int(float(tokens)),
200
+ remaining_requests=int(float(requests_)) if requests_ is not None else None,
201
+ reset_seconds=_parse_reset(reset) if reset is not None else None,
202
+ remaining_input_tokens=int(float(in_)) if in_ is not None else None,
203
+ remaining_output_tokens=int(float(out)) if out is not None else None,
204
+ limit_tokens=int(float(limit)) if limit is not None else None,
205
+ )
206
+ self._read_at = self._clock()
@@ -62,12 +62,17 @@ class Estimator:
62
62
  return max(self.min_estimate, self.base_estimate(input_tokens) + self.epsilon())
63
63
 
64
64
  def sigma(self, fallback_estimate: int) -> float:
65
- """Std. dev. of realized consumption.
65
+ """Std. dev. of the estimation RESIDUALS (realized − base estimate).
66
+
67
+ Raw consumption trends upward as the context grows; measuring its
68
+ spread would mistake systematic growth for uncertainty and inflate
69
+ the safety margin (and the waits) enormously. The residuals absorb
70
+ the trend, so their spread is the honest uncertainty of the estimate.
66
71
 
67
72
  Falls back to a fraction of the current estimate until enough samples
68
73
  exist (< 4), so early decisions stay prudent.
69
74
  """
70
- h = self._realized
75
+ h = self._all_errors
71
76
  if len(h) >= 4:
72
77
  mean = sum(h) / len(h)
73
78
  return (sum((x - mean) ** 2 for x in h) / len(h)) ** 0.5
@@ -0,0 +1,75 @@
1
+ """Online detection of the provider's refill regime (D8).
2
+
3
+ Two refill physics exist behind rate limits:
4
+
5
+ * **continuous token bucket** — tokens trickle back at a steady rate
6
+ (Groq, OpenAI): refill-aware waits are exact;
7
+ * **fixed window** — nothing returns until a boundary instant, then the
8
+ whole budget at once: the only safe wait is the full reset.
9
+
10
+ Nobody documents which one a provider uses, but it can be *measured*:
11
+ take two telemetry readings with no real call in between and look at the
12
+ slope of ``remaining``. Rising at roughly the header-implied rate →
13
+ continuous. Flat (or nibbled slightly by the pings themselves) → fixed.
14
+
15
+ The detector votes on each valid sample pair and reports the majority of
16
+ the recent votes; with insufficient evidence it says ``unknown``, and the
17
+ caller stays conservative (full-reset waits).
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ from collections import deque
23
+ from typing import Deque
24
+
25
+ __all__ = ["RegimeDetector"]
26
+
27
+
28
+ class RegimeDetector:
29
+ """Classify the refill regime from (Δt, Δremaining) observations.
30
+
31
+ Args:
32
+ min_gap_s: sample pairs closer than this are noise and are ignored.
33
+ window: how many recent votes participate in the majority.
34
+ min_votes: evidence required before leaving ``unknown``.
35
+ """
36
+
37
+ CONTINUOUS_RATIO = 0.5 # observed/implied rate above this → continuous
38
+ FIXED_RATIO = 0.1 # below this → fixed window
39
+
40
+ def __init__(self, min_gap_s: float = 2.0, window: int = 3,
41
+ min_votes: int = 2) -> None:
42
+ self.min_gap_s = min_gap_s
43
+ self.min_votes = min_votes
44
+ self._votes: Deque[str] = deque(maxlen=window)
45
+
46
+ def feed(self, dt: float, dr: float, implied_rate: float) -> None:
47
+ """Record one observation.
48
+
49
+ Args:
50
+ dt: seconds between the two readings (no real call in between).
51
+ dr: change in remaining tokens over that gap.
52
+ implied_rate: ``(limit - remaining) / reset_seconds`` at the
53
+ first reading — what a continuous bucket would refill at.
54
+ """
55
+ if dt < self.min_gap_s or implied_rate <= 0:
56
+ return
57
+ ratio = (dr / dt) / implied_rate
58
+ if ratio >= self.CONTINUOUS_RATIO:
59
+ self._votes.append("continuous")
60
+ elif ratio <= self.FIXED_RATIO:
61
+ self._votes.append("fixed")
62
+ # ambiguous ratios cast no vote
63
+
64
+ @property
65
+ def regime(self) -> str:
66
+ """``"continuous"``, ``"fixed"``, or ``"unknown"`` (be conservative)."""
67
+ if len(self._votes) < self.min_votes:
68
+ return "unknown"
69
+ cont = sum(1 for v in self._votes if v == "continuous")
70
+ fixed = len(self._votes) - cont
71
+ if cont > fixed:
72
+ return "continuous"
73
+ if fixed > cont:
74
+ return "fixed"
75
+ return "unknown"
@@ -52,6 +52,12 @@ class Budget:
52
52
  input and output separately (Anthropic does).
53
53
  remaining_output_tokens: output-token budget, if reported separately —
54
54
  a reasoning-heavy agent can drain this while input budget abounds.
55
+ limit_tokens: the bucket capacity (TPM limit). Together with
56
+ ``reset_seconds`` it yields the refill rate, enabling
57
+ refill-aware waits.
58
+ refill_regime: ``"continuous"`` / ``"fixed"`` / None, as measured by
59
+ :class:`~agentpause.refill.RegimeDetector`. Fixed-window providers
60
+ get full-reset waits regardless of the refill math.
55
61
  """
56
62
 
57
63
  remaining_tokens: int
@@ -59,6 +65,8 @@ class Budget:
59
65
  reset_seconds: Optional[float] = None
60
66
  remaining_input_tokens: Optional[int] = None
61
67
  remaining_output_tokens: Optional[int] = None
68
+ limit_tokens: Optional[int] = None
69
+ refill_regime: Optional[str] = None
62
70
 
63
71
 
64
72
  @dataclass
@@ -69,6 +77,38 @@ class Decision:
69
77
  budget: Budget
70
78
  estimated: int
71
79
  sigma: float
80
+ wait_seconds: Optional[float] = None # suggested wait, when action == "wait"
81
+
82
+
83
+ _WAIT_BUFFER = 1.15 # wait 15% longer than the bare refill math suggests
84
+
85
+
86
+ def _refill_wait(budget: Budget, needed: float) -> Optional[float]:
87
+ """Seconds until the bucket holds ``needed`` tokens, from headers alone.
88
+
89
+ Token buckets refill continuously; the refill rate falls out of the
90
+ headers with no extra assumptions:
91
+
92
+ rate = (limit - remaining) / reset_seconds
93
+
94
+ so the wait for a deficit is ``deficit / rate`` (buffered, and capped at
95
+ the full reset — never wait longer than a guaranteed-full bucket).
96
+ Returns None when the headers don't allow the math.
97
+ """
98
+ if budget.reset_seconds is None:
99
+ return None
100
+ if budget.refill_regime == "fixed":
101
+ return budget.reset_seconds # nothing trickles back: wait it out
102
+ if budget.limit_tokens is None or budget.reset_seconds <= 0:
103
+ return budget.reset_seconds
104
+ missing = budget.limit_tokens - budget.remaining_tokens
105
+ if missing <= 0:
106
+ return budget.reset_seconds
107
+ rate = missing / budget.reset_seconds
108
+ deficit = needed - budget.remaining_tokens
109
+ if deficit <= 0:
110
+ return 0.5
111
+ return min(budget.reset_seconds, deficit / rate * _WAIT_BUFFER)
72
112
 
73
113
 
74
114
  def decide(
@@ -100,12 +140,21 @@ def decide(
100
140
  output_fit = (budget.remaining_output_tokens is None or estimated_output is None
101
141
  or estimated_output <= budget.remaining_output_tokens)
102
142
  if tokens_fit and requests_fit and input_fit and output_fit:
103
- action = "continue"
104
- elif budget.reset_seconds is not None and budget.reset_seconds <= wait_threshold_s:
105
- action = "wait"
106
- else:
107
- action = "checkpoint"
108
- return Decision(action=action, budget=budget, estimated=estimated, sigma=sigma)
143
+ return Decision("continue", budget, estimated, sigma)
144
+ # compute the ACTUAL wait needed, then compare THAT to the threshold:
145
+ # a bucket full in 50s may hold enough for the next call after 10s
146
+ wait_s: Optional[float] = None
147
+ if budget.reset_seconds is not None:
148
+ if not tokens_fit and requests_fit and input_fit and output_fit:
149
+ # pure token deficit: wait only until the bucket refills enough
150
+ wait_s = _refill_wait(budget, needed=estimated + k * sigma)
151
+ else:
152
+ # requests or split dimensions exhausted: refill rate unknown,
153
+ # stay conservative and wait out the full reset
154
+ wait_s = budget.reset_seconds
155
+ if wait_s is not None and wait_s <= wait_threshold_s:
156
+ return Decision("wait", budget, estimated, sigma, wait_seconds=wait_s)
157
+ return Decision("checkpoint", budget, estimated, sigma)
109
158
 
110
159
 
111
160
  @dataclass
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentpause
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  Summary: Predictive scheduling for autonomous LLM agents: suspend gracefully before rate limits, resume without losing work.
5
5
  Author: Andrea Lelio Ciampolillo
6
6
  License: MIT
@@ -40,6 +40,21 @@ plugin for self-hosted runtimes (llama.cpp, vLLM).
40
40
  | End-to-end A/B (thinking models, 4B/8B) | recovery **54×–93× faster**; total task time −19% |
41
41
 
42
42
  Reproduce them yourself with a free key: `python scripts/benchmark_groq.py`.
43
+ Live run (2026-07-08, Groq free tier, 12 steps, ~1k-token context,
44
+ refill-aware waiting enabled):
45
+
46
+ | | reactive baseline | agentpause |
47
+ |---|---|---|
48
+ | 429 errors suffered | 7 | **0** |
49
+ | steps redone | 7 | **0** |
50
+ | tokens re-sent (waste) | 12,520 | **0** |
51
+ | wall-clock | 79 s | 145 s* |
52
+
53
+ \* zero waste costs some waiting. Refill-aware math already cuts most waits
54
+ from ~50 s to ~10 s; the residue is a benchmark artifact (condition B starts
55
+ right after A has drained the shared TPM window) plus the safety cap when a
56
+ late-task call needs nearly the whole bucket. Waste is money on paid tiers —
57
+ waiting is free.
43
58
 
44
59
  > Status: **early alpha (0.1)**. Core components, the high-level
45
60
  > `PredictiveScheduler` API, the LiteLLM adapter (any provider), the LangGraph
@@ -6,6 +6,7 @@ src/agentpause/breaker.py
6
6
  src/agentpause/errors.py
7
7
  src/agentpause/estimator.py
8
8
  src/agentpause/fallback.py
9
+ src/agentpause/refill.py
9
10
  src/agentpause/retry.py
10
11
  src/agentpause/risk.py
11
12
  src/agentpause/scheduler.py
@@ -18,6 +19,7 @@ src/agentpause.egg-info/top_level.txt
18
19
  src/agentpause/adapters/__init__.py
19
20
  src/agentpause/adapters/langgraph.py
20
21
  src/agentpause/adapters/litellm.py
22
+ src/agentpause/adapters/openai_compat.py
21
23
  tests/test_async.py
22
24
  tests/test_breaker.py
23
25
  tests/test_estimator.py
@@ -26,7 +28,10 @@ tests/test_io_tpm.py
26
28
  tests/test_langgraph_adapter.py
27
29
  tests/test_litellm_adapter.py
28
30
  tests/test_money_and_hooks.py
31
+ tests/test_openai_compat.py
29
32
  tests/test_quantile.py
33
+ tests/test_refill_wait.py
34
+ tests/test_regime.py
30
35
  tests/test_risk.py
31
36
  tests/test_robustness.py
32
37
  tests/test_scheduler.py
@@ -0,0 +1,98 @@
1
+ """Behavioral contract for the direct OpenAI-compatible adapter.
2
+
3
+ Exists because litellm currently drops provider response headers
4
+ (BerriAI/litellm#11749): telemetry must not depend on a middle layer.
5
+ All tests inject a fake transport — offline, no keys.
6
+ """
7
+
8
+ import pytest
9
+
10
+ from agentpause import BackendError, PredictiveScheduler, RateLimitHit, StateStore
11
+ from agentpause.adapters.openai_compat import OpenAICompatAdapter
12
+
13
+ HEADERS = {
14
+ "x-ratelimit-remaining-tokens": "5000",
15
+ "x-ratelimit-remaining-requests": "12",
16
+ "x-ratelimit-reset-tokens": "7.66s",
17
+ }
18
+ BODY = {"choices": [{"message": {"content": "hello"}}], "usage": {"total_tokens": 42}}
19
+
20
+
21
+ class FakePost:
22
+ def __init__(self, status=200, headers=None, body=None):
23
+ self.status, self.headers = status, HEADERS if headers is None else headers
24
+ self.body = BODY if body is None else body
25
+ self.calls = []
26
+
27
+ def __call__(self, url, headers, payload):
28
+ self.calls.append((url, payload))
29
+ return self.status, self.headers, self.body
30
+
31
+
32
+ def make_adapter(post, **kw):
33
+ return OpenAICompatAdapter("m", base_url="https://api.example.com/v1",
34
+ api_key="k", post_fn=post, **kw)
35
+
36
+
37
+ def test_backend_returns_reply_and_tokens():
38
+ post = FakePost()
39
+ a = make_adapter(post)
40
+ assert a.backend([{"role": "user", "content": "hi"}]) == ("hello", 42)
41
+ url, payload = post.calls[0]
42
+ assert url.endswith("/chat/completions")
43
+ assert payload["model"] == "m"
44
+
45
+
46
+ def test_budget_parsed_from_real_http_headers():
47
+ a = make_adapter(FakePost())
48
+ b = a.budget() # pings once
49
+ assert b.remaining_tokens == 5000
50
+ assert b.remaining_requests == 12
51
+ assert b.reset_seconds == pytest.approx(7.66)
52
+
53
+
54
+ def test_backend_refreshes_budget_no_extra_ping():
55
+ post = FakePost()
56
+ a = make_adapter(post)
57
+ a.backend([{"role": "user", "content": "hi"}])
58
+ assert a.telemetry() == 5000
59
+ assert len(post.calls) == 1
60
+
61
+
62
+ def test_429_maps_to_rate_limit_hit():
63
+ post = FakePost(status=429, headers=dict(HEADERS, **{"retry-after": "9"}))
64
+ a = make_adapter(post)
65
+ with pytest.raises(RateLimitHit) as exc:
66
+ a.backend([{"role": "user", "content": "hi"}])
67
+ assert exc.value.retry_after == 9.0
68
+
69
+
70
+ def test_5xx_is_retriable_4xx_is_not():
71
+ a = make_adapter(FakePost(status=503))
72
+ with pytest.raises(BackendError) as exc:
73
+ a.backend([])
74
+ assert exc.value.retriable is True
75
+ a = make_adapter(FakePost(status=400))
76
+ with pytest.raises(BackendError) as exc:
77
+ a.backend([])
78
+ assert exc.value.retriable is False
79
+
80
+
81
+ def test_for_model_maps_known_prefixes(monkeypatch):
82
+ monkeypatch.setenv("GROQ_API_KEY", "k1")
83
+ monkeypatch.setenv("OPENAI_API_KEY", "k2")
84
+ g = OpenAICompatAdapter.for_model("groq/llama-3.1-8b-instant")
85
+ assert g.base_url == "https://api.groq.com/openai/v1"
86
+ assert g.model == "llama-3.1-8b-instant"
87
+ o = OpenAICompatAdapter.for_model("gpt-4o-mini")
88
+ assert o.base_url == "https://api.openai.com/v1"
89
+
90
+
91
+ def test_scheduler_integration(tmp_path):
92
+ a = make_adapter(FakePost())
93
+ sched = PredictiveScheduler(backend=a.backend, telemetry=a.budget,
94
+ store=StateStore(str(tmp_path)))
95
+ with sched.session("t") as s:
96
+ s.add_user("q")
97
+ assert s.next_action().action == "continue"
98
+ assert s.call() == "hello"
@@ -0,0 +1,77 @@
1
+ """Behavioral contract for D7: refill-aware waiting.
2
+
3
+ Providers with token-bucket limits refill CONTINUOUSLY: waiting for the full
4
+ window reset is wasteful. From the headers alone we can derive the refill
5
+ rate — (limit - remaining) / reset_seconds — and wait only until the bucket
6
+ holds enough for the NEXT call (plus a safety buffer).
7
+
8
+ Live-benchmark motivation (2026-07-08, Groq): full-reset waits cost 162s
9
+ wall-clock vs 78s for the reactive baseline; refill-aware waits close the gap.
10
+ """
11
+
12
+ import pytest
13
+
14
+ from agentpause import Budget, decide
15
+
16
+
17
+ def test_wait_is_proportional_to_the_deficit():
18
+ # bucket: 6000 cap, 1000 left, full in 50s → refill rate 100 tok/s
19
+ b = Budget(remaining_tokens=1000, reset_seconds=50.0, limit_tokens=6000)
20
+ # need 2000 (estimated+k·sigma≈2000) → deficit 1000 → 10s × 1.15 buffer
21
+ d = decide(b, estimated=1800, sigma=100.0, k=2.0)
22
+ assert d.action == "wait"
23
+ assert d.wait_seconds == pytest.approx((2000 - 1000) / 100.0 * 1.15)
24
+
25
+
26
+ def test_wait_never_exceeds_full_reset():
27
+ b = Budget(remaining_tokens=0, reset_seconds=50.0, limit_tokens=6000)
28
+ d = decide(b, estimated=1_000_000, sigma=0.0, # absurd need
29
+ wait_threshold_s=60.0) # generous threshold
30
+ assert d.action == "wait"
31
+ assert d.wait_seconds == 50.0 # capped at the reset
32
+
33
+
34
+ def test_long_effective_wait_still_checkpoints():
35
+ """Refill math says 48s but the threshold is 15s → suspend, don't idle."""
36
+ b = Budget(remaining_tokens=0, reset_seconds=50.0, limit_tokens=6000)
37
+ d = decide(b, estimated=5000, sigma=0.0, wait_threshold_s=15.0)
38
+ assert d.action == "checkpoint"
39
+
40
+
41
+ def test_short_effective_wait_beats_a_long_reset():
42
+ """The 2026-07-08 benchmark case: full reset 50s, but the next call only
43
+ needs ~10s of refill — with refill-aware math we wait instead of idling."""
44
+ b = Budget(remaining_tokens=1000, reset_seconds=50.0, limit_tokens=6000)
45
+ d = decide(b, estimated=1800, sigma=100.0, k=2.0, wait_threshold_s=15.0)
46
+ assert d.action == "wait"
47
+ assert d.wait_seconds < 15.0
48
+
49
+
50
+ def test_without_limit_header_falls_back_to_full_reset():
51
+ b = Budget(remaining_tokens=100, reset_seconds=8.0)
52
+ d = decide(b, estimated=1000, sigma=100.0)
53
+ assert d.action == "wait"
54
+ assert d.wait_seconds == 8.0
55
+
56
+
57
+ def test_exhausted_requests_wait_the_full_reset():
58
+ """RPM refill rate is unknown: stay conservative on that dimension."""
59
+ b = Budget(remaining_tokens=100_000, remaining_requests=0,
60
+ reset_seconds=8.0, limit_tokens=200_000)
61
+ d = decide(b, estimated=100, sigma=10.0)
62
+ assert d.action == "wait"
63
+ assert d.wait_seconds == 8.0
64
+
65
+
66
+ def test_continue_carries_no_wait():
67
+ b = Budget(remaining_tokens=100_000, limit_tokens=200_000, reset_seconds=30.0)
68
+ d = decide(b, estimated=100, sigma=10.0)
69
+ assert d.action == "continue"
70
+ assert d.wait_seconds is None
71
+
72
+
73
+ def test_checkpoint_carries_no_wait():
74
+ b = Budget(remaining_tokens=10, reset_seconds=300.0, limit_tokens=6000)
75
+ d = decide(b, estimated=1000, sigma=100.0, wait_threshold_s=15.0)
76
+ assert d.action == "checkpoint"
77
+ assert d.wait_seconds is None
@@ -0,0 +1,98 @@
1
+ """Behavioral contract for D8: online refill-regime detection + chunked waits.
2
+
3
+ Two refill physics exist in the wild:
4
+ * continuous token bucket (Groq, OpenAI): tokens trickle back — refill-aware
5
+ waits are exact;
6
+ * fixed window: nothing refills until a boundary instant, then everything —
7
+ refill math would under-wait there.
8
+
9
+ The detector classifies the regime from observation alone: consecutive
10
+ telemetry readings with no real call in between. Rising remaining ≈ implied
11
+ rate → continuous; flat → fixed; not enough data → unknown (conservative).
12
+ Chunked waiting supplies those samples for free while ALSO letting the agent
13
+ resume as soon as the bucket actually holds enough.
14
+ """
15
+
16
+ import pytest
17
+
18
+ from agentpause import Budget, decide
19
+ from agentpause.adapters.langgraph import AgentPauseGuard
20
+ from agentpause.refill import RegimeDetector
21
+
22
+
23
+ # -- the detector ---------------------------------------------------------------
24
+
25
+ def test_unknown_until_enough_votes():
26
+ d = RegimeDetector()
27
+ assert d.regime == "unknown"
28
+ d.feed(dt=10.0, dr=1000, implied_rate=100.0) # one vote is not enough
29
+ assert d.regime == "unknown"
30
+
31
+
32
+ def test_rising_remaining_means_continuous():
33
+ d = RegimeDetector()
34
+ d.feed(dt=10.0, dr=950, implied_rate=100.0) # ~95% of implied rate
35
+ d.feed(dt=10.0, dr=1020, implied_rate=100.0)
36
+ assert d.regime == "continuous"
37
+
38
+
39
+ def test_flat_remaining_means_fixed_window():
40
+ d = RegimeDetector()
41
+ d.feed(dt=10.0, dr=0, implied_rate=100.0)
42
+ d.feed(dt=10.0, dr=-11, implied_rate=100.0) # pings nibble a few tokens
43
+ assert d.regime == "fixed"
44
+
45
+
46
+ def test_short_gaps_and_zero_rates_are_ignored():
47
+ d = RegimeDetector(min_gap_s=2.0)
48
+ d.feed(dt=0.5, dr=500, implied_rate=100.0) # gap too short: noise
49
+ d.feed(dt=10.0, dr=500, implied_rate=0.0) # no implied rate: unusable
50
+ assert d.regime == "unknown"
51
+
52
+
53
+ def test_ambiguous_ratios_do_not_vote():
54
+ d = RegimeDetector()
55
+ d.feed(dt=10.0, dr=300, implied_rate=100.0) # 30%: neither clearly
56
+ d.feed(dt=10.0, dr=300, implied_rate=100.0)
57
+ assert d.regime == "unknown"
58
+
59
+
60
+ # -- regime drives the wait math ---------------------------------------------------
61
+
62
+ def test_fixed_regime_waits_the_full_reset():
63
+ """Same numbers that would give a ~10s refill wait — but the provider is
64
+ fixed-window, so the only safe wait is the full reset."""
65
+ b = Budget(remaining_tokens=1000, reset_seconds=12.0, limit_tokens=6000,
66
+ refill_regime="fixed")
67
+ d = decide(b, estimated=1800, sigma=100.0, k=2.0, wait_threshold_s=15.0)
68
+ assert d.action == "wait"
69
+ assert d.wait_seconds == 12.0
70
+
71
+
72
+ def test_continuous_regime_keeps_refill_math():
73
+ b = Budget(remaining_tokens=1000, reset_seconds=50.0, limit_tokens=6000,
74
+ refill_regime="continuous")
75
+ d = decide(b, estimated=1800, sigma=100.0, k=2.0)
76
+ assert d.action == "wait"
77
+ assert d.wait_seconds == pytest.approx(1000 / 100.0 * 1.15)
78
+
79
+
80
+ # -- chunked waiting in the guard ----------------------------------------------------
81
+
82
+ def test_guard_sleeps_in_chunks_and_resumes_early():
83
+ """A 40s computed wait must not be one blind sleep: cap at chunk_s and
84
+ re-read — if the bucket refills sooner, the agent resumes sooner."""
85
+ budget = Budget(remaining_tokens=0, reset_seconds=14.0, limit_tokens=200_000)
86
+ sleeps = []
87
+
88
+ def fake_sleep(s):
89
+ sleeps.append(s)
90
+ if len(sleeps) == 2: # refilled after ~2 chunks
91
+ budget.remaining_tokens = 150_000
92
+
93
+ guard = AgentPauseGuard(telemetry=lambda: budget,
94
+ interrupt_fn=lambda p: None,
95
+ sleep_fn=fake_sleep, chunk_s=5.0)
96
+ guard.check([{"role": "user", "content": "q " * 50}])
97
+ assert len(sleeps) == 2 # resumed at the SECOND re-read
98
+ assert all(s <= 5.5 for s in sleeps) # never a blind long sleep
File without changes
File without changes