mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
learn: exp-005 graph slice -- slice in the prompt (opt-in), program-arm fill, manifest and units
Implemented with Codex. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
1b2c7993a7
commit
0acbc7ec41
@@ -599,3 +599,61 @@ def test_experiment_rows_share_one_json_type_per_column(tmp_path: Path) -> None:
|
||||
if value is not None:
|
||||
shapes.setdefault(name, set()).add(_json_shape(value))
|
||||
assert {name: sorted(s) for name, s in shapes.items() if len(s) > 1} == {}
|
||||
|
||||
|
||||
def test_graph_slice_manifest_freezes_incumbents_units_and_single_prompt_override() -> None:
|
||||
manifest = ex.load_manifest(ROOT / "plane/experiments/exp-005-graph-slice.yaml")
|
||||
incumbent = ex.load_manifest(ROOT / "plane/experiments/exp-004-pooled-block-gate.yaml")
|
||||
assert manifest.kind == "program" and manifest.programs == incumbent.programs
|
||||
assert manifest.families == incumbent.families and manifest.folds == incumbent.folds
|
||||
assert manifest.arms["A"].spec == {} and manifest.arms["A"].replay_only
|
||||
assert manifest.arms["B"].spec == {"slice": True} and manifest.paid_arms == ["B"]
|
||||
assert manifest.budget_total_usd == manifest.arms["B"].budget_usd == 6.0
|
||||
assert manifest.units_file.read_bytes() == incumbent.units_file.read_bytes()
|
||||
assert len(ex.read_units(manifest.units_file)) == 501
|
||||
assert [m.name for m in manifest.primary] == ["within_pair_auc"]
|
||||
assert [m.name for m in manifest.secondary] == ["top_30_precision", "unsure_rate", "usd_per_candidate"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ci,cost_a,cost_b,expected",
|
||||
[
|
||||
((0.01, 0.2), 1.0, 1.5, "adopt"),
|
||||
((0.01, 0.2), 1.0, 1.5001, "inconclusive"),
|
||||
((0.01, 0.2), None, None, "inconclusive"),
|
||||
((-0.2, -0.01), 1.0, 2.0, "reject"),
|
||||
((0.0, 0.2), 1.0, 1.0, "inconclusive"),
|
||||
((-0.2, 0.1), 1.0, 1.0, "inconclusive"),
|
||||
],
|
||||
)
|
||||
def test_graph_slice_adoption_cost_guard(ci: tuple[float, float], cost_a: float | None, cost_b: float | None, expected: str) -> None:
|
||||
manifest = ex.load_manifest(ROOT / "plane/experiments/exp-005-graph-slice.yaml")
|
||||
est = ex.Estimate("within_pair_auc", 0.5, 0.6, 0.1, ci, 0.1, 100, 10, 100)
|
||||
cost = ex.Estimate("usd_per_candidate", cost_a, cost_b, None, None, None, 100, 10, 100)
|
||||
assert ex.decide(manifest, {"within_pair_auc": est}, cost, None)[0] == expected
|
||||
|
||||
|
||||
def test_graph_slice_secondary_metrics_use_pooled_scores_and_unsure() -> None:
|
||||
rows = [
|
||||
ex.ArmUnit(f"u{i:02}", "g", "injection", int(i < 15), "", float(40 - i), "unsure" if i < 4 else "not_vulnerable", 0.01, None)
|
||||
for i in range(40)
|
||||
]
|
||||
assert ex.metric_value(ex.Metric("top_30_precision"), rows) == 0.5
|
||||
assert ex.metric_value(ex.Metric("top_30_precision"), rows[:29]) is None
|
||||
assert ex.metric_value(ex.Metric("unsure_rate"), rows) == 0.1
|
||||
assert ex.binary_value(ex.Metric("unsure_rate"), rows[0]) is True
|
||||
|
||||
|
||||
def test_pair_bootstrap_counts_repeated_repository_draws(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
rows = crafted(groups=2, b_better=True)
|
||||
pairs = ex.paired_units(rows)
|
||||
original = ex.metric_value
|
||||
seen: list[int] = []
|
||||
|
||||
def measured(metric: ex.Metric, units: list[ex.ArmUnit]) -> float | None:
|
||||
seen.append(len({(u.group, u.family, u.pair) for u in units}))
|
||||
return original(metric, units)
|
||||
|
||||
monkeypatch.setattr(ex, "metric_value", measured)
|
||||
ex.paired_bootstrap(pairs, ex.Metric("within_pair_auc"), resamples=20)
|
||||
assert len(seen) == 42 and set(seen) == {4}
|
||||
|
||||
@@ -289,3 +289,86 @@ def test_bounded_provider_caps_output_and_disables_unmetered_retries(monkeypatch
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sampling", ["temperature", "retrieval_ensemble"])
|
||||
def test_program_fill_selected_arm_replicates_retries_and_replay(tmp_path: Path, sampling: str) -> None:
|
||||
from test_learn_experiments import artifact, committed, example_rows, put_program, write_manifest
|
||||
|
||||
from openultrasast.learn.fill import fill_experiment
|
||||
from openultrasast.learn.program import Caller
|
||||
|
||||
store = FileStore(tmp_path / "memory")
|
||||
store.put_rows(example_rows())
|
||||
path = write_manifest(tmp_path, put_program(store, artifact()), repeats=2)
|
||||
path.write_text(path.read_text().replace("sampling: retrieval_ensemble", f"sampling: {sampling}, slice: true"))
|
||||
manifest = ex.load_manifest(path)
|
||||
units = ex.freeze_units(store, manifest, EXAMPLES)[:4]
|
||||
ex.write_units(manifest.units_file, units)
|
||||
ex.register(store, manifest, commit_of=committed)
|
||||
before = snapshot(tmp_path)
|
||||
dry = fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, dry_run=True)
|
||||
assert dry["arms"] == ["B"] and total(dry, "missing") == 4 * len(units)
|
||||
assert snapshot(tmp_path) == before
|
||||
zero = fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, budget_usd=0, client=ScriptedChat())
|
||||
assert zero["status"] == "unfinished" and zero["client_calls"] == 0 and snapshot(tmp_path) == before
|
||||
with pytest.raises(ex.ExperimentError, match="manifest ceiling"):
|
||||
fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, budget_usd=2.01, client=ScriptedChat())
|
||||
calls = 0
|
||||
|
||||
def retry_once(messages: Any, temperature: Any) -> str:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return "invalid json" if calls == 1 else default_answer(messages, temperature)
|
||||
|
||||
chat = ScriptedChat(retry_once)
|
||||
filled = fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, budget_usd=1, client=chat)
|
||||
assert filled["status"] == "done" and total(filled, "still_missing") == 0
|
||||
assert filled["client_calls"] == 4 * len(units) + 1
|
||||
assert all("Data flow inside the function:" in str(c["messages"]) for c in chat.calls)
|
||||
assert not ex.load_outcomes(store, manifest.id)
|
||||
arm_a = fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, arms=["A"], dry_run=True)
|
||||
assert arm_a["arms"] == ["A"] and total(arm_a, "cached") == 0
|
||||
fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, arms=["A"], budget_usd=1, client=ScriptedChat())
|
||||
callers = {name: Caller(None, "deepseek-flash", PRICES, store=store) for name in manifest.arms}
|
||||
summary = ex.run(store, manifest, EXAMPLES, callers, excerpt_text, stop_early=False)
|
||||
assert summary.status == "done" and summary.decided == {"A": 2 * len(units), "B": 2 * len(units)}
|
||||
assert all(c.calls == 0 for c in callers.values())
|
||||
before = snapshot(tmp_path)
|
||||
with pytest.raises(ex.ExperimentError, match="known arm"):
|
||||
fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, arms=["unknown"], dry_run=True)
|
||||
assert snapshot(tmp_path) == before
|
||||
|
||||
|
||||
def test_program_cli_selects_fill_arm_and_runs_replay_without_endpoint(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
from test_learn_experiments import prepared as prepared_program
|
||||
|
||||
import openultrasast.learn.examples as examples
|
||||
import openultrasast.model.endpoint as endpoint
|
||||
from openultrasast.learn.fill import fill_experiment
|
||||
|
||||
store, manifest, units = prepared_program(tmp_path)
|
||||
real_get_blob = FileStore.get_blob
|
||||
|
||||
def get_blob(self: FileStore, kind: str, sha: str, **kwargs: Any) -> bytes | None:
|
||||
if kind == "excerpts":
|
||||
code = CODE.get(sha)
|
||||
return code.encode() if code 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("resolved an endpoint"))
|
||||
args = [str(manifest.path), "--memory", f"file://{tmp_path / 'memory'}"]
|
||||
before = snapshot(tmp_path)
|
||||
assert main(["learn", "experiment", "fill", *args, "--arm", "B", "--dry-run"]) == 0
|
||||
assert json.loads(capsys.readouterr().out)["arms"] == ["B"] and snapshot(tmp_path) == before
|
||||
assert main(["learn", "experiment", "run", *args, "--replay-only"]) == 2 # missing cache, no endpoint
|
||||
capsys.readouterr()
|
||||
fill_experiment(store, manifest, EXAMPLES, PRICES, excerpt_text, arms=["A", "B"], client=ScriptedChat(), budget_usd=1)
|
||||
assert main(["learn", "experiment", "run", *args, "--replay-only", "--no-early-stop"]) == 0
|
||||
report = json.loads(capsys.readouterr().out)
|
||||
assert report["decided"] == {"A": len(units), "B": len(units)}
|
||||
assert all(arm["client_calls"] == 0 for arm in report["spend"].values())
|
||||
|
||||
@@ -118,8 +118,6 @@ def test_learn_slice_default_prompt_frozen_before_change() -> None:
|
||||
"f6b5cfb5a8b7cb1591d7b4534e254ad2a349452a03e825bdcf250e5f60b4e80b"
|
||||
)
|
||||
assert baseline.ProgramSpec().slice is False
|
||||
reserved, _ = baseline.program(slice=True).prepare(baseline.candidate(), baseline.FOLD)
|
||||
assert reserved.messages == prompt.messages # no prompt wiring before audit review
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -192,3 +190,38 @@ def test_learn_slice_python_fallback_order(monkeypatch: pytest.MonkeyPatch) -> N
|
||||
|
||||
def test_learn_slice_bare_arrow_parameter() -> None:
|
||||
assert sliced("const f = x => eval(x);", "javascript").steps
|
||||
|
||||
|
||||
def test_learn_slice_prompt_flow_examples_citations_and_request_identity() -> None:
|
||||
import test_learn_program as baseline
|
||||
|
||||
from openultrasast.learn.program import parse_answer, render, request_key
|
||||
|
||||
example = replace(baseline.EXAMPLES[0], language="python", family="injection")
|
||||
code = " 40 def f(x):\n 41 execute(x)\n"
|
||||
demo = baseline.Demo(example.id, example.excerpt_sha, 1, "injection", "vulnerable", 0.9, "L41 executes x.", (41,))
|
||||
spec = baseline.ProgramSpec(demos=(demo,), slice=True)
|
||||
candidate = baseline.candidate(code)
|
||||
prompt = render(spec, candidate, [example], [demo], {example.id: example}, lambda _: code)
|
||||
block = render_slice(function_slice(code, "python", "injection"), "python")
|
||||
assert "L40 source:" in block and "L41 sink:" in block
|
||||
assert str(prompt.messages[1]["content"]).count(code.rstrip() + "\n" + block + "\nSignals:") == 2
|
||||
assert code.rstrip() + "\n" + block + "\nSignals:" in prompt.prefix
|
||||
assert prompt.valid_lines == frozenset({40, 41})
|
||||
answer = json.dumps(
|
||||
{"verdict": "vulnerable", "confidence": 0.9, "family": "injection", "rationale": "L41 executes x.", "cited_lines": [40, 41]}
|
||||
)
|
||||
assert parse_answer(answer, prompt.valid_lines).cited_lines == (40, 41)
|
||||
default = render(replace(spec, slice=False), candidate, [example], [demo], {example.id: example}, lambda _: code)
|
||||
assert request_key(spec.model, "params", prompt.messages, 0.0, "0", True) != request_key(
|
||||
spec.model, "params", default.messages, 0.0, "0", True
|
||||
)
|
||||
|
||||
|
||||
def test_learn_slice_prompt_not_analysable() -> None:
|
||||
import test_learn_program as baseline
|
||||
|
||||
candidate = replace(baseline.candidate(), language="php")
|
||||
prompt, _ = baseline.program(slice=True).prepare(candidate, baseline.FOLD)
|
||||
tail = str(prompt.messages[1]["content"]).split("Candidate:\n")[1]
|
||||
assert candidate.code.rstrip() + "\nData flow inside the function: not analysed (php is not supported).\nSignals:" in tail
|
||||
|
||||
Reference in New Issue
Block a user