mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
learn: experiment fill -- fill a rule experiment's response cache with a hard ceiling, dry-run estimate
Implemented with Codex; review fix: the bounded client keeps the configured thinking setting so a fill answers as the run would. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
212a391872
commit
8d14b0c13f
@@ -252,6 +252,32 @@ The reading is on [Where we stand](where-we-stand.md#2-our-approach-to-detection
|
||||
| `ousast learn curve --program ID` | the metric over memory size on a fixed candidate subset |
|
||||
| `ousast learn audit-leaks` | how well each signal alone separates the sides of a pair (no model) |
|
||||
| `ousast learn experiment register\|units\|run\|analyse MANIFEST` | A/B experiments from a pre-registered manifest (`plane/experiments/<id>.yaml`): record its digest, freeze the units (no model), run the arms paired (arm A may replay from the cache at $0), analyse (repository bootstrap, exact McNemar, two looks, the adoption verdict) |
|
||||
| `ousast learn experiment fill MANIFEST --budget-usd X` | Fill missing responses for a registered `kind: rule` experiment's frozen units before its zero-cost run; `--dry-run` inspects and estimates without calls or writes |
|
||||
|
||||
`experiment fill` accepts `--memory URL`, `--profile static|plane` (default `static`),
|
||||
`--manifest-model PATH` (default `plane/models/deepseek-flash.yaml`) and `--out PATH`.
|
||||
Freeze, commit and register the units first. Fill uses the incumbent programs and the
|
||||
same retrieval, sampling and cache keys as replay. The first missing request is the
|
||||
smoke call; a successful smoke response is kept. Any provider failure stops the fill.
|
||||
|
||||
The JSON report on stdout gives each family's units, distinct requests, initially
|
||||
cached and missing responses, newly filled responses and still-missing responses,
|
||||
plus total client calls and metered USD. `--out` also saves it, except during dry-run,
|
||||
which writes nothing. Exit codes are 0 for completion or dry-run, 1 for a provider/cache
|
||||
failure, 2 for invalid setup and 3 for budget exhaustion. Resume a partial fill with a
|
||||
new invocation and budget; that budget covers only the new invocation, including smoke.
|
||||
|
||||
Dry-run prices each missing prompt using the program's input-token estimate and the
|
||||
family's mean cached completion-token usage, falling back to **300 output tokens**.
|
||||
It reports potential retries separately as `conditional_retries`: an unseen answer
|
||||
may need one parse retry, whose necessity cannot be known without calling the model.
|
||||
Known invalid cached answers include their retry requests in the inventory.
|
||||
|
||||
Fill reserves at least **$0.01** and the largest observed call cost before each call.
|
||||
For the production client it also reserves a conservative input-byte token bound plus
|
||||
the program's enforced output-token limit; hidden provider retries are disabled.
|
||||
Consequently a small budget may stop before making any call, and a fill may finish
|
||||
with unused budget. Calls without token usage stop as failures, never as free calls.
|
||||
|
||||
The memory store is the plane's (`OUSAST_MEMORY`; see [ops/ax/README.md](ops/ax/README.md)).
|
||||
The harvests that produced the plane signals ran as plane Runs generated by
|
||||
|
||||
@@ -60,18 +60,24 @@ def add_commands(learn_sub: Any) -> None:
|
||||
("register", "record the manifest's and the units file's sha256 as the experiment row (both committed at HEAD)"),
|
||||
("units", "freeze the units of a program experiment from the memory (no model): the families' evaluation candidates"),
|
||||
("run", "run every arm on every unit, paired; arm A replays from the response cache when the manifest says so"),
|
||||
("fill", "fill a rule experiment's response cache under a separate budget, or estimate with --dry-run"),
|
||||
("analyse", "paired cluster bootstrap, exact McNemar, the two looks and the adoption verdict (counts and intervals only)"),
|
||||
):
|
||||
sub = experiment_sub.add_parser(name, help=text)
|
||||
sub.add_argument("manifest", type=Path, help="plane/experiments/<id>.yaml")
|
||||
sub.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
||||
if name in ("units", "run"):
|
||||
if name in ("units", "run", "fill"):
|
||||
sub.add_argument("--profile", choices=("static", "plane"), default="static")
|
||||
if name == "run":
|
||||
if name in ("run", "fill"):
|
||||
sub.add_argument(
|
||||
"--manifest-model", type=Path, default=Path("plane/models/deepseek-flash.yaml"), help="the chat Model (prices)"
|
||||
)
|
||||
if name == "run":
|
||||
sub.add_argument("--no-early-stop", action="store_true", help="take no first look during the run")
|
||||
if name == "fill":
|
||||
sub.add_argument("--budget-usd", type=float, help="required for paid fill; includes the smoke call")
|
||||
sub.add_argument("--dry-run", action="store_true", help="inspect and estimate only; no calls or writes, including --out")
|
||||
sub.add_argument("--out", type=Path, help="also write the fill report as JSON (ignored for --dry-run)")
|
||||
if name in ("run", "analyse"):
|
||||
sub.add_argument("--resamples", type=int, help="bootstrap resamples (default: the manifest's)")
|
||||
sub.add_argument("--out", type=Path, help="write the counts-only report here (analyse --record: benchmarks/experiments/<id>/)")
|
||||
@@ -154,6 +160,33 @@ def _experiment(args: argparse.Namespace) -> int:
|
||||
data = store.get_blob("excerpts", sha)
|
||||
return data.decode("utf-8") if data is not None else None
|
||||
|
||||
if args.experiment_command == "fill":
|
||||
from ..tool_hunter import ChatClient
|
||||
from .fill import fill_rule
|
||||
|
||||
def resolve_fill_client() -> ChatClient:
|
||||
from ..config import load_config
|
||||
from ..model.endpoint import resolve_chat_endpoint
|
||||
|
||||
resolved = resolve_chat_endpoint(load_config(None))
|
||||
if resolved is None:
|
||||
raise ex.ExperimentError("no chat endpoint configured; nothing was spent")
|
||||
return resolved[0]
|
||||
|
||||
report = fill_rule(
|
||||
store,
|
||||
manifest,
|
||||
memory,
|
||||
parameters,
|
||||
excerpt_text,
|
||||
vectors,
|
||||
client_factory=resolve_fill_client,
|
||||
budget_usd=args.budget_usd,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
_write(report, None if args.dry_run else args.out)
|
||||
return {"done": 0, "dry_run": 0, "unfinished": 3, "failed": 1}[report["status"]]
|
||||
|
||||
if manifest.kind == "rule": # zero cost: cache only, behind a zero-ceiling meter; no endpoint is resolved
|
||||
caller, meter = ex.zero_cost_caller(_arm_model(manifest, manifest.arms["A"], store), parameters, store)
|
||||
summary = ex.run_rule(store, manifest, memory, caller, meter, excerpt_text, vectors)
|
||||
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Fill rule-experiment responses without writing scores or changing the frozen population.
|
||||
|
||||
Estimates use Prompt.estimated_tokens and the family's mean cached completion usage,
|
||||
or 300 output tokens when no usage is available. Unknown answers may require one retry;
|
||||
dry-run counts those separately as conditional requests, never fabricating an answer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from ..model.endpoint import DeepSeekChatClient
|
||||
from ..plane.budget import BudgetExhausted, MeteredClient, prices_from
|
||||
from ..plane.memory import MemoryStore
|
||||
from ..tool_hunter import ChatClient, ChatResponse
|
||||
from . import experiments as ex
|
||||
from .examples import Example
|
||||
from .program import Caller, Program, Prompt, parse_answer, request_key
|
||||
|
||||
DEFAULT_OUTPUT_TOKENS = 300
|
||||
DEFAULT_MARGIN_USD = 0.01
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Request:
|
||||
family: str
|
||||
prompt: Prompt
|
||||
temperature: float
|
||||
sample: str
|
||||
|
||||
def key(self, caller: Caller) -> str:
|
||||
return request_key(caller.model, caller.params_digest, self.prompt.messages, self.temperature, self.sample, True)
|
||||
|
||||
def ask(self, caller: Caller) -> str:
|
||||
return caller.ask(self.prompt.messages, temperature=self.temperature, sample=self.sample)
|
||||
|
||||
def retry(self) -> Request:
|
||||
return Request(self.family, self.prompt, 0.0, f"{self.sample}:retry")
|
||||
|
||||
|
||||
class FillMeter(MeteredClient):
|
||||
"""Reserve a whole call before admission, including the first (smoke) call.
|
||||
|
||||
The fixed cent is a floor, never replaced by a smaller observed cost. For the
|
||||
production adapter, UTF-8 bytes conservatively bound input tokens (plus 256 framing
|
||||
tokens per message), and the provider enforces the output limit. Injected clients
|
||||
must honour the same token/usage contract. No hidden transport retries are allowed.
|
||||
"""
|
||||
|
||||
def __init__(self, client: ChatClient, parameters: Mapping[str, Any], budget: float, output_limit: int) -> None:
|
||||
super().__init__(client, prices=parameters, budget_usd=budget)
|
||||
self.limit = budget
|
||||
self.output_limit = output_limit
|
||||
self.largest_call_usd = 0.0
|
||||
|
||||
def complete(self, **kwargs: Any) -> ChatResponse:
|
||||
assert self.prices is not None
|
||||
messages = kwargs["messages"]
|
||||
input_bound = sum(len(str(m["content"]).encode("utf-8")) + 256 for m in messages)
|
||||
bound = (
|
||||
input_bound * max(self.prices.input_per_m, self.prices.cache_hit_per_m) + self.output_limit * self.prices.output_per_m
|
||||
) / 1_000_000
|
||||
reserve = max(DEFAULT_MARGIN_USD, self.largest_call_usd, bound)
|
||||
self.budget_usd = max(0.0, self.limit - reserve)
|
||||
before = self.usd or 0.0
|
||||
tokens_before = sum(self.usage.values())
|
||||
try:
|
||||
response = super().complete(**kwargs)
|
||||
except BudgetExhausted:
|
||||
raise
|
||||
except Exception:
|
||||
self.calls += 1 # attempted provider calls count even when no response was returned
|
||||
raise
|
||||
finally:
|
||||
self.largest_call_usd = max(self.largest_call_usd, (self.usd or 0.0) - before)
|
||||
if sum(self.usage.values()) <= tokens_before:
|
||||
raise RuntimeError("provider returned no metered token usage")
|
||||
return response
|
||||
|
||||
|
||||
def fill_rule(
|
||||
store: MemoryStore,
|
||||
manifest: ex.Manifest,
|
||||
memory: Sequence[Example],
|
||||
parameters: Mapping[str, Any],
|
||||
excerpt_text: Callable[[str], str | None],
|
||||
vectors: Mapping[str, Sequence[float]] | None = None,
|
||||
*,
|
||||
client: ChatClient | None = None,
|
||||
client_factory: Callable[[], ChatClient] | None = None,
|
||||
budget_usd: float | None = None,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Inspect or fill every frozen unit. The first cache miss is the smoke call.
|
||||
|
||||
A successful smoke is retained in the cache; any provider error stops immediately
|
||||
and is reported without its message, which may contain credentials. A factory is
|
||||
resolved only after validation, and only when a paid request is actually missing.
|
||||
"""
|
||||
from .evaluate import candidate_from_example
|
||||
|
||||
if manifest.kind != "rule":
|
||||
raise ex.ExperimentError("experiment fill supports kind: rule only")
|
||||
if not dry_run and (budget_usd is None or not math.isfinite(budget_usd) or budget_usd < 0):
|
||||
raise ex.ExperimentError("--budget-usd must be a finite non-negative number (or use --dry-run)")
|
||||
prices = prices_from(parameters)
|
||||
if prices is None or any(not math.isfinite(p) or p < 0 for p in (prices.input_per_m, prices.output_per_m, prices.cache_hit_per_m)):
|
||||
raise ex.ExperimentError("fill requires finite, non-negative Model prices")
|
||||
ex.check_registered(store, manifest)
|
||||
units, folds, _ = ex._checked_units(manifest, memory)
|
||||
programs = {f: Program(ex.arm_spec(manifest, manifest.arms["A"], f, store), memory, excerpt_text, vectors) for f in manifest.families}
|
||||
models = {p.spec.model for p in programs.values()}
|
||||
if len(models) != 1:
|
||||
raise ex.ExperimentError("the rule programs must use the same model")
|
||||
caller = Caller(None, next(iter(models)), parameters, store=store)
|
||||
index = {e.id: e for e in memory}
|
||||
candidate_of = candidate_from_example(excerpt_text, vectors)
|
||||
requests: list[Request] = []
|
||||
for unit in sorted(units, key=lambda u: (u.group, u.family, u.unit)):
|
||||
candidate = candidate_of(index[unit.unit])
|
||||
if candidate is None:
|
||||
raise ex.ExperimentError(f"{manifest.id}: unit {unit.unit[:12]} has no excerpt in the store")
|
||||
for prompt, _, temperature, sample in programs[unit.family].prepare_samples(
|
||||
candidate, folds[unit.fold], seed=int(manifest.folds.get("seed", 0))
|
||||
):
|
||||
requests.append(Request(unit.family, prompt, temperature, sample))
|
||||
|
||||
# Request counts are distinct cache keys per family, even when multiple units share one.
|
||||
def inventory() -> dict[str, dict[str, Request]]:
|
||||
result: dict[str, dict[str, Request]] = {f: {} for f in manifest.families}
|
||||
for req in requests:
|
||||
result[req.family][req.key(caller)] = req
|
||||
entry = caller._cached(req.key(caller))
|
||||
if entry is not None and parse_answer(entry.get("content"), req.prompt.valid_lines) is None:
|
||||
retry = req.retry()
|
||||
result[req.family][retry.key(caller)] = retry
|
||||
return result
|
||||
|
||||
initial = inventory()
|
||||
initially_cached = {key for cells in initial.values() for key in cells if caller._cached(key) is not None}
|
||||
families: dict[str, Any] = {}
|
||||
for family, cells in initial.items():
|
||||
outputs = [
|
||||
float((entry.get("usage") or {})["completion_tokens"])
|
||||
for key in cells
|
||||
if (entry := caller._cached(key)) is not None and "completion_tokens" in (entry.get("usage") or {})
|
||||
]
|
||||
mean_output = sum(outputs) / len(outputs) if outputs else DEFAULT_OUTPUT_TOKENS
|
||||
missing = [req for key, req in cells.items() if key not in initially_cached]
|
||||
families[family] = {
|
||||
"units": sum(u.family == family for u in units),
|
||||
"requests": len(cells),
|
||||
"cached": len(cells) - len(missing),
|
||||
"missing": len(missing),
|
||||
"filled": 0,
|
||||
"still_missing": len(missing),
|
||||
"estimated_usd": sum(
|
||||
(r.prompt.estimated_tokens * prices.input_per_m + mean_output * prices.output_per_m) / 1_000_000 for r in missing
|
||||
),
|
||||
"estimated_output_tokens": mean_output,
|
||||
"conditional_retries": sum(not r.sample.endswith(":retry") for r in missing),
|
||||
}
|
||||
report: dict[str, Any] = {
|
||||
"experiment": manifest.id,
|
||||
"status": "dry_run" if dry_run else "done",
|
||||
"families": families,
|
||||
"client_calls": 0,
|
||||
"usd": 0.0,
|
||||
"reason": "",
|
||||
}
|
||||
if dry_run or not any(c["missing"] for c in families.values()):
|
||||
return report
|
||||
meter: FillMeter | None = None
|
||||
try:
|
||||
if client is None and client_factory is not None:
|
||||
client = client_factory()
|
||||
if client is None:
|
||||
raise ex.ExperimentError("no chat endpoint configured; nothing was spent")
|
||||
output_limit = max(p.spec.max_output_tokens for p in programs.values())
|
||||
if isinstance(client, DeepSeekChatClient):
|
||||
client = client.bounded(output_limit)
|
||||
assert budget_usd is not None
|
||||
meter = FillMeter(client, parameters, budget_usd, output_limit)
|
||||
caller.client = meter
|
||||
for req in requests:
|
||||
content = req.ask(caller)
|
||||
if parse_answer(content, req.prompt.valid_lines) is None:
|
||||
req.retry().ask(caller)
|
||||
except BudgetExhausted:
|
||||
report.update(status="unfinished", reason="budget exhausted; resume with a new fill budget")
|
||||
except ex.ExperimentError:
|
||||
raise
|
||||
except Exception:
|
||||
# Never echo provider exceptions: response bodies and URLs may contain secrets.
|
||||
report.update(status="failed", reason="provider or cache failure; fill stopped")
|
||||
final = inventory()
|
||||
for family, cells in final.items():
|
||||
cached = {key for key in cells if caller._cached(key) is not None}
|
||||
families[family].update(requests=len(cells), filled=len(cached - initially_cached), still_missing=len(cells) - len(cached))
|
||||
if meter is not None:
|
||||
report.update(client_calls=meter.calls, usd=meter.usd)
|
||||
return report
|
||||
@@ -506,20 +506,24 @@ class Program:
|
||||
out.append((prompt, got))
|
||||
return out
|
||||
|
||||
def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision:
|
||||
def prepare_samples(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> list[tuple[Prompt, Retrieval, float, str]]:
|
||||
"""The exact prompts, temperatures and sample labels shared by execution and cache planning."""
|
||||
if self.spec.sampling == RETRIEVAL_ENSEMBLE:
|
||||
prepared = self.prepare_ensemble(candidate, fold, seed=seed)
|
||||
answers = [
|
||||
sample_answer(prompt, caller, temperature=TEMPERATURES[0], sample=str(j), families=self.families)
|
||||
for j, (prompt, _) in enumerate(prepared)
|
||||
]
|
||||
got = prepared[0][1]
|
||||
return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours)
|
||||
return [(prompt, got, TEMPERATURES[0], str(j)) for j, (prompt, got) in enumerate(prepared)]
|
||||
if self.spec.sampling != SAMPLING[0]:
|
||||
raise ValueError(f"unknown sampling {self.spec.sampling!r} (one of {', '.join(SAMPLING)})")
|
||||
prompt, got = self.prepare(candidate, fold, seed=seed)
|
||||
decision = classify(prompt, caller, self.spec.k, families=self.families)
|
||||
return replace(decision, retrieval=got.mode, neighbours=got.neighbours)
|
||||
return [(prompt, got, TEMPERATURES[0] if j == 0 else TEMPERATURES[1], str(j)) for j in range(self.spec.k)]
|
||||
|
||||
def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision:
|
||||
prepared = self.prepare_samples(candidate, fold, seed=seed)
|
||||
answers = [
|
||||
sample_answer(prompt, caller, temperature=temperature, sample=sample, families=self.families)
|
||||
for prompt, _, temperature, sample in prepared
|
||||
]
|
||||
got = prepared[0][1]
|
||||
return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours)
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
@@ -96,6 +96,21 @@ class DeepSeekChatClient:
|
||||
self._disable_thinking = disable_thinking
|
||||
self.endpoint = endpoint
|
||||
self.usage: list[dict[str, object]] = [] # one entry per call, for the round's cost meter
|
||||
self._output_limit: int | None = None
|
||||
|
||||
def bounded(self, max_output_tokens: int) -> DeepSeekChatClient:
|
||||
"""A single-attempt, output-capped copy for callers reserving a hard per-call budget.
|
||||
|
||||
Empty JSON retries belong to the caller's meter, and transport failures must
|
||||
stop that caller before another potentially billable attempt.
|
||||
"""
|
||||
from dataclasses import replace
|
||||
|
||||
if max_output_tokens <= 0:
|
||||
raise ValueError("max_output_tokens must be positive")
|
||||
client = DeepSeekChatClient(replace(self._client, max_attempts=1), disable_thinking=self._disable_thinking, endpoint=self.endpoint)
|
||||
client._output_limit = max_output_tokens
|
||||
return client
|
||||
|
||||
def complete(
|
||||
self,
|
||||
@@ -117,7 +132,7 @@ class DeepSeekChatClient:
|
||||
logprobs=logprobs,
|
||||
temperature=temperature,
|
||||
)
|
||||
if json_object and not (response.content or "").strip():
|
||||
if self._output_limit is None and json_object and not (response.content or "").strip():
|
||||
response = self._call(
|
||||
model=model,
|
||||
messages=messages,
|
||||
@@ -147,6 +162,8 @@ class DeepSeekChatClient:
|
||||
temperature: float | None = None,
|
||||
) -> ChatResponse:
|
||||
extra: dict[str, object] = {}
|
||||
if self._output_limit is not None:
|
||||
extra["max_tokens"] = self._output_limit
|
||||
if self._disable_thinking:
|
||||
extra["thinking"] = {"type": "disabled"} # otherwise temperature is silently ignored
|
||||
if temperature is not None:
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
"""Rule cache filling uses scripted clients and an isolated local store; no provider is contacted."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from learn_fixtures import ScriptedChat, default_answer
|
||||
from test_learn_experiments import CODE, EXAMPLES, PRICES, excerpt_text, paid
|
||||
from test_learn_rules import prepared, write_rule_manifest
|
||||
|
||||
from openultrasast.cli import main
|
||||
from openultrasast.learn import experiments as ex
|
||||
from openultrasast.learn.fill import fill_rule
|
||||
from openultrasast.learn.program import Candidate, Program
|
||||
from openultrasast.learn.retrieve import Target
|
||||
from openultrasast.plane.memory import FileStore
|
||||
|
||||
|
||||
def snapshot(root: Path) -> dict[str, bytes]:
|
||||
return {str(p.relative_to(root)): p.read_bytes() for p in root.rglob("*") if p.is_file()}
|
||||
|
||||
|
||||
def seed_one(store: FileStore, manifest: ex.Manifest, unit: ex.Unit) -> int:
|
||||
chat, caller = paid()
|
||||
caller.store = store
|
||||
example = next(e for e in EXAMPLES if e.id == unit.unit)
|
||||
folds = {f.name: f for f in ex.program_folds(EXAMPLES, manifest)[0]}
|
||||
program = Program(ex.arm_spec(manifest, manifest.arms["A"], unit.family, store), EXAMPLES, excerpt_text)
|
||||
program.decide(Candidate(Target.of(example), CODE[example.excerpt_sha], "python"), folds[unit.fold], caller, seed=0)
|
||||
return len(chat.calls)
|
||||
|
||||
|
||||
def total(report: dict[str, Any], key: str) -> int:
|
||||
return sum(f[key] for f in report["families"].values())
|
||||
|
||||
|
||||
def test_dry_run_makes_zero_calls_and_writes_nothing(tmp_path: Path) -> None:
|
||||
store, manifest, units = prepared(tmp_path)
|
||||
cached = seed_one(store, manifest, units[0])
|
||||
chat = ScriptedChat()
|
||||
before = snapshot(tmp_path)
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, dry_run=True)
|
||||
assert report["status"] == "dry_run" and report["client_calls"] == 0 and report["usd"] == 0
|
||||
assert not chat.calls and snapshot(tmp_path) == before
|
||||
assert total(report, "units") == len(units)
|
||||
assert total(report, "requests") == 2 * len(units)
|
||||
assert total(report, "cached") == cached
|
||||
assert total(report, "missing") == 2 * len(units) - cached
|
||||
assert total(report, "filled") == 0
|
||||
assert all(f["estimated_usd"] > 0 for f in report["families"].values() if f["missing"])
|
||||
|
||||
|
||||
def test_fill_only_missing_requests_then_replay_is_complete(tmp_path: Path) -> None:
|
||||
store, manifest, units = prepared(tmp_path)
|
||||
cached = seed_one(store, manifest, units[0])
|
||||
chat = ScriptedChat()
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=1.0)
|
||||
assert report["status"] == "done" and total(report, "still_missing") == 0
|
||||
assert total(report, "filled") == report["client_calls"] == len(chat.calls) == 2 * len(units) - cached
|
||||
assert 0 < report["usd"] < 1.0
|
||||
before = snapshot(tmp_path)
|
||||
again = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=1.0)
|
||||
assert again["client_calls"] == 0 and total(again, "filled") == 0 and snapshot(tmp_path) == before
|
||||
caller, meter = ex.zero_cost_caller("deepseek-flash", PRICES, store)
|
||||
replay = ex.run_rule(store, manifest, EXAMPLES, caller, meter, excerpt_text)
|
||||
assert replay.decided == {"score": len(units)} and replay.skipped["replay_miss"] == 0
|
||||
assert meter.calls == 0 and meter.usd == 0
|
||||
|
||||
|
||||
class FixedCostChat(ScriptedChat):
|
||||
def complete(self, **kwargs: Any) -> Any:
|
||||
response = super().complete(**kwargs)
|
||||
self.usage[-1] = {"prompt_tokens": 0, "completion_tokens": 6_000}
|
||||
return response
|
||||
|
||||
|
||||
def test_hard_ceiling_stops_before_a_non_multiple_budget_is_exceeded(tmp_path: Path) -> None:
|
||||
store, manifest, _ = prepared(tmp_path)
|
||||
chat = FixedCostChat()
|
||||
prices = {"cache_hit_per_m": 0.0, "input_per_m": 0.0, "output_per_m": 1.0}
|
||||
report = fill_rule(store, manifest, EXAMPLES, prices, excerpt_text, client=chat, budget_usd=0.025)
|
||||
assert report["status"] == "unfinished"
|
||||
assert 0 < report["usd"] <= 0.025
|
||||
assert report["usd"] == pytest.approx(len(chat.calls) * 0.006)
|
||||
assert report["client_calls"] == total(report, "filled") == len(chat.calls)
|
||||
assert total(report, "still_missing") > 0
|
||||
assert total(report, "filled") + total(report, "still_missing") == total(report, "missing")
|
||||
|
||||
|
||||
def test_smoke_failure_stops_before_any_fill_and_redacts_provider_error(tmp_path: Path) -> None:
|
||||
store, manifest, _ = prepared(tmp_path)
|
||||
|
||||
def fail(*args: Any) -> str:
|
||||
raise RuntimeError("provider unavailable SECRET_DO_NOT_PRINT")
|
||||
|
||||
chat = ScriptedChat(fail)
|
||||
before = snapshot(tmp_path)
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=1.0)
|
||||
assert report["status"] == "failed" and report["client_calls"] == len(chat.calls) == 1
|
||||
assert total(report, "filled") == 0 and total(report, "still_missing") == total(report, "missing")
|
||||
assert snapshot(tmp_path) == before
|
||||
assert "SECRET_DO_NOT_PRINT" not in json.dumps(report)
|
||||
|
||||
|
||||
def test_pending_units_refuse_before_calls_or_writes(tmp_path: Path) -> None:
|
||||
store, original, _ = prepared(tmp_path)
|
||||
path = write_rule_manifest(tmp_path, original.programs["injection"], "exp-902-pending")
|
||||
manifest = ex.load_manifest(path)
|
||||
ex.register(store, manifest, commit_of=lambda _: "c" * 40)
|
||||
chat = ScriptedChat()
|
||||
before = snapshot(tmp_path)
|
||||
with pytest.raises(ex.ExperimentError, match="pending|frozen|units"):
|
||||
fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=1.0)
|
||||
assert not chat.calls and snapshot(tmp_path) == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("budget", [None, -1.0, float("nan"), float("inf")])
|
||||
def test_paid_fill_requires_a_finite_nonnegative_budget(tmp_path: Path, budget: float | None) -> None:
|
||||
store, manifest, _ = prepared(tmp_path)
|
||||
chat = ScriptedChat()
|
||||
before = snapshot(tmp_path)
|
||||
with pytest.raises(ex.ExperimentError, match="budget-usd"):
|
||||
fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=budget)
|
||||
assert not chat.calls and snapshot(tmp_path) == before
|
||||
|
||||
|
||||
def test_zero_budget_stops_before_smoke(tmp_path: Path) -> None:
|
||||
store, manifest, _ = prepared(tmp_path)
|
||||
chat = ScriptedChat()
|
||||
before = snapshot(tmp_path)
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=0.0)
|
||||
assert report["status"] == "unfinished" and report["client_calls"] == 0 and report["usd"] == 0
|
||||
assert not chat.calls and snapshot(tmp_path) == before
|
||||
|
||||
|
||||
def test_parse_retry_requests_are_cached_for_identical_replay(tmp_path: Path) -> None:
|
||||
store, manifest, units = prepared(tmp_path)
|
||||
answers = 0
|
||||
|
||||
def first_invalid(messages: Any, temperature: Any) -> str:
|
||||
nonlocal answers
|
||||
answers += 1
|
||||
return "invalid json" if answers == 1 else default_answer(messages, temperature)
|
||||
|
||||
chat = ScriptedChat(first_invalid)
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, client=chat, budget_usd=1.0)
|
||||
assert report["status"] == "done" and total(report, "still_missing") == 0
|
||||
assert len(chat.calls) == 2 * len(units) + 1
|
||||
caller, meter = ex.zero_cost_caller("deepseek-flash", PRICES, store)
|
||||
replay = ex.run_rule(store, manifest, EXAMPLES, caller, meter, excerpt_text)
|
||||
assert replay.decided == {"score": len(units)} and replay.skipped["replay_miss"] == 0
|
||||
assert meter.calls == 0
|
||||
|
||||
|
||||
def test_cli_dry_run_needs_no_endpoint_and_does_not_write_out(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
import openultrasast.learn.examples as examples
|
||||
import openultrasast.model.endpoint as endpoint
|
||||
|
||||
_, manifest, _ = prepared(tmp_path)
|
||||
real_get_blob = FileStore.get_blob
|
||||
|
||||
def get_blob(self: FileStore, kind: str, sha: str, **kwargs: Any) -> bytes | None:
|
||||
if kind == "excerpts":
|
||||
text = CODE.get(sha)
|
||||
return text.encode() if text is not None else None
|
||||
return real_get_blob(self, kind, sha, **kwargs)
|
||||
|
||||
monkeypatch.setattr(FileStore, "get_blob", get_blob)
|
||||
monkeypatch.setattr(examples, "load_examples", lambda *_: EXAMPLES)
|
||||
monkeypatch.setattr(endpoint, "resolve_chat_endpoint", lambda *_: pytest.fail("dry run resolved an endpoint"))
|
||||
out = tmp_path / "report.json"
|
||||
before = snapshot(tmp_path)
|
||||
code = main(
|
||||
["learn", "experiment", "fill", str(manifest.path), "--dry-run", "--memory", f"file://{tmp_path / 'memory'}", "--out", str(out)]
|
||||
)
|
||||
assert code == 0 and not out.exists() and snapshot(tmp_path) == before
|
||||
report = json.loads(capsys.readouterr().out)
|
||||
assert report["client_calls"] == 0 and report["status"] == "dry_run"
|
||||
|
||||
|
||||
def test_growing_call_costs_stay_inside_the_initial_reserve(tmp_path: Path) -> None:
|
||||
store, manifest, _ = prepared(tmp_path)
|
||||
|
||||
class GrowingCostChat(ScriptedChat):
|
||||
def complete(self, **kwargs: Any) -> Any:
|
||||
response = super().complete(**kwargs)
|
||||
self.usage[-1] = {"prompt_tokens": 0, "completion_tokens": min(len(self.calls), 9) * 1_000}
|
||||
return response
|
||||
|
||||
chat = GrowingCostChat()
|
||||
prices = {"cache_hit_per_m": 0.0, "input_per_m": 0.0, "output_per_m": 1.0}
|
||||
report = fill_rule(store, manifest, EXAMPLES, prices, excerpt_text, client=chat, budget_usd=0.025)
|
||||
assert report["status"] == "unfinished" and 0 < report["usd"] <= 0.025
|
||||
assert report["usd"] == pytest.approx(sum(int(row["completion_tokens"]) for row in chat.usage) / 1_000_000)
|
||||
assert total(report, "filled") == len(chat.calls) and total(report, "still_missing") > 0
|
||||
|
||||
|
||||
def test_dry_run_prices_prompt_tokens_and_family_cached_output_mean(tmp_path: Path) -> None:
|
||||
from openultrasast.learn.fill import DEFAULT_OUTPUT_TOKENS
|
||||
|
||||
store, manifest, units = prepared(tmp_path)
|
||||
seed_one(store, manifest, units[0])
|
||||
report = fill_rule(store, manifest, EXAMPLES, PRICES, excerpt_text, dry_run=True)
|
||||
folds = {f.name: f for f in ex.program_folds(EXAMPLES, manifest)[0]}
|
||||
index = {e.id: e for e in EXAMPLES}
|
||||
expected = dict.fromkeys(manifest.families, 0.0)
|
||||
for unit in units[1:]:
|
||||
example = index[unit.unit]
|
||||
program = Program(ex.arm_spec(manifest, manifest.arms["A"], unit.family, store), EXAMPLES, excerpt_text)
|
||||
prompt, _ = program.prepare(Candidate(Target.of(example), CODE[example.excerpt_sha], "python"), folds[unit.fold], seed=0)
|
||||
output_tokens = 60 if unit.family == units[0].family else DEFAULT_OUTPUT_TOKENS
|
||||
expected[unit.family] += (
|
||||
program.spec.k * (prompt.estimated_tokens * PRICES["input_per_m"] + output_tokens * PRICES["output_per_m"]) / 1_000_000
|
||||
)
|
||||
for family, amount in expected.items():
|
||||
assert report["families"][family]["estimated_usd"] == pytest.approx(amount)
|
||||
|
||||
|
||||
def test_bounded_provider_caps_output_and_disables_unmetered_retries(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from openultrasast.model.endpoint import DeepSeekChatClient
|
||||
from openultrasast.provider.openrouter import OpenRouterChatClient
|
||||
|
||||
raw_calls: list[dict[str, Any]] = []
|
||||
|
||||
def raw(self: OpenRouterChatClient, **kwargs: Any) -> dict[str, object]:
|
||||
raw_calls.append(kwargs)
|
||||
assert self.max_attempts == 1
|
||||
return {"choices": [{"message": {"role": "assistant", "content": ""}}], "usage": {"prompt_tokens": 20, "completion_tokens": 3}}
|
||||
|
||||
monkeypatch.setattr(OpenRouterChatClient, "complete_chat_raw", raw)
|
||||
transport = OpenRouterChatClient(api_key="test-placeholder", max_attempts=3)
|
||||
original = DeepSeekChatClient(transport, disable_thinking=True) # the default: config.models.thinking is False
|
||||
bounded = original.bounded(300)
|
||||
response = bounded.complete(model="deepseek-flash", messages=[{"role": "user", "content": "JSON"}], tools=[], json_object=True)
|
||||
assert response.content == "" and len(raw_calls) == 1
|
||||
assert raw_calls[0]["extra_body"]["max_tokens"] == 300
|
||||
# The bounded copy keeps the configured client's thinking setting, so a fill answers as the run would.
|
||||
assert raw_calls[0]["extra_body"]["thinking"] == {"type": "disabled"}
|
||||
thinking_on = DeepSeekChatClient(transport, disable_thinking=False).bounded(300)
|
||||
thinking_on.complete(model="deepseek-flash", messages=[{"role": "user", "content": "JSON"}], tools=[], json_object=True)
|
||||
assert "thinking" not in raw_calls[-1]["extra_body"]
|
||||
assert bounded.usage == [{"prompt_tokens": 20, "completion_tokens": 3}]
|
||||
assert original.usage == [] and transport.max_attempts == 3
|
||||
Reference in New Issue
Block a user