mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 11:49:28 +00:00
Marius CiepluchandClaude Opus 5.5
dd7ebc119b
kind: rule experiments score once from the response cache behind a zero-ceiling meter (a miss is excluded and counted, never asked) and apply each arm's BLOCK rule to the same held-out scores. Arm B chooses a per-family raw-score threshold on the other outer folds only (lowest t with Wilson 95% LB >= 0.95); a test fails if a fold's own labels move its threshold. Manifest and frozen units (353 units, 146 groups, 244 paired) for exp-003. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
395 lines
22 KiB
Python
395 lines
22 KiB
Python
"""``ousast learn ...`` beyond ``labels``: the local memory and the compiled program (design sections 4-5).
|
|
|
|
Host only, maintainer surface. The parser is registered from :mod:`..cli`; every heavy import is lazy. No command
|
|
here spends money unless it is given a real client: ``compile``/``evaluate``/``curve`` and ``memory embed`` run under
|
|
an explicit ``--budget-usd`` ceiling through the metered clients.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def add_commands(learn_sub: Any) -> None:
|
|
memory = learn_sub.add_parser("memory", help="the decision engine's local memory of labelled examples")
|
|
memory_sub = memory.add_subparsers(dest="memory_command", required=True)
|
|
build = memory_sub.add_parser("build", help="label + feature record + excerpt -> example rows (no model)")
|
|
build.add_argument("--labels", type=Path, required=True, help="labels-<sha>.jsonl written by `ousast learn labels --out`")
|
|
build.add_argument("--features", type=Path, action="append", default=[], help="feature records (JSONL with repo, pin); repeatable")
|
|
build.add_argument("--sources", type=Path, help="the source list (default: the packaged learn/sources.toml)")
|
|
build.add_argument("--root", type=Path, default=Path("."), help="the repository root the source paths are relative to")
|
|
build.add_argument("--cache", type=Path, help="clones and checkouts (default ~/.cache/openultrasast)")
|
|
build.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
|
build.add_argument("--profile", choices=("static", "plane"), default="static")
|
|
build.add_argument("--units", type=Path, help="the harvest units.json: the pin each pair side's record is stored under")
|
|
build.add_argument("--dry-run", action="store_true", help="count, write nothing")
|
|
embed = memory_sub.add_parser("embed", help="embed the stored excerpts (OpenRouter, metered; cached by excerpt sha)")
|
|
embed.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
|
embed.add_argument("--model", default="openai/text-embedding-3-small")
|
|
embed.add_argument("--manifest", type=Path, default=Path("plane/models/openrouter-embedding.yaml"), help="the Model with the prices")
|
|
embed.add_argument("--budget-usd", type=float, help="the ceiling; required unless --dry-run")
|
|
embed.add_argument("--budget-calls", type=int)
|
|
embed.add_argument("--batch", type=int, default=32)
|
|
embed.add_argument("--dry-run", action="store_true", help="count and estimate the tokens; no call")
|
|
for name, text in (
|
|
("compile", "compile the program on the compile split (bootstrap few-shot, instruction search)"),
|
|
("evaluate", "evaluate a compiled program on the outer folds: cross-fitted calibration, nested operating points"),
|
|
("curve", "the program's metric over memory size on a fixed candidate subset"),
|
|
):
|
|
sub = learn_sub.add_parser(name, help=text)
|
|
sub.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
|
sub.add_argument("--profile", choices=("static", "plane"), default="static")
|
|
sub.add_argument("--spec", type=Path, help="the pre-registered compile spec (default: the packaged learn/compile.toml)")
|
|
sub.add_argument("--manifest", type=Path, default=Path("plane/models/deepseek-flash.yaml"), help="the chat Model with the prices")
|
|
sub.add_argument("--budget-usd", type=float, help="the ceiling of a paid run; required unless --replay-only")
|
|
sub.add_argument("--budget-calls", type=int)
|
|
sub.add_argument("--replay-only", action="store_true", help="no client: every response must come from the cache")
|
|
sub.add_argument("--out", type=Path, help="write the report (counts and intervals only) here")
|
|
if name != "compile":
|
|
sub.add_argument("--program", required=True, help="the compiled program id (programs/<id>.json in the store)")
|
|
if name == "curve":
|
|
sub.add_argument("--subset", type=int, default=200)
|
|
|
|
experiment = learn_sub.add_parser("experiment", help="A/B experiments: register a manifest, run its arms paired, analyse")
|
|
experiment_sub = experiment.add_subparsers(dest="experiment_command", required=True)
|
|
for name, text in (
|
|
("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"),
|
|
("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"):
|
|
sub.add_argument("--profile", choices=("static", "plane"), default="static")
|
|
if name == "run":
|
|
sub.add_argument(
|
|
"--manifest-model", type=Path, default=Path("plane/models/deepseek-flash.yaml"), help="the chat Model (prices)"
|
|
)
|
|
sub.add_argument("--no-early-stop", action="store_true", help="take no first look during the 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>/)")
|
|
if name == "analyse":
|
|
sub.add_argument("--record", action="store_true", help="also write the experiment_result row")
|
|
|
|
leaks = learn_sub.add_parser("audit-leaks", help="how well each signal alone separates the sides of a pair (no model)")
|
|
leaks.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
|
leaks.add_argument("--units", type=Path, help="the harvest units.json (unit -> case and side) for feature rows without a side")
|
|
leaks.add_argument("--profile", choices=("static", "plane"), default="plane")
|
|
leaks.add_argument("--kind", default="features", help="the row kind to audit (features, or example)")
|
|
leaks.add_argument("--threshold", type=float, help="flag a separation at least this (default 0.20)")
|
|
leaks.add_argument("--min-pairs", type=int, help="over at least this many pairs (default 10)")
|
|
leaks.add_argument("--out", type=Path, help="write the counts-only report here")
|
|
|
|
|
|
def run(args: argparse.Namespace) -> int:
|
|
if args.learn_command == "audit-leaks":
|
|
return _audit_leaks(args)
|
|
if args.learn_command == "memory" and args.memory_command == "build":
|
|
return _memory_build(args)
|
|
if args.learn_command == "memory" and args.memory_command == "embed":
|
|
return _memory_embed(args)
|
|
if args.learn_command in ("compile", "evaluate", "curve"):
|
|
return _program_command(args)
|
|
if args.learn_command == "experiment":
|
|
return _experiment(args)
|
|
print(f"learn: unknown command {args.learn_command}", file=sys.stderr)
|
|
return 2
|
|
|
|
|
|
def _experiment(args: argparse.Namespace) -> int:
|
|
"""register | units | run | analyse. Exit 2 on a refused manifest (unregistered, modified, uncommitted), 3 when a
|
|
ceiling stops the run (``unfinished``: a rerun resumes from the outcome rows)."""
|
|
from datetime import date
|
|
|
|
from ..plane.budget import MeteredClient
|
|
from ..plane.memory import open_store
|
|
from . import experiments as ex
|
|
from .program import Caller
|
|
|
|
store = open_store(args.memory)
|
|
try:
|
|
manifest = ex.load_manifest(args.manifest)
|
|
if args.experiment_command == "register":
|
|
row = ex.register(store, manifest, created=date.today().isoformat())
|
|
shown = ("experiment", "manifest_digest", "units_status", "units_digest", "pin", "arms")
|
|
print(json.dumps({k: row[k] for k in shown}, indent=2))
|
|
return 0
|
|
if args.experiment_command == "analyse":
|
|
row = ex.check_registered(store, manifest, need_units=manifest.kind == "rule")
|
|
provenance = {"code_commit": ex.committed_at_head(manifest.path)}
|
|
outcomes = ex.load_outcomes(store, manifest.id)
|
|
if manifest.kind == "rule":
|
|
units = ex.read_units(manifest.units_file)
|
|
report = ex.analyse_rule(manifest, outcomes, units, resamples=args.resamples, provenance=provenance)
|
|
else:
|
|
report = ex.analyse(manifest, outcomes, resamples=args.resamples, provenance=provenance)
|
|
if args.record:
|
|
result = ex.rule_result_row if manifest.kind == "rule" else ex.result_row
|
|
store.put_row(result(manifest, str(row["pin"]), report))
|
|
_write(report, args.out or (ex.RESULT_DIR / manifest.id / "result.json" if args.record else None))
|
|
return 0
|
|
from .embeddings import EmbeddingCache
|
|
from .examples import load_examples
|
|
|
|
memory = load_examples(store, args.profile)
|
|
if args.experiment_command == "units":
|
|
units = ex.freeze_units(store, manifest, memory)
|
|
digest = ex.write_units(manifest.units_file, units)
|
|
paired = sum(1 for u in units if u.pair)
|
|
print(json.dumps({"units": len(units), "paired": paired, "groups": len({u.group for u in units}), "sha256": digest,
|
|
"path": str(manifest.units_file), "next": "commit the file, then register again"}, indent=2)) # fmt: skip
|
|
return 0
|
|
cache = EmbeddingCache(store)
|
|
vectors = {sha: v for sha in cache.names() if (v := cache.get(sha)) is not None} or None
|
|
parameters = _model_parameters(args.manifest_model)
|
|
|
|
def excerpt_text(sha: str) -> str | None:
|
|
data = store.get_blob("excerpts", sha)
|
|
return data.decode("utf-8") if data is not None else None
|
|
|
|
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)
|
|
_write(summary.as_dict(), args.out)
|
|
return 0
|
|
callers: dict[str, Caller] = {}
|
|
for name, arm in manifest.arms.items():
|
|
if arm.replay_only:
|
|
callers[name] = Caller(None, _arm_model(manifest, arm, store), parameters, store=store)
|
|
continue
|
|
from ..config import load_config
|
|
from ..model.endpoint import resolve_chat_endpoint
|
|
|
|
resolved = resolve_chat_endpoint(load_config(None))
|
|
if resolved is None:
|
|
raise SystemExit("learn experiment run: no chat endpoint configured (DEEPSEEK_API_KEY); nothing was spent")
|
|
client = MeteredClient(resolved[0], prices=parameters, budget_usd=arm.budget_usd)
|
|
callers[name] = Caller(client, _arm_model(manifest, arm, store), parameters, store=store)
|
|
|
|
early = not args.no_early_stop
|
|
summary = ex.run(store, manifest, memory, callers, excerpt_text, vectors, stop_early=early, resamples=args.resamples)
|
|
except ex.ExperimentError as exc:
|
|
print(f"learn experiment {args.experiment_command}: {exc}", file=sys.stderr)
|
|
return 2
|
|
_write(summary.as_dict(), args.out)
|
|
if summary.status == "unfinished":
|
|
print(f"learn experiment run: unfinished -- {summary.reason}", file=sys.stderr)
|
|
return 3
|
|
return 0
|
|
|
|
|
|
def _arm_model(manifest: Any, arm: Any, store: Any) -> str:
|
|
from .compile import load_program
|
|
|
|
if "model" in arm.spec:
|
|
return str(arm.spec["model"])
|
|
pid = next(iter({**manifest.programs, **arm.programs}.values()), None)
|
|
return str(load_program(store, pid)["classify"]["model"]) if pid else "deepseek-flash"
|
|
|
|
|
|
def _write(report: dict[str, Any], out: Path | None) -> None:
|
|
text = json.dumps(report, indent=2, sort_keys=True, default=str) + "\n"
|
|
if out is not None:
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_text(text, encoding="utf-8")
|
|
print(text, end="")
|
|
|
|
|
|
def _caller(args: argparse.Namespace, store: Any, model: str) -> Any:
|
|
"""The metered, cached caller; ``None`` client under --replay-only. Exits 2 without a ceiling or a key."""
|
|
from ..plane.budget import MeteredClient
|
|
from .program import Caller
|
|
|
|
parameters = _model_parameters(args.manifest)
|
|
if args.replay_only:
|
|
return Caller(None, model, parameters, store=store)
|
|
if args.budget_usd is None:
|
|
raise SystemExit("learn: --budget-usd is required (a paid run starts only under a ceiling), or --replay-only")
|
|
from ..config import load_config
|
|
from ..model.endpoint import resolve_chat_endpoint
|
|
|
|
resolved = resolve_chat_endpoint(load_config(None))
|
|
if resolved is None:
|
|
raise SystemExit("learn: no chat endpoint configured (DEEPSEEK_API_KEY); nothing was spent")
|
|
client = MeteredClient(resolved[0], prices=parameters, budget_usd=args.budget_usd, budget_calls=args.budget_calls)
|
|
return Caller(client, model, parameters, store=store)
|
|
|
|
|
|
def _identities(store: Any) -> dict[str, tuple[str, ...]]:
|
|
"""Per example id, the strings its demonstration's rationale must not name: group, repository and path parts."""
|
|
from .compile import group_identities
|
|
from .examples import EXAMPLE_KIND
|
|
|
|
out: dict[str, tuple[str, ...]] = {}
|
|
for record in store.rows(kind=EXAMPLE_KIND):
|
|
row = record.row
|
|
path = str(row.get("candidate", "")).partition("::")[0]
|
|
out[str(row["id"])] = (*group_identities(str(row.get("group") or "")), str(row["repo"]), path, path.rsplit("/", 1)[-1])
|
|
return out
|
|
|
|
|
|
def _program_command(args: argparse.Namespace) -> int:
|
|
"""compile | evaluate | curve. Exit 3 when the ceiling stops the run (``unfinished``, no artifact, no number)."""
|
|
from datetime import date
|
|
|
|
from ..plane.budget import BudgetExhausted
|
|
from ..plane.memory import open_store
|
|
from .compile import DEFAULT_SPEC, compile_program, load_program, load_spec, spec_of
|
|
from .embeddings import EmbeddingCache
|
|
from .evaluate import candidate_from_example, evaluate, learning_curve, predict, targets, unevaluable_families
|
|
from .examples import load_examples
|
|
from .folds import compile_split, outer_folds
|
|
from .program import Program
|
|
|
|
store = open_store(args.memory)
|
|
spec = load_spec(args.spec or DEFAULT_SPEC)
|
|
memory = load_examples(store, args.profile)
|
|
cache = EmbeddingCache(store)
|
|
vectors = {sha: v for sha in cache.names() if (v := cache.get(sha)) is not None} or None
|
|
|
|
def excerpt_text(sha: str) -> str | None:
|
|
data = store.get_blob("excerpts", sha)
|
|
return data.decode("utf-8") if data is not None else None
|
|
|
|
model = str(spec["classify"].get("model", "deepseek-flash"))
|
|
caller = _caller(args, store, model)
|
|
report: dict[str, Any]
|
|
try:
|
|
if args.learn_command == "compile":
|
|
ids = _identities(store)
|
|
result = compile_program(
|
|
memory, caller, excerpt_text, spec=spec, profile=args.profile, vectors=vectors,
|
|
identities=lambda e: ids.get(e.id, ()), store=store, created=date.today().isoformat(),
|
|
) # fmt: skip
|
|
report = {"status": result.status, "program_id": result.program_id, "calls": result.calls, "replayed": result.replayed,
|
|
"reason": result.reason, "bootstrap": result.bootstrap} # fmt: skip
|
|
else:
|
|
artifact = load_program(store, args.program)
|
|
program = Program(spec_of(artifact), memory, excerpt_text, vectors)
|
|
split = compile_split(memory, fraction=float(spec["split"].get("compile_fraction", 0.25)), seed=spec.seed)
|
|
folds = outer_folds(memory, seed=spec.seed, exclude=split.groups)
|
|
pairs = targets(folds, memory, skip=unevaluable_families(memory))
|
|
candidate_of = candidate_from_example(excerpt_text, vectors)
|
|
if args.learn_command == "evaluate":
|
|
report = evaluate(predict(program, folds, caller, candidate_of, seed=spec.seed, pairs=pairs), seed=spec.seed)
|
|
else:
|
|
report = {"curve": learning_curve(program, pairs[: args.subset], caller, candidate_of, seed=spec.seed)}
|
|
report["program_id"] = args.program
|
|
except BudgetExhausted as exc:
|
|
print(f"learn {args.learn_command}: unfinished -- {exc}", file=sys.stderr)
|
|
return 3
|
|
report["spend"] = {"usd": caller.usd(), "client_calls": caller.calls, "replayed": caller.replayed}
|
|
text = json.dumps(report, indent=2, sort_keys=True, default=str) + "\n"
|
|
if args.out is not None:
|
|
args.out.parent.mkdir(parents=True, exist_ok=True)
|
|
args.out.write_text(text, encoding="utf-8")
|
|
print(text, end="")
|
|
return 0
|
|
|
|
|
|
def _model_parameters(path: Path) -> dict[str, Any]:
|
|
from ..plane.manifests import load_manifests
|
|
|
|
models = load_manifests([path]).models
|
|
if len(models) != 1:
|
|
raise SystemExit(f"{path}: expected one Model, found {sorted(models)}")
|
|
return dict(next(iter(models.values())).parameters)
|
|
|
|
|
|
def _memory_embed(args: argparse.Namespace) -> int:
|
|
"""Exit 2 without a ceiling, 3 when the ceiling stops the run (``unfinished``; a rerun resumes from the cache)."""
|
|
from ..plane.budget import BudgetExhausted, MeteredEmbeddingClient
|
|
from ..plane.memory import open_store
|
|
from .embeddings import embed_excerpts
|
|
from .examples import EXAMPLE_KIND
|
|
|
|
store = open_store(args.memory)
|
|
shas = sorted({str(r.row["excerpt_sha"]) for r in store.rows(kind=EXAMPLE_KIND)})
|
|
client = None
|
|
if not args.dry_run:
|
|
if args.budget_usd is None:
|
|
print("learn memory embed: --budget-usd is required (a paid run starts only under a ceiling)", file=sys.stderr)
|
|
return 2
|
|
from ..provider.openrouter import OpenRouterEmbeddingClient
|
|
|
|
client = MeteredEmbeddingClient(
|
|
OpenRouterEmbeddingClient.from_env(), prices=_model_parameters(args.manifest), budget_usd=args.budget_usd,
|
|
budget_calls=args.budget_calls,
|
|
) # fmt: skip
|
|
try:
|
|
report = embed_excerpts(store, shas, client, model=args.model, batch=args.batch)
|
|
except BudgetExhausted as exc:
|
|
print(f"learn memory embed: unfinished -- {exc}", file=sys.stderr)
|
|
return 3
|
|
print(json.dumps(report.as_dict(), indent=2, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
def _memory_build(args: argparse.Namespace) -> int:
|
|
"""Exit 2 on a refused label (``origin: user``, a source outside the list) or a source with >10% unread excerpts."""
|
|
from ..plane.memory import open_store
|
|
from .examples import ExampleBuildError, SideReader, build_examples, pair_pins, read_jsonl, records_by_key
|
|
from .labels import DEFAULT_CACHE, DEFAULT_SOURCES, load_sources
|
|
|
|
sources = load_sources(args.sources or DEFAULT_SOURCES)
|
|
labels = read_jsonl(args.labels)
|
|
store = None if args.dry_run else open_store(args.memory)
|
|
rows = [row for path in args.features for row in read_jsonl(path) if row.get("profile", args.profile) == args.profile]
|
|
if store is not None:
|
|
# one profile only: the static and plane records of a candidate share a key, and the other would overwrite it
|
|
rows += [r.row for r in store.rows(kind="features") if r.row.get("profile") == args.profile]
|
|
reader = SideReader(args.root, args.cache or DEFAULT_CACHE, sources, labels)
|
|
try:
|
|
build = build_examples(
|
|
labels, records_by_key(rows), reader, sources=sources, store=store, profile=args.profile, license_of=reader.license,
|
|
pins=pair_pins(json.loads(args.units.read_text(encoding="utf-8"))) if args.units is not None else None,
|
|
) # fmt: skip
|
|
except ExampleBuildError as exc:
|
|
print(f"learn memory build: {exc}", file=sys.stderr)
|
|
return 2
|
|
summary = build.summary()
|
|
print(json.dumps(summary, indent=2, sort_keys=True))
|
|
if summary["failing_sources"]:
|
|
print(
|
|
f"learn memory build: more than 10% of the recorded candidates have no excerpt in {summary['failing_sources']}", file=sys.stderr
|
|
)
|
|
return 2
|
|
return 0
|
|
|
|
|
|
def _audit_leaks(args: argparse.Namespace) -> int:
|
|
"""Exit 2 when the store holds no pair to audit (an unread store is not a clean one)."""
|
|
from ..plane.memory import open_store
|
|
from .leaks import MIN_PAIRS, THRESHOLD, audit, pair_rows, sides_from_units
|
|
|
|
store = open_store(args.memory)
|
|
rows = [r.row for r in store.rows(kind=args.kind)]
|
|
units = json.loads(args.units.read_text(encoding="utf-8")) if args.units is not None else {}
|
|
pairs, skipped = pair_rows(rows, sides_from_units(units), args.profile)
|
|
if not pairs:
|
|
print(f"learn audit-leaks: {len(rows)} {args.kind} rows in {store.describe()}, no pair to audit {skipped}", file=sys.stderr)
|
|
return 2
|
|
threshold = THRESHOLD if args.threshold is None else args.threshold
|
|
report = audit(pairs, args.profile, threshold=threshold, min_pairs=MIN_PAIRS if args.min_pairs is None else args.min_pairs)
|
|
report = {"rows_read": len(rows), "kind": args.kind, "skipped": skipped, **report}
|
|
text = json.dumps(report, indent=1, sort_keys=True) + "\n"
|
|
if args.out is not None:
|
|
args.out.parent.mkdir(parents=True, exist_ok=True)
|
|
args.out.write_text(text, encoding="utf-8")
|
|
print(json.dumps({"rows_read": len(rows), "pairs": report["pairs"], "flagged_signals": report["flagged_signals"]}, sort_keys=True))
|
|
for cell in report["flagged"]:
|
|
print(
|
|
f" {cell['separation']:.3f} {cell['signal']} [{cell['scope']}] {cell['correct']}-{cell['wrong']} of {cell['pairs']} pairs"
|
|
)
|
|
return 0
|
|
|
|
|
|
__all__ = ["add_commands", "run"]
|