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.
- {agentpause-0.2.0/src/agentpause.egg-info → agentpause-0.2.1}/PKG-INFO +16 -1
- {agentpause-0.2.0 → agentpause-0.2.1}/README.md +15 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/pyproject.toml +1 -1
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/__init__.py +3 -1
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/langgraph.py +9 -3
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/litellm.py +13 -0
- agentpause-0.2.1/src/agentpause/adapters/openai_compat.py +206 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/estimator.py +7 -2
- agentpause-0.2.1/src/agentpause/refill.py +75 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/risk.py +55 -6
- {agentpause-0.2.0 → agentpause-0.2.1/src/agentpause.egg-info}/PKG-INFO +16 -1
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/SOURCES.txt +5 -0
- agentpause-0.2.1/tests/test_openai_compat.py +98 -0
- agentpause-0.2.1/tests/test_refill_wait.py +77 -0
- agentpause-0.2.1/tests/test_regime.py +98 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/LICENSE +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/setup.cfg +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/adapters/__init__.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/breaker.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/errors.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/fallback.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/retry.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/scheduler.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause/state.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/dependency_links.txt +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/requires.txt +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/src/agentpause.egg-info/top_level.txt +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_async.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_breaker.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_estimator.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_fallback.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_io_tpm.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_langgraph_adapter.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_litellm_adapter.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_money_and_hooks.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_quantile.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_risk.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_robustness.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_scheduler.py +0 -0
- {agentpause-0.2.0 → agentpause-0.2.1}/tests/test_state.py +0 -0
- {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.
|
|
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.
|
|
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.
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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.
|
|
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
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
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.
|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|