mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
feat(learn): the classifier program and its response cache (decision-engine 6.4)
learn/program.py: the declared signature (code, signals, roles, family in; verdict, family, confidence, rationale with cited lines out), the baseline instruction, rendering as [system: instruction + demonstrations] -- byte-identical for every candidate, so the provider's prefix cache can hit -- then [user: retrieved examples with their labels, the candidate without one]. Signals pass validate_x first (no EXCLUDED_FIELDS, missing instruments in words). Answers are one JSON object; one that does not parse or cites no line inside the excerpt is asked once more at temperature 0 and otherwise becomes unsure/parse_failed, never not_vulnerable. Classify(k): sample 1 at 0.0, the rest at 0.7; score s = mean p_j; majority verdict, ties unsure; a majority unsure cannot BLOCK. Caller caches each response as responses/<sha>.json keyed by (model, Model parameters digest, messages, temperature, sample, json_object) and replays at $0; spend is priced from the recorded usage once per distinct request, so a replay reports the same usd. Program.decide asserts the evaluation boundary on the rendered ids before any call. MeteredClient.complete passes temperature/logprobs through (only when set); DeepSeekChatClient sends temperature beside thinking=disabled. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
f5180a14ec
commit
ea61dc1bb3
+47
-1
@@ -3,11 +3,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from collections.abc import Mapping, Sequence
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from openultrasast.learn.examples import Example
|
||||
from openultrasast.learn.features import NOT_APPLICABLE, Part, engine_part, quick_part, record
|
||||
from openultrasast.tool_hunter import ChatResponse
|
||||
|
||||
|
||||
def x_record(*, hits: int = 0, engine: str = "ran", findings: int = 0, language: str = "python", prior_hits: int = 0) -> dict[str, Any]:
|
||||
@@ -62,3 +64,47 @@ def corpus(groups: int = 12, per_group: int = 4) -> list[Example]:
|
||||
|
||||
def by_id(examples: Sequence[Example]) -> Mapping[str, Example]:
|
||||
return {e.id: e for e in examples}
|
||||
|
||||
|
||||
def candidate_lines(messages: Sequence[Mapping[str, object]]) -> list[int]:
|
||||
"""The line numbers of the candidate's code in a rendered prompt."""
|
||||
text = str(messages[-1]["content"]).rsplit("Candidate:\n", 1)[1]
|
||||
code = text.split("Code:\n", 1)[1].split("\nSignals:\n", 1)[0]
|
||||
return [int(line.split()[0]) for line in code.splitlines() if line.strip() and line.split()[0].isdigit()]
|
||||
|
||||
|
||||
def default_answer(messages: Sequence[Mapping[str, object]], temperature: float | None) -> str:
|
||||
"""Vulnerable when the candidate's code executes something, else not; cites the candidate's first line."""
|
||||
text = str(messages[-1]["content"]).rsplit("Candidate:\n", 1)[1]
|
||||
vulnerable = "execute(" in text
|
||||
lines = candidate_lines(messages)
|
||||
return json.dumps(
|
||||
{
|
||||
"verdict": "vulnerable" if vulnerable else "not_vulnerable", "family": "injection", "confidence": 0.8 if vulnerable else 0.7,
|
||||
"rationale": "The value reaches the call on the cited line.", "cited_lines": lines[:1],
|
||||
}
|
||||
) # fmt: skip
|
||||
|
||||
|
||||
class ScriptedChat:
|
||||
"""A chat client whose answers a function of the prompt scripts; records every call and a usage row per call."""
|
||||
|
||||
def __init__(self, answer: Callable[[Sequence[Mapping[str, object]], float | None], str] = default_answer) -> None:
|
||||
self.answer = answer
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.usage: list[dict[str, object]] = []
|
||||
|
||||
def complete(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]],
|
||||
timeout_seconds: int = 60,
|
||||
json_object: bool = False,
|
||||
temperature: float | None = None,
|
||||
) -> ChatResponse:
|
||||
self.calls.append({"model": model, "messages": messages, "temperature": temperature, "json_object": json_object})
|
||||
prompt = sum(len(str(m["content"])) for m in messages) // 4
|
||||
self.usage.append({"prompt_tokens": prompt, "prompt_cache_hit_tokens": prompt // 2, "completion_tokens": 60})
|
||||
return ChatResponse(content=self.answer(messages, temperature))
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
"""The program (learned-decision-engine design 4.1, task 6.4): rendering without label or identity fields, the
|
||||
byte-identical prefix, answer parsing, the vote, ``unsure`` never blocks, the metered client's temperature and
|
||||
ceiling, and the response cache's replay. Every client is scripted."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from learn_fixtures import ScriptedChat, by_id, corpus, x_record
|
||||
|
||||
from openultrasast.learn import program as program_module
|
||||
from openultrasast.learn.folds import outer_folds
|
||||
from openultrasast.learn.program import (
|
||||
PARSE_FAILED,
|
||||
Answer,
|
||||
Caller,
|
||||
Candidate,
|
||||
Demo,
|
||||
Program,
|
||||
ProgramSpec,
|
||||
ReplayMiss,
|
||||
aggregate,
|
||||
classify,
|
||||
line_numbers,
|
||||
parse_answer,
|
||||
render,
|
||||
request_key,
|
||||
)
|
||||
from openultrasast.learn.retrieve import BoundaryViolation, Retrieval, Target
|
||||
from openultrasast.learn.schema import EXCLUDED_FIELDS, FeatureRecordError
|
||||
from openultrasast.model.endpoint import DeepSeekChatClient
|
||||
from openultrasast.plane.budget import BudgetExhausted, MeteredClient
|
||||
from openultrasast.plane.memory import FileStore
|
||||
|
||||
PRICES = {"cache_hit_per_m": 0.014, "input_per_m": 0.44, "output_per_m": 1.32}
|
||||
EXAMPLES = corpus(groups=10)
|
||||
INDEX = by_id(EXAMPLES)
|
||||
CODE = {
|
||||
e.excerpt_sha: f" 10 def f{i}(q):\n 11 return {'execute(q)' if e.label else 'escape(q)'}\n" for i, e in enumerate(EXAMPLES)
|
||||
}
|
||||
|
||||
|
||||
def excerpt_text(sha: str) -> str | None:
|
||||
return CODE.get(sha)
|
||||
|
||||
|
||||
def candidate(code: str = " 40 def view(q):\n 41 return execute(q)\n", group: str = "cand/repo", **kw: object) -> Candidate:
|
||||
rec = x_record(**kw) # type: ignore[arg-type]
|
||||
return Candidate(Target(group, "injection", "static", rec["x"], rec["instruments"]), code, "python")
|
||||
|
||||
|
||||
def program(**spec: object) -> Program:
|
||||
return Program(ProgramSpec(**spec), EXAMPLES, excerpt_text) # type: ignore[arg-type]
|
||||
|
||||
|
||||
FOLD = replace(outer_folds(EXAMPLES, k=5)[0], eval_groups=frozenset({"cand/repo"}))
|
||||
|
||||
|
||||
def test_prompts_hold_no_excluded_field_and_no_candidate_label() -> None:
|
||||
prompt, got = program().prepare(candidate(), FOLD)
|
||||
text = "\n".join(str(m["content"]) for m in prompt.messages)
|
||||
for name in EXCLUDED_FIELDS:
|
||||
assert not re.search(rf"^{re.escape(name)}\s*[:=]", text, re.MULTILINE), name
|
||||
candidate_part = text.rsplit("Candidate:\n", 1)[1]
|
||||
assert "label" not in candidate_part and "vulnerable" not in candidate_part
|
||||
assert len(got.examples) == 6 and text.count("(label: ") == 6
|
||||
leaky = candidate()
|
||||
x = {**leaky.target.x, "site_match": True}
|
||||
with pytest.raises(FeatureRecordError, match="site_match"):
|
||||
render(ProgramSpec(), replace(leaky, target=replace(leaky.target, x=x)), [], [], INDEX, excerpt_text)
|
||||
|
||||
|
||||
def test_missing_instruments_are_said_in_words() -> None:
|
||||
prompt, _ = program().prepare(candidate(engine="none"), FOLD)
|
||||
tail = str(prompt.messages[-1]["content"]).rsplit("Candidate:\n", 1)[1]
|
||||
assert "engine: no coverage" in tail and "delta: not applicable" in tail and "eng.findings" not in tail
|
||||
|
||||
|
||||
def test_the_prefix_is_byte_identical_across_candidates() -> None:
|
||||
demo_example = EXAMPLES[0]
|
||||
demo = Demo(demo_example.id, demo_example.excerpt_sha, 0, "injection", "not_vulnerable", 0.8, "Line 11 escapes q.", (11,))
|
||||
prog = program(demos=(demo,))
|
||||
a, _ = prog.prepare(candidate(), FOLD)
|
||||
b, _ = prog.prepare(candidate(" 7 def other(x):\n 8 return x\n", hits=3), FOLD)
|
||||
assert a.messages[0] == b.messages[0] and a.prefix.encode() == b.prefix.encode()
|
||||
assert "Demonstration 1:" in a.prefix and a.messages[1] != b.messages[1]
|
||||
|
||||
|
||||
def test_parsing_valid_invalid_and_out_of_range_citations() -> None:
|
||||
lines = frozenset({40, 41})
|
||||
ok = parse_answer(
|
||||
'{"verdict": "vulnerable", "family": "injection", "confidence": 0.9, "rationale": "r", "cited_lines": [41, 99]}', lines
|
||||
)
|
||||
assert ok == Answer("vulnerable", "injection", 0.9, "r", (41,))
|
||||
fenced = '```json\n{"verdict": "unsure", "family": "x", "confidence": 0.5, "rationale": "r", "cited_lines": [40]}\n```'
|
||||
assert parse_answer(fenced, lines, families=("injection",)) == Answer("unsure", "none", 0.5, "r", (40,))
|
||||
for bad in (
|
||||
"not json", '["vulnerable"]', '{"verdict": "maybe", "confidence": 0.5, "rationale": "r", "cited_lines": [40]}',
|
||||
'{"verdict": "vulnerable", "confidence": 1.5, "rationale": "r", "cited_lines": [40]}',
|
||||
'{"verdict": "vulnerable", "confidence": 0.5, "rationale": "", "cited_lines": [40]}',
|
||||
'{"verdict": "vulnerable", "confidence": 0.5, "rationale": "r", "cited_lines": [99]}', "",
|
||||
): # fmt: skip
|
||||
assert parse_answer(bad, lines) is None, bad
|
||||
assert line_numbers(" 40 a\n 41 b\ndiff:\n@@ -1 +1 @@\n 42 not a line") == {40, 41}
|
||||
|
||||
|
||||
def test_unparsable_answers_are_retried_once_then_unsure_never_not_vulnerable() -> None:
|
||||
chat = ScriptedChat(lambda messages, t: "I think it is fine.")
|
||||
caller = Caller(MeteredClient(chat, prices=PRICES), "deepseek-flash", PRICES)
|
||||
prompt, _ = program().prepare(candidate(), FOLD)
|
||||
decision = classify(prompt, caller, 1)
|
||||
assert len(chat.calls) == 2 and [c["temperature"] for c in chat.calls] == [0.0, 0.0]
|
||||
assert decision.verdict == "unsure" and decision.parse_failed == 1 and decision.s == 0.5 and not decision.blockable
|
||||
|
||||
|
||||
def test_vote_score_and_unsure_never_blocks() -> None:
|
||||
v = Answer("vulnerable", "injection", 0.9, "r", (1,))
|
||||
n = Answer("not_vulnerable", "injection", 0.8, "r", (1,))
|
||||
u = Answer("unsure", "none", 0.4, "r", (1,))
|
||||
d = aggregate([v, v, v, n, u, PARSE_FAILED])
|
||||
assert d.verdict == "vulnerable" and d.s == pytest.approx((0.9 * 3 + 0.2 + 0.5 + 0.5) / 6) and d.blockable
|
||||
assert d.votes == {"vulnerable": 3, "not_vulnerable": 1, "unsure": 2} and d.parse_failed == 1
|
||||
assert aggregate([v, v, u, PARSE_FAILED]).verdict == "unsure" # parse failures vote unsure: a tie is unsure
|
||||
tie = aggregate([v, n])
|
||||
assert tie.verdict == "unsure" and not tie.blockable
|
||||
unsure = aggregate([u, u, v])
|
||||
assert unsure.verdict == "unsure" and not unsure.blockable
|
||||
|
||||
|
||||
def test_metered_client_receives_temperature_and_stops_at_its_ceiling() -> None:
|
||||
chat = ScriptedChat()
|
||||
caller = Caller(MeteredClient(chat, prices=PRICES, budget_usd=10.0), "deepseek-flash", PRICES)
|
||||
decision = program(k=3).decide(candidate(), FOLD, caller)
|
||||
assert [c["temperature"] for c in chat.calls] == [0.0, 0.7, 0.7] and all(c["json_object"] for c in chat.calls)
|
||||
assert decision.verdict == "vulnerable" and decision.votes["vulnerable"] == 3 and decision.retrieval == "signals_only"
|
||||
assert decision.neighbours["label"] == {"0": 3, "1": 3}
|
||||
capped = Caller(MeteredClient(ScriptedChat(), prices=PRICES, budget_calls=2), "deepseek-flash", PRICES)
|
||||
with pytest.raises(BudgetExhausted):
|
||||
program(k=3).decide(candidate(), FOLD, capped)
|
||||
|
||||
|
||||
def test_deepseek_client_forwards_temperature() -> None:
|
||||
seen: list[dict] = []
|
||||
|
||||
class Raw:
|
||||
def complete_chat_raw(self, **kw: object) -> dict:
|
||||
seen.append(dict(kw["extra_body"])) # type: ignore[arg-type]
|
||||
return {"choices": [{"message": {"content": "{}"}}], "usage": {"prompt_tokens": 3}}
|
||||
|
||||
client = DeepSeekChatClient(Raw()) # type: ignore[arg-type]
|
||||
MeteredClient(client).complete(model="m", messages=[], tools=[], json_object=True, temperature=0.7)
|
||||
assert seen[0]["temperature"] == 0.7 and seen[0]["thinking"] == {"type": "disabled"}
|
||||
client.complete(model="m", messages=[], tools=[])
|
||||
assert "temperature" not in seen[1]
|
||||
|
||||
|
||||
def test_response_cache_replays_at_zero_cost(tmp_path: Path) -> None:
|
||||
store = FileStore(tmp_path / "memory")
|
||||
chat = ScriptedChat()
|
||||
first = Caller(MeteredClient(chat, prices=PRICES), "deepseek-flash", PRICES, store=store)
|
||||
one = program(k=2).decide(candidate(), FOLD, first)
|
||||
assert first.calls == 2 and len(store.blob_names("responses")) == 2
|
||||
replay = Caller(None, "deepseek-flash", PRICES, store=store)
|
||||
two = program(k=2).decide(candidate(), FOLD, replay)
|
||||
assert replay.calls == 0 and replay.replayed == 2 and one == two and replay.usd() == first.usd() and first.usd()
|
||||
assert replay.cache_digest() == first.cache_digest()
|
||||
with pytest.raises(ReplayMiss):
|
||||
program(k=2).decide(candidate(" 1 def z():\n 2 return 0\n"), FOLD, replay)
|
||||
key = request_key("deepseek-flash", first.params_digest, [{"role": "user", "content": "x"}], 0.0, "0", True)
|
||||
assert key != request_key("deepseek-flash", first.params_digest, [{"role": "user", "content": "x"}], 0.7, "0", True)
|
||||
stored = json.loads(store.get_blob("responses", store.blob_names("responses")[0], verify=False) or b"{}")
|
||||
assert set(stored) == {"content", "model", "usage"}
|
||||
|
||||
|
||||
def test_the_boundary_is_asserted_before_any_call(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
sibling = next(e for e in EXAMPLES if e.group == EXAMPLES[0].group and e.id != EXAMPLES[0].id)
|
||||
target_group = EXAMPLES[0].group
|
||||
|
||||
def leaky(*args: object, **kw: object) -> Retrieval:
|
||||
return Retrieval((sibling,), "signals_only", "short", 0, 1, 0)
|
||||
|
||||
monkeypatch.setattr(program_module, "retrieve", leaky)
|
||||
chat = ScriptedChat()
|
||||
with pytest.raises(BoundaryViolation):
|
||||
program().decide(candidate(group=target_group), FOLD, Caller(MeteredClient(chat), "deepseek-flash"))
|
||||
assert chat.calls == []
|
||||
Reference in New Issue
Block a user