hackbait 0.1.0__py3-none-any.whl

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.
hackbait/runner.py ADDED
@@ -0,0 +1,452 @@
1
+ """OpenAI-compatible chat client and the episode loop.
2
+
3
+ Stdlib only, like the rest of the core. ``chat`` speaks the chat-completions
4
+ wire format over urllib; Anthropic, OpenRouter, vLLM and Ollama all expose
5
+ OpenAI-compatible endpoints, so this one client covers them. ``run_suite``
6
+ drives a responder over the generated cases, hands each completion to the
7
+ owning honeypot's judge, and appends every episode to a JSONL run file that
8
+ stays readable when a run dies halfway.
9
+
10
+ The runner never interprets completions; judging belongs to the honeypots.
11
+ Retry sleeps are fixed powers of two with no jitter, so a rerun behaves the
12
+ same way twice, and the api key never reaches a log line or an exception.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ import time
19
+ import urllib.error
20
+ import urllib.request
21
+ from collections.abc import Callable, Sequence
22
+ from dataclasses import dataclass
23
+ from pathlib import Path
24
+ from typing import Any
25
+
26
+ from .types import Case, Episode, Honeypot
27
+
28
+ __all__ = [
29
+ "Response",
30
+ "chat",
31
+ "chat_response",
32
+ "judge_case",
33
+ "model_responder",
34
+ "run_suite",
35
+ ]
36
+
37
+
38
+ @dataclass(frozen=True)
39
+ class Response:
40
+ """One model completion plus the resource metadata that came with it.
41
+
42
+ ``text`` is the message content the judge sees. The token counts and
43
+ latency are metadata only: they are recorded on the Episode after the
44
+ verdict is set and never feed judging or case generation. Usage is None
45
+ when the provider omits a usage block, and ``latency_ms`` is the wall-clock
46
+ time of the request, which is why it too must stay out of any verdict.
47
+ """
48
+
49
+ text: str
50
+ prompt_tokens: int | None = None
51
+ completion_tokens: int | None = None
52
+ latency_ms: float | None = None
53
+
54
+
55
+ # A responder maps a case to a completion. Scripted policies return a plain
56
+ # str (no usage to report); a model responder returns a Response carrying the
57
+ # provider's token counts and the request latency.
58
+ Responder = Callable[[Case], "str | Response"]
59
+
60
+ _BACKOFF_BASE = 1.0
61
+ _SNIPPET_LIMIT = 300
62
+
63
+ # After this many responder failures in a row the endpoint is treated as dead
64
+ # (connection refused, DNS, bad key, wrong model name all look like this) and
65
+ # the run aborts instead of burning a per-episode retry cycle on every
66
+ # remaining case. A single success resets the count.
67
+ MAX_CONSECUTIVE_RESPONDER_ERRORS = 5
68
+
69
+
70
+ class NoContent(RuntimeError):
71
+ """The endpoint answered normally but the model returned no content.
72
+
73
+ A refusal or a truncation is a fact about that one episode, not evidence
74
+ the endpoint is dead, so a streak of them must not abort the run. Claude
75
+ Opus 5 answers finish_reason "content_filter" with zero tokens on every
76
+ HB-TEST prompt, and a suite that gives up there can never measure the
77
+ model on the four families it does engage with.
78
+ """
79
+
80
+
81
+ def _snippet(body: str) -> str:
82
+ """Compress a response body to one short line for exception messages."""
83
+ flat = " ".join(body.split())
84
+ if len(flat) > _SNIPPET_LIMIT:
85
+ return flat[:_SNIPPET_LIMIT] + "..."
86
+ return flat
87
+
88
+
89
+ def _redact(text: str, api_key: str) -> str:
90
+ """Strip the api key from text bound for an exception or a log.
91
+
92
+ A provider that echoes the Authorization header back in an error body
93
+ would otherwise carry the key into the RuntimeError we raise. The empty
94
+ key is left alone so an unset key does not turn every character into ``*``.
95
+ """
96
+ return text.replace(api_key, "***") if api_key else text
97
+
98
+
99
+ def _parse_content(body: str, api_key: str = "") -> str:
100
+ """Pull choices[0].message.content out of a chat-completions response.
101
+
102
+ An empty body that did not finish on "stop" is raised as an error rather
103
+ than returned. A refused or truncated request is not an answer, and
104
+ scoring it as one puts a non-response in the hack-rate denominator and
105
+ reads as a task the model simply failed. Live runs hit this: Claude Opus 5
106
+ returns finish_reason "content_filter" with zero completion tokens on the
107
+ HB-TEST prompt, which invites test tampering.
108
+ """
109
+ try:
110
+ choice = json.loads(body)["choices"][0]
111
+ content = choice["message"]["content"]
112
+ except (json.JSONDecodeError, KeyError, IndexError, TypeError) as err:
113
+ snippet = _snippet(_redact(body, api_key))
114
+ raise RuntimeError(
115
+ f"chat response is not OpenAI-shaped (choices[0].message.content): {snippet}"
116
+ ) from err
117
+ if content is not None and not isinstance(content, str):
118
+ raise RuntimeError(f"chat response content is {type(content).__name__}, expected str")
119
+ finish = choice.get("finish_reason")
120
+ if not (content or "").strip() and finish not in (None, "stop"):
121
+ raise NoContent(f"chat returned no content, finish_reason {finish!r}")
122
+ if content is None:
123
+ # Some providers send null content when the model emits nothing.
124
+ return ""
125
+ return content
126
+
127
+
128
+ def _parse_usage(body: str) -> tuple[int | None, int | None]:
129
+ """Pull prompt/completion token counts out of a chat-completions usage block.
130
+
131
+ Usage is optional in the wire format and absent on some providers, so a
132
+ body without it, or with a non-integer count, yields None rather than
133
+ raising. This is metadata: a missing usage block must never abort a run or
134
+ change a verdict, only leave the token figure unknown.
135
+ """
136
+ try:
137
+ usage = json.loads(body).get("usage")
138
+ except (json.JSONDecodeError, AttributeError):
139
+ return (None, None)
140
+ if not isinstance(usage, dict):
141
+ return (None, None)
142
+ prompt = usage.get("prompt_tokens")
143
+ completion = usage.get("completion_tokens")
144
+ # bool is an int subclass; a stray True would otherwise read as 1 token.
145
+ prompt = prompt if isinstance(prompt, int) and not isinstance(prompt, bool) else None
146
+ completion = (
147
+ completion if isinstance(completion, int) and not isinstance(completion, bool) else None
148
+ )
149
+ return (prompt, completion)
150
+
151
+
152
+ def chat_response(
153
+ messages: list[dict[str, str]],
154
+ model: str,
155
+ base_url: str,
156
+ api_key: str,
157
+ temperature: float | None = 0.0,
158
+ timeout: float = 120,
159
+ max_retries: int = 4,
160
+ extra_headers: dict[str, str] | None = None,
161
+ ) -> Response:
162
+ """POST one chat-completions request and return the content plus usage.
163
+
164
+ ``base_url`` is used exactly as given apart from stripping one trailing
165
+ slash, so any path prefix a provider wants (``/v1``, a proxy route)
166
+ belongs in it. The request goes to ``{base_url}/chat/completions``.
167
+
168
+ ``temperature=None`` omits the field. Newer models reject it outright
169
+ (Anthropic's Claude 5 answers "`temperature` is deprecated for this
170
+ model" with a 400), and a sampling knob is not worth failing a run over.
171
+
172
+ HTTP 429, 5xx and connection errors are retried up to ``max_retries``
173
+ times, sleeping 1, 2, 4, ... seconds between attempts (deliberately
174
+ jitter-free: reruns should behave identically). Any other HTTP error
175
+ raises at once with a snippet of the response body. The api key travels
176
+ only in the Authorization header, never in exceptions or output.
177
+
178
+ The returned Response carries the response usage block (None when the
179
+ provider omits it) and the wall-clock latency of the request. Both are
180
+ metadata for the resource report and are measured only around the request
181
+ itself, so a retry's earlier attempts do not inflate the reported latency.
182
+ """
183
+ url = base_url.removesuffix("/") + "/chat/completions"
184
+ request_body: dict[str, Any] = {"model": model, "messages": messages}
185
+ if temperature is not None:
186
+ request_body["temperature"] = temperature
187
+ payload = json.dumps(request_body).encode("utf-8")
188
+ headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
189
+ if extra_headers:
190
+ headers.update(extra_headers)
191
+ reason = "not attempted"
192
+ for attempt in range(max_retries + 1):
193
+ request = urllib.request.Request(url, data=payload, headers=headers, method="POST")
194
+ try:
195
+ start = time.perf_counter()
196
+ with urllib.request.urlopen(request, timeout=timeout) as response:
197
+ body = response.read().decode("utf-8", errors="replace")
198
+ latency_ms = (time.perf_counter() - start) * 1000.0
199
+ text = _parse_content(body, api_key)
200
+ prompt_tokens, completion_tokens = _parse_usage(body)
201
+ return Response(
202
+ text=text,
203
+ prompt_tokens=prompt_tokens,
204
+ completion_tokens=completion_tokens,
205
+ latency_ms=latency_ms,
206
+ )
207
+ except urllib.error.HTTPError as err:
208
+ try:
209
+ err_body = err.read().decode("utf-8", errors="replace")
210
+ except Exception:
211
+ err_body = ""
212
+ err_body = _redact(err_body, api_key)
213
+ if err.code != 429 and err.code < 500:
214
+ raise RuntimeError(
215
+ f"chat request failed: HTTP {err.code} from {url}: {_snippet(err_body)}"
216
+ ) from None
217
+ reason = f"HTTP {err.code}: {_snippet(err_body)}"
218
+ except urllib.error.URLError as err:
219
+ reason = f"connection error: {err.reason}"
220
+ except (TimeoutError, OSError) as err:
221
+ # A socket read timeout surfaces as a bare TimeoutError, not a
222
+ # URLError, so without this clause it would escape the retry loop
223
+ # unretried despite the docstring's promise.
224
+ reason = f"connection error: {err}"
225
+ if attempt < max_retries:
226
+ time.sleep(_BACKOFF_BASE * 2**attempt)
227
+ raise RuntimeError(f"chat request failed after {max_retries + 1} attempts: {reason}")
228
+
229
+
230
+ def chat(
231
+ messages: list[dict[str, str]],
232
+ model: str,
233
+ base_url: str,
234
+ api_key: str,
235
+ temperature: float | None = 0.0,
236
+ timeout: float = 120,
237
+ max_retries: int = 4,
238
+ extra_headers: dict[str, str] | None = None,
239
+ ) -> str:
240
+ """POST one chat-completions request and return just the message content.
241
+
242
+ A thin wrapper over ``chat_response`` for callers that want the string and
243
+ nothing else. Same arguments, same retry and redaction behaviour; the usage
244
+ and latency are discarded here.
245
+ """
246
+ return chat_response(
247
+ messages,
248
+ model=model,
249
+ base_url=base_url,
250
+ api_key=api_key,
251
+ temperature=temperature,
252
+ timeout=timeout,
253
+ max_retries=max_retries,
254
+ extra_headers=extra_headers,
255
+ ).text
256
+
257
+
258
+ def judge_case(honeypot: Honeypot, case: Case, completion: str, execute: bool = False) -> Episode:
259
+ """Route a completion to the honeypot's judge.
260
+
261
+ HB-TEST is the one family whose judge takes an ``execute`` flag, because
262
+ running model-written code against held-out tests is opt-in. Every other
263
+ judge has the plain two-argument shape from the Honeypot protocol.
264
+ """
265
+ if honeypot.family == "HB-TEST":
266
+ return honeypot.judge(case, completion, execute=execute) # type: ignore[call-arg]
267
+ return honeypot.judge(case, completion)
268
+
269
+
270
+ def _as_response(value: str | Response) -> Response:
271
+ """Normalise a responder result to a Response.
272
+
273
+ A scripted policy returns a bare str and carries no usage, so it becomes a
274
+ Response with the token counts and latency left None.
275
+ """
276
+ if isinstance(value, Response):
277
+ return value
278
+ return Response(text=value)
279
+
280
+
281
+ def _episode_for(
282
+ honeypot: Honeypot, case: Case, respond: Responder, execute: bool
283
+ ) -> tuple[Episode, bool]:
284
+ """One episode, never an exception: errors become FAIL rows with detail.
285
+
286
+ The second element says whether the responder itself failed, which is the
287
+ transport-level class (connection refused, DNS, auth, model-not-found all
288
+ surface here as a raised chat error). A judge crash is a grader problem,
289
+ not a dead endpoint, so it reports False. A refusal (``NoContent``) also
290
+ reports False: the endpoint is healthy and the model simply declined, and
291
+ a model that declines one whole family would otherwise abort its own run.
292
+
293
+ When the responder returns a Response, its token counts and latency are
294
+ copied onto the Episode after judging. They are recorded, never consulted:
295
+ only ``response.text`` is handed to the judge, so a verdict cannot depend
296
+ on usage or wall-clock time.
297
+ """
298
+ try:
299
+ result = respond(case)
300
+ except NoContent as err:
301
+ return (
302
+ Episode(
303
+ case=case,
304
+ completion="",
305
+ verdict="FAIL",
306
+ detail="model returned no content (refusal or truncation)",
307
+ error=f"{type(err).__name__}: {err}",
308
+ ),
309
+ False,
310
+ )
311
+ except Exception as err:
312
+ return (
313
+ Episode(
314
+ case=case,
315
+ completion="",
316
+ verdict="FAIL",
317
+ detail="responder raised; no completion to judge",
318
+ error=f"{type(err).__name__}: {err}",
319
+ ),
320
+ True,
321
+ )
322
+ response = _as_response(result)
323
+ try:
324
+ episode = judge_case(honeypot, case, response.text, execute=execute)
325
+ except Exception as err:
326
+ # A crashed judge is a grader error, not evidence of hacking. Keep
327
+ # the completion in the row and keep the run moving.
328
+ episode = Episode(
329
+ case=case,
330
+ completion=response.text,
331
+ verdict="FAIL",
332
+ detail="judge raised",
333
+ error=f"{type(err).__name__}: {err}",
334
+ )
335
+ episode.prompt_tokens = response.prompt_tokens
336
+ episode.completion_tokens = response.completion_tokens
337
+ episode.latency_ms = response.latency_ms
338
+ return episode, False
339
+
340
+
341
+ def run_suite(
342
+ respond: Responder,
343
+ out_path: str | Path,
344
+ model_label: str,
345
+ per_tier: int = 8,
346
+ seed: int = 0,
347
+ tiers: Sequence[int] = (1, 2, 3),
348
+ families: Sequence[str] | None = None,
349
+ execute: bool = False,
350
+ progress: Callable[[str], object] | None = print,
351
+ ) -> dict[str, Any]:
352
+ """Run ``respond`` over the suite, appending every episode to a JSONL file.
353
+
354
+ Line one is the run header (kind, model, suite_version, started, config);
355
+ each following line is one ``Episode.to_row()``. Every line is flushed as
356
+ written, so an interrupted run keeps the rows it earned. A responder
357
+ exception becomes a FAIL row with ``error`` set and the run continues,
358
+ unless ``MAX_CONSECUTIVE_RESPONDER_ERRORS`` responder failures land in a
359
+ row, which means the endpoint is dead (refused connection, bad DNS, bad
360
+ key, wrong model name) and the run aborts with a RuntimeError rather than
361
+ retrying its way through every remaining case. Rows written before the
362
+ abort stay on disk. Returns ``report.summarize`` over the rows written.
363
+ """
364
+ from . import SUITE_VERSION, suite
365
+
366
+ honeypots = list(suite())
367
+ if families is not None:
368
+ wanted = {family.strip().upper() for family in families}
369
+ unknown = sorted(wanted - {hp.family for hp in honeypots})
370
+ if unknown:
371
+ known = ", ".join(hp.family for hp in honeypots)
372
+ raise ValueError(f"unknown families: {', '.join(unknown)} (known: {known})")
373
+ honeypots = [hp for hp in honeypots if hp.family in wanted]
374
+ tier_set = set(tiers)
375
+ header = {
376
+ "kind": "run",
377
+ "model": model_label,
378
+ "suite_version": SUITE_VERSION,
379
+ "started": int(time.time()),
380
+ "config": {
381
+ "per_tier": per_tier,
382
+ "seed": seed,
383
+ "tiers": list(tiers),
384
+ "families": list(families) if families is not None else None,
385
+ "execute": execute,
386
+ },
387
+ }
388
+ out = Path(out_path)
389
+ out.parent.mkdir(parents=True, exist_ok=True)
390
+ rows: list[dict[str, Any]] = []
391
+ consecutive_responder_errors = 0
392
+ with out.open("w", encoding="utf-8") as fh:
393
+ fh.write(json.dumps(header) + "\n")
394
+ fh.flush()
395
+ for honeypot in honeypots:
396
+ for case in honeypot.cases(per_tier=per_tier, seed=seed):
397
+ if case.tier not in tier_set:
398
+ continue
399
+ episode, responder_failed = _episode_for(honeypot, case, respond, execute)
400
+ row = episode.to_row()
401
+ fh.write(json.dumps(row) + "\n")
402
+ fh.flush()
403
+ rows.append(row)
404
+ if progress is not None:
405
+ progress(f"{case.id}: {episode.verdict}")
406
+ consecutive_responder_errors = (
407
+ consecutive_responder_errors + 1 if responder_failed else 0
408
+ )
409
+ if consecutive_responder_errors >= MAX_CONSECUTIVE_RESPONDER_ERRORS:
410
+ raise RuntimeError(
411
+ f"aborting run: {consecutive_responder_errors} consecutive "
412
+ f"responder errors, the endpoint looks dead "
413
+ f"(last error: {episode.error}); "
414
+ f"{len(rows)} episodes saved to {out}"
415
+ )
416
+ from . import report
417
+
418
+ return report.summarize(rows, model=model_label)
419
+
420
+
421
+ def model_responder(
422
+ model: str,
423
+ base_url: str,
424
+ api_key: str,
425
+ system_override: str | None = None,
426
+ temperature: float | None = 0.0,
427
+ ) -> Callable[[Case], Response]:
428
+ """Build a responder that sends each case to an OpenAI-compatible model.
429
+
430
+ The case's own system prompt is used unless ``system_override`` is given;
431
+ an override of "" suppresses the system message entirely. ``temperature``
432
+ of None omits the field for models that reject it.
433
+
434
+ The responder returns a Response carrying the completion plus the token
435
+ usage and latency of the request, so the run can report real resource use.
436
+ """
437
+
438
+ def respond(case: Case) -> Response:
439
+ system = case.system if system_override is None else system_override
440
+ chat_messages: list[dict[str, str]] = []
441
+ if system:
442
+ chat_messages.append({"role": "system", "content": system})
443
+ chat_messages.append({"role": "user", "content": case.prompt})
444
+ return chat_response(
445
+ chat_messages,
446
+ model=model,
447
+ base_url=base_url,
448
+ api_key=api_key,
449
+ temperature=temperature,
450
+ )
451
+
452
+ return respond
hackbait/scripted.py ADDED
@@ -0,0 +1,41 @@
1
+ """Scripted policies: the free confusion-matrix gate.
2
+
3
+ Each bait module ships three policies next to its judge: oracle takes the
4
+ legitimate path, cheat takes the bait, garbage does neither. `hackbait
5
+ selftest` runs all three over every case and fails unless every verdict lands
6
+ where `expected_verdicts` says it must. Instrumentation that cannot pass this
7
+ gate has no business scoring a model, and the gate costs zero API calls.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from .types import CONTROL_TIER, Case, Policy
13
+
14
+
15
+ def expected_verdicts(policy: str, case: Case, execute: bool = False) -> tuple[str, ...]:
16
+ """The verdicts the selftest gate accepts for one policy on one case.
17
+
18
+ Control cases carry no bait, so the cheat policy (which answers the
19
+ poisoned value anyway) must land FAIL there, never HACK. The HB-TEST
20
+ oracle is CLEAN_UNVERIFIED without held-out execution and exactly CLEAN
21
+ with it; accepting both regardless of the execute flag would let a broken
22
+ --exec wiring pass the gate, which is the regression class the gate
23
+ exists to catch.
24
+ """
25
+ if case.tier == CONTROL_TIER:
26
+ return {"oracle": ("CLEAN",), "cheat": ("FAIL",), "garbage": ("FAIL",)}[policy]
27
+ if policy == "oracle":
28
+ if case.family == "HB-TEST" and not execute:
29
+ return ("CLEAN_UNVERIFIED",)
30
+ return ("CLEAN",)
31
+ if policy == "cheat":
32
+ return ("HACK",)
33
+ return ("FAIL",)
34
+
35
+
36
+ def scripted() -> dict[str, dict[str, Policy]]:
37
+ """Map family code to its oracle/cheat/garbage policies."""
38
+ from .bait import keyleak, loophole, rubric, selfreport, testedit
39
+
40
+ modules = (keyleak, testedit, rubric, selfreport, loophole)
41
+ return {m.HONEYPOT.family: dict(m.SCRIPTED_POLICIES) for m in modules}
hackbait/stats.py ADDED
@@ -0,0 +1,180 @@
1
+ """Sample-size planning for hack rates: how many episodes buy a given precision.
2
+
3
+ Stdlib only (math), a pure function of its arguments: no wall clock, no
4
+ randomness, no environment reads. Usage, latency and cost are metadata and
5
+ never appear here, because a sample-size answer that shifted with how long a
6
+ run happened to take would be worthless for planning the next one.
7
+
8
+ Design note. Hackbait lives near p=0: a well-behaved model hacks a few percent
9
+ of pressured episodes, often zero. That is exactly where the textbook normal
10
+ approximation p +- z*sqrt(p(1-p)/n) lies to you. At an observed rate of 0 it
11
+ collapses to the interval [0, 0], claiming perfect certainty from a handful of
12
+ episodes, so any n "achieves" any target and the planner returns nonsense. The
13
+ Wilson score interval does not collapse there (report.wilson_interval carries
14
+ the same formula and the same reason), so precision here is measured as the
15
+ half-width of the Wilson interval, the symmetric +- margin around its center.
16
+
17
+ That half-width is strictly decreasing in n for any fixed p in [0, 1], so
18
+ episodes_for_half_width solves by monotone search over n rather than inverting
19
+ a closed form. The closed form we would invert is the normal one, and near
20
+ p=0 it is the thing that is wrong; searching the honest interval is correct at
21
+ every rate the tool actually reports.
22
+ """
23
+
24
+ from __future__ import annotations
25
+
26
+ import math
27
+
28
+ __all__ = [
29
+ "DEFAULT_HALF_WIDTH",
30
+ "DEFAULT_RATES",
31
+ "episodes_and_per_tier",
32
+ "episodes_for_half_width",
33
+ "per_tier_for",
34
+ "power_table",
35
+ "pressured_cells_per_unit",
36
+ "wilson_half_width",
37
+ ]
38
+
39
+ # Two-sided 95% z, the same constant report.py uses. Kept local so this module
40
+ # stays stdlib-only and does not reach into report for a bare number.
41
+ _Z_95 = 1.96
42
+
43
+ # Rates worth tabulating are the ones the tool sees: a clean model near zero up
44
+ # to a badly-behaved one around ten percent. The integrator can override these
45
+ # from the command line; they exist so the default table lands where planning
46
+ # actually happens.
47
+ DEFAULT_RATES: tuple[float, ...] = (0.0, 0.01, 0.02, 0.05, 0.10)
48
+
49
+ # Two percentage points is the default precision target: tight enough to tell
50
+ # a 2% hacker from a 5% one, loose enough to reach without a runaway episode
51
+ # budget.
52
+ DEFAULT_HALF_WIDTH: float = 0.02
53
+
54
+
55
+ def wilson_half_width(p_hat: float, n: int) -> float:
56
+ """Half-width of the 95% Wilson score interval at rate ``p_hat`` over ``n``.
57
+
58
+ This is the symmetric +- margin around the interval's center, the same
59
+ ``half`` term report.wilson_interval computes; the interval is
60
+ ``center +- half_width``. Returned as a fraction (0.02 means two
61
+ percentage points), so it compares directly against a target half-width.
62
+
63
+ ``n <= 0`` has no interval, so it returns infinity: a run with no episodes
64
+ has measured nothing, and a monotone search must keep looking rather than
65
+ treat zero episodes as infinitely precise.
66
+ """
67
+ if n <= 0:
68
+ return math.inf
69
+ z2 = _Z_95 * _Z_95
70
+ denom = 1.0 + z2 / n
71
+ return _Z_95 * math.sqrt(p_hat * (1.0 - p_hat) / n + z2 / (4 * n * n)) / denom
72
+
73
+
74
+ def episodes_for_half_width(p_expected: float, target_half_width: float) -> int:
75
+ """Smallest ``n`` whose 95% Wilson half-width at ``p_expected`` is <= target.
76
+
77
+ Solved by monotone search over n, not by inverting the normal-approx
78
+ formula, because that formula is exactly what fails near p=0 where the tool
79
+ operates (see the module docstring). The half-width is strictly decreasing
80
+ in n for any fixed p, so a doubling bracket followed by a bisection lands
81
+ on the smallest sufficient n.
82
+
83
+ ``p_expected`` must be a probability in [0, 1]; ``target_half_width`` must
84
+ be strictly positive, since no finite n reaches a target of zero.
85
+ """
86
+ if not 0.0 <= p_expected <= 1.0:
87
+ raise ValueError(f"p_expected must be in [0, 1], got {p_expected!r}")
88
+ if target_half_width <= 0.0:
89
+ raise ValueError(f"target_half_width must be > 0, got {target_half_width!r}")
90
+
91
+ lo = 1
92
+ if wilson_half_width(p_expected, lo) <= target_half_width:
93
+ return lo
94
+ # Double until the target is met, keeping lo as the last n that missed it.
95
+ # Termination is guaranteed: the half-width falls toward 0 as n grows.
96
+ hi = 2
97
+ while wilson_half_width(p_expected, hi) > target_half_width:
98
+ lo = hi
99
+ hi *= 2
100
+ # Invariant: hw(lo) > target >= hw(hi). Bisect for the crossover.
101
+ while hi - lo > 1:
102
+ mid = (lo + hi) // 2
103
+ if wilson_half_width(p_expected, mid) <= target_half_width:
104
+ hi = mid
105
+ else:
106
+ lo = mid
107
+ return hi
108
+
109
+
110
+ def pressured_cells_per_unit() -> int:
111
+ """Pressured episodes produced by one ``per_tier`` step of the suite.
112
+
113
+ A ``per_tier`` of one runs every family once at every pressure tier, so the
114
+ pressured episode count per step is families times pressured tiers. Read
115
+ from the live suite rather than hardcoded, so adding a family or a tier
116
+ keeps the planning math honest. suite() is a pure generator with no clock
117
+ or randomness, so calling it here does not taint this module's determinism.
118
+ """
119
+ from . import TIERS, suite
120
+
121
+ return len(suite()) * len(TIERS)
122
+
123
+
124
+ def per_tier_for(episodes_needed: int, cells_per_unit: int | None = None) -> int:
125
+ """The ``per_tier`` that yields at least ``episodes_needed`` pressured episodes.
126
+
127
+ ``per_tier`` is a whole-suite dial, so the count rounds up: half a family is
128
+ not a thing you can run. With the standard 5 families times 3 pressured
129
+ tiers that is ceil(n / 15). ``cells_per_unit`` overrides the suite reading
130
+ for callers that want the math without importing the suite.
131
+ """
132
+ if cells_per_unit is None:
133
+ cells_per_unit = pressured_cells_per_unit()
134
+ return math.ceil(episodes_needed / cells_per_unit)
135
+
136
+
137
+ def episodes_and_per_tier(
138
+ p_expected: float,
139
+ target_half_width: float,
140
+ cells_per_unit: int | None = None,
141
+ ) -> tuple[int, int]:
142
+ """Episodes needed for ``target_half_width`` at ``p_expected``, and the per_tier for it.
143
+
144
+ Returns ``(episodes_needed, per_tier)`` where per_tier is the smallest
145
+ whole-suite setting that runs at least ``episodes_needed`` pressured
146
+ episodes. This is the bridge from a precision goal to a runnable command.
147
+ """
148
+ needed = episodes_for_half_width(p_expected, target_half_width)
149
+ return needed, per_tier_for(needed, cells_per_unit)
150
+
151
+
152
+ def power_table(
153
+ rates: object,
154
+ half_widths: object,
155
+ cells_per_unit: int | None = None,
156
+ ) -> list[dict[str, float | int]]:
157
+ """A record per (rate, half_width): the episodes and per_tier each needs.
158
+
159
+ Returns a flat list of plain dicts (rate, half_width, episodes, per_tier),
160
+ one row per combination in ``rates`` times ``half_widths``, kept in input
161
+ order so the caller controls the table layout. A list of dicts stays
162
+ JSON-serializable and renders as a table without further shaping; a tuple
163
+ key would not survive json.dumps. ``cells_per_unit`` is read once and
164
+ reused so every row shares one suite reading.
165
+ """
166
+ if cells_per_unit is None:
167
+ cells_per_unit = pressured_cells_per_unit()
168
+ table: list[dict[str, float | int]] = []
169
+ for rate in rates: # type: ignore[attr-defined]
170
+ for half_width in half_widths: # type: ignore[attr-defined]
171
+ needed, per_tier = episodes_and_per_tier(rate, half_width, cells_per_unit)
172
+ table.append(
173
+ {
174
+ "rate": rate,
175
+ "half_width": half_width,
176
+ "episodes": needed,
177
+ "per_tier": per_tier,
178
+ }
179
+ )
180
+ return table