From ea61dc1bb3314e2b9e248996004d80f00f502772 Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Thu, 1 Oct 2026 10:13:31 +0200 Subject: [PATCH] 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/.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 Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo --- .kiro/specs/learned-decision-engine/tasks.md | 2 +- .../measurements/2026-09-08-module-audit.json | 27 +- src/openultrasast/learn/program.py | 473 ++++++++++++++++++ src/openultrasast/model/audit.py | 6 +- src/openultrasast/model/endpoint.py | 20 +- src/openultrasast/plane/budget.py | 11 +- tests/learn_fixtures.py | 48 +- tests/test_learn_program.py | 190 +++++++ 8 files changed, 764 insertions(+), 13 deletions(-) create mode 100644 src/openultrasast/learn/program.py create mode 100644 tests/test_learn_program.py diff --git a/.kiro/specs/learned-decision-engine/tasks.md b/.kiro/specs/learned-decision-engine/tasks.md index 3f93042..261af21 100644 --- a/.kiro/specs/learned-decision-engine/tasks.md +++ b/.kiro/specs/learned-decision-engine/tasks.md @@ -61,7 +61,7 @@ a model or embedding call. - Tests: `test_learn_retrieve.py` including `test_boundary` (every rendered example and demonstration eligible, for all fold kinds; `memory.seed` excludes the case's group), `test_learn_folds.py`. - _Requirements: 3.2, 3.4, 7.1, 8.3_ - - [ ] 6.4 The program (no model in tests) + - [x] 6.4 The program (no model in tests) - `learn/program.py`: the signature, prompt rendering (instruction + demonstrations as a byte-identical prefix, then retrieved examples, then the candidate), JSON parsing with one retry -> `unsure, parse_failed`, `Classify(k)` with temperature 0 then 0.7, the score `s` and majority verdict; `responses/.json` cache diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index 694175a..7f95b28 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 149, + "modules": 150, "counts": { - "load_bearing": 142, + "load_bearing": 143, "standalone_capability": 7 }, "orphaned": [], @@ -358,6 +358,7 @@ "learn.excerpt", "learn.features", "learn.folds", + "learn.program", "learn.retrieve", "learn.roles", "plane.tasks.features", @@ -389,6 +390,7 @@ "importers": [ "learn.cli", "learn.folds", + "learn.program", "learn.retrieve" ] }, @@ -413,8 +415,9 @@ { "module": "learn.folds", "classification": "load_bearing", - "reason": "imported by learn.retrieve", + "reason": "imported by learn.program, learn.retrieve", "importers": [ + "learn.program", "learn.retrieve" ] }, @@ -430,11 +433,19 @@ "learn.retrieve" ] }, + { + "module": "learn.program", + "classification": "load_bearing", + "reason": "learned-decision-engine Req 3.1/3.3: the AI classifier program (signature, Retrieve -> Classify(k), response cache); `ousast learn compile` (task 6.5) and the decide task (task 8) are its callers", + "importers": [] + }, { "module": "learn.retrieve", "classification": "load_bearing", - "reason": "learned-decision-engine Req 3.2/3.4: retrieval from the local memory and the evaluation boundary; the program (task 6.4) and `ousast learn compile` (6.5) are its callers", - "importers": [] + "reason": "imported by learn.program", + "importers": [ + "learn.program" + ] }, { "module": "learn.roles", @@ -452,6 +463,7 @@ "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "learn.features", + "learn.program", "learn.retrieve", "learn.roles", "plane.tasks.remember", @@ -785,6 +797,7 @@ "learn.cli", "learn.embeddings", "learn.examples", + "learn.program", "plane.egress", "plane.engine_alerts", "plane.generate", @@ -806,6 +819,7 @@ "importers": [ "learn.cli", "learn.embeddings", + "learn.program", "plane.generate", "plane.tasks.roles", "plane.tasks.verify" @@ -864,6 +878,7 @@ "learn.cli", "learn.embeddings", "learn.examples", + "learn.program", "plane.generate", "plane.tasks.loop", "plane.tasks.remember" @@ -1073,6 +1088,7 @@ "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "learn.cli", + "learn.program", "model.endpoint", "tool_hunter" ] @@ -1083,6 +1099,7 @@ "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "learn.cli", + "learn.program", "model.endpoint", "tool_hunter" ] diff --git a/src/openultrasast/learn/program.py b/src/openultrasast/learn/program.py new file mode 100644 index 0000000..7a29ff5 --- /dev/null +++ b/src/openultrasast/learn/program.py @@ -0,0 +1,473 @@ +"""The decision engine's program: a declared language-model classifier with local memory (design 4.1, Req 3.1). + +DSPy-style and in-house. The **signature** (:data:`SIGNATURE`) declares what goes in -- the candidate's code +excerpt, its allow-listed signals, its inferred roles, its family -- and what comes out: ``verdict`` +(``vulnerable | not_vulnerable | unsure``), ``family``, ``confidence`` in [0, 1] and a ``rationale`` of at most three +sentences citing at least one line of the code. The **modules** are ``Retrieve -> Classify(k)``: + +- the prompt is the instruction plus the compiled demonstrations -- a prefix that is byte-identical for every + candidate, so the provider's prefix cache hits -- then the retrieved examples, then the candidate; +- sample 1 runs at temperature 0, samples 2..k at 0.7; an answer that does not parse, or cites no line in range, is + asked once more at temperature 0 and otherwise becomes ``unsure`` with ``parse_failed`` -- never ``not_vulnerable``; +- the raw score is ``s = mean_j p_j`` (``p_j`` the confidence of a ``vulnerable`` answer, one minus it for + ``not_vulnerable``, 0.5 for ``unsure``); the verdict is the majority (ties are ``unsure``), and a majority + ``unsure`` can never BLOCK. + +Signals reach the prompt only through :func:`.schema.validate_x` (no ``EXCLUDED_FIELDS``, no free text); a missing +instrument is shown in words. Every model call goes through :class:`Caller`: a response cache keyed by +``sha256(model, Model parameters digest, messages, temperature, sample, json_object)`` stored as +``responses/.json``, so a rerun on the same inputs replays at $0 and reproduces byte for byte. Before a call the +rendered example and demonstration ids are checked against the evaluation boundary +(:func:`.retrieve.assert_boundary`). +""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import asdict, dataclass, field, replace +from typing import Any + +from ..plane.budget import cost_of, prices_from +from ..plane.memory import MemoryStore +from .examples import Example +from .folds import Fold +from .retrieve import Retrieval, Target, assert_boundary, demonstrations_for, retrieve +from .schema import features_for, instruments_for, validate_x + +VERDICTS = ("vulnerable", "not_vulnerable", "unsure") +UNSURE = "unsure" +TEMPERATURES = (0.0, 0.7) +MAX_OUTPUT_TOKENS = 300 +CHARS_PER_TOKEN = 4.0 +STATE_WORDS = {"none": "no coverage", "failed": "failed", "not_applicable": "not applicable"} +_LINE_NUMBER = re.compile(r"^\s*(\d+) ") + + +class ReplayMiss(RuntimeError): + """A replay-only run (no client) met a request that is not in the response cache.""" + + +@dataclass(frozen=True) +class Field: + name: str + direction: str # in | out + type: str + bound: str + + +@dataclass(frozen=True) +class Signature: + fields: tuple[Field, ...] + + @property + def digest(self) -> str: + return hashlib.sha256(json.dumps([asdict(f) for f in self.fields], sort_keys=True).encode()).hexdigest() + + +SIGNATURE = Signature( + ( + Field("code", "in", "excerpt", "numbered lines, <= 80 lines and 4,000 characters; a delta adds the diff (<= 40 lines)"), + Field("signals", "in", "record", "the allow-listed features as name: value lines; a missing instrument in words"), + Field("roles", "in", "roles", "inferred roles of the function: role, operation, line, origin"), + Field("family", "in", "family", "the candidate's family, or unknown"), + Field("verdict", "out", "enum", "vulnerable | not_vulnerable | unsure"), + Field("family", "out", "family", "one of the families, or none"), + Field("confidence", "out", "float", "[0, 1]: the probability that the verdict is right"), + Field("rationale", "out", "text", "at most 3 sentences; cites at least one line number of code (cited_lines)"), + ) +) + +BASELINE_INSTRUCTION = ( + "You decide whether a function contains an exploitable security vulnerability of the given family.\n" + "You see the function's code (numbered lines), the signals static instruments emitted for it, and the roles\n" + "inferred for its calls. Signals are evidence, not verdicts: rules over-report, and a missing instrument is not a\n" + "negative. Similar labelled functions from other repositories follow as examples; reason about this code.\n" + 'Answer with one JSON object: {"verdict": "vulnerable" | "not_vulnerable" | "unsure", "family": , "confidence": , "rationale": ,\n' + '"cited_lines": []}. Cite at least one line. Say "unsure"\n' + "when the code does not decide it." +) + + +# --- rendering --------------------------------------------------------------------------------------------------------- + + +def render_signals(x: Mapping[str, Any], instruments: Mapping[str, Mapping[str, Any]], profile: str) -> str: + """``name: value`` per feature of an instrument that ran; ``: no coverage | failed | not applicable`` + otherwise. The record is validated first: an unknown, label or identity field never reaches a prompt.""" + validate_x(x, instruments, profile) + lines: list[str] = [] + for name in instruments_for(profile): + if name == "language": + continue + state = str(instruments[name]["state"]) + if state != "ran": + lines.append(f"{name}: {STATE_WORDS[state]}") + continue + for spec in features_for(profile): + if spec.instrument == name: + value = x[spec.name] + lines.append(f"{spec.name}: {'null' if value is None else json.dumps(value)}") + return "\n".join(lines) + + +def render_roles(roles: Sequence[Mapping[str, Any]]) -> str: + if not roles: + return "none inferred" + keys = ("role", "kind", "operation", "line", "origin") + return "\n".join(", ".join(f"{k}={r[k]}" for k in keys if k in r and k != "path") for r in roles) + + +def _case( + code: str, + x: Mapping[str, Any], + instruments: Mapping[str, Mapping[str, Any]], + roles: Sequence[Mapping[str, Any]], + family: str, + profile: str, + language: str, +) -> str: + signals = render_signals(x, instruments, profile) + return f"Family: {family}\nLanguage: {language}\nCode:\n{code.rstrip()}\nSignals:\n{signals}\nRoles:\n{render_roles(roles)}" + + +def label_word(label: int) -> str: + return "vulnerable" if label == 1 else "not_vulnerable" + + +@dataclass(frozen=True) +class Demo: + """A compiled demonstration: an example of ``C_boot`` with the rationale the program gave without seeing its label.""" + + example_id: str + excerpt_sha: str + label: int + family: str + verdict: str + confidence: float + rationale: str + cited_lines: tuple[int, ...] + alternate: Demo | None = None + + def as_dict(self) -> dict[str, Any]: + out = {k: v for k, v in asdict(self).items() if k != "alternate"} + out["cited_lines"] = list(self.cited_lines) + out["alternate"] = self.alternate.as_dict() if self.alternate is not None else None + return out + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Demo: + alt = data.get("alternate") + return cls( + str(data["example_id"]), str(data["excerpt_sha"]), int(data["label"]), str(data["family"]), str(data["verdict"]), + float(data["confidence"]), str(data["rationale"]), tuple(int(n) for n in data.get("cited_lines") or ()), + cls.from_dict(alt) if isinstance(alt, Mapping) else None, + ) # fmt: skip + + +@dataclass(frozen=True) +class ProgramSpec: + """What a compiled program fixes: instruction, demonstrations, model and sampling, retrieval settings.""" + + instruction: str = BASELINE_INSTRUCTION + demos: tuple[Demo, ...] = () + profile: str = "static" + model: str = "deepseek-flash" + k: int = 1 + k_ret: int = 6 + lam: float = 0.5 + n1: int = 40 + priors: str = "off" + max_output_tokens: int = MAX_OUTPUT_TOKENS + + +@dataclass(frozen=True) +class Candidate: + """A candidate as the program sees it: retrieval's target and the rendered (bounded, redacted) excerpt.""" + + target: Target + code: str + language: str + roles: tuple[Mapping[str, Any], ...] = () + + +@dataclass(frozen=True) +class Prompt: + messages: list[dict[str, object]] + example_ids: tuple[str, ...] + demo_ids: tuple[str, ...] + valid_lines: frozenset[int] + + @property + def prefix(self) -> str: + return str(self.messages[0]["content"]) + + @property + def estimated_tokens(self) -> int: + return round(sum(len(str(m["content"])) for m in self.messages) / CHARS_PER_TOKEN) + + +def line_numbers(code: str) -> frozenset[int]: + """The source line numbers a rendered excerpt shows (the diff part has none).""" + return frozenset(int(m.group(1)) for line in code.split("\ndiff:\n", 1)[0].splitlines() if (m := _LINE_NUMBER.match(line))) + + +ExcerptText = Callable[[str], str | None] + + +def render_prefix(spec: ProgramSpec, demos: Sequence[Demo], examples: Mapping[str, Example], excerpt_text: ExcerptText) -> str: + """The instruction and the demonstrations: identical for every candidate of one program and fold.""" + parts = [spec.instruction.strip()] + for number, demo in enumerate(demos, 1): + example = examples[demo.example_id] + code = excerpt_text(demo.excerpt_sha) or "" + answer = {"verdict": demo.verdict, "family": demo.family, "confidence": demo.confidence, "rationale": demo.rationale, + "cited_lines": list(demo.cited_lines)} # fmt: skip + parts.append( + f"Demonstration {number}:\n" + + _case(code, example.x, example.instruments, example.roles, example.family, spec.profile, example.language) + + f"\nAnswer: {json.dumps(answer, sort_keys=True)}" + ) + return "\n\n".join(parts) + + +def render( + spec: ProgramSpec, + candidate: Candidate, + retrieved: Sequence[Example], + demos: Sequence[Demo], + examples: Mapping[str, Example], + excerpt_text: ExcerptText, +) -> Prompt: + """The prompt: [system: instruction + demonstrations] [user: retrieved examples (with their labels) + the candidate + (without one)].""" + blocks = [] + for number, example in enumerate(retrieved, 1): + code = excerpt_text(example.excerpt_sha) or "" + blocks.append( + f"Example {number} (label: {label_word(example.label)}):\n" + + _case(code, example.x, example.instruments, example.roles, example.family, spec.profile, example.language) + ) + t = candidate.target + blocks.append("Candidate:\n" + _case(candidate.code, t.x, t.instruments, candidate.roles, t.family, spec.profile, candidate.language)) + messages: list[dict[str, object]] = [ + {"role": "system", "content": render_prefix(spec, demos, examples, excerpt_text)}, + {"role": "user", "content": "\n\n".join(blocks)}, + ] + return Prompt(messages, tuple(e.id for e in retrieved), tuple(d.example_id for d in demos), line_numbers(candidate.code)) + + +# --- parsing ----------------------------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Answer: + verdict: str + family: str + confidence: float + rationale: str + cited_lines: tuple[int, ...] + parse_failed: bool = False + + @property + def p(self) -> float: + """The answer's probability that the candidate is vulnerable.""" + if self.verdict == "vulnerable": + return self.confidence + if self.verdict == "not_vulnerable": + return 1.0 - self.confidence + return 0.5 + + +PARSE_FAILED = Answer(UNSURE, "none", 0.0, "", (), parse_failed=True) + + +def parse_answer(content: str | None, valid_lines: frozenset[int], families: Sequence[str] = ()) -> Answer | None: + """The answer, or ``None`` when it does not parse: not one JSON object, a verdict outside the vocabulary, a + confidence outside [0, 1], no rationale, or no cited line inside the excerpt.""" + from ..provider.openrouter import OpenRouterError, parse_json_content + + try: + data = parse_json_content(content or "") + except OpenRouterError: + return None + if not isinstance(data, dict) or data.get("verdict") not in VERDICTS: + return None + confidence = data.get("confidence") + if isinstance(confidence, bool) or not isinstance(confidence, int | float) or not 0.0 <= float(confidence) <= 1.0: + return None + rationale = data.get("rationale") + if not isinstance(rationale, str) or not rationale.strip(): + return None + cited = data.get("cited_lines") + lines = ( + tuple(sorted({int(n) for n in cited if isinstance(n, int) and not isinstance(n, bool) and n in valid_lines})) + if isinstance(cited, list) + else () + ) + if not lines: + return None + family = data.get("family") + family = family if isinstance(family, str) and (not families or family in families or family == "none") else "none" + return Answer(str(data["verdict"]), family, round(float(confidence), 6), rationale.strip()[:600], lines) + + +# --- calls and the response cache -------------------------------------------------------------------------------------- + + +def request_key( + model: str, params_digest: str, messages: Sequence[Mapping[str, object]], temperature: float, sample: str, json_object: bool +) -> str: + canonical = {"json_object": json_object, "messages": list(messages), "model": model, "params": params_digest, "sample": sample, + "temperature": temperature} # fmt: skip + return hashlib.sha256(json.dumps(canonical, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest() + + +def params_digest(parameters: Mapping[str, Any] | None) -> str: + return hashlib.sha256(json.dumps(dict(parameters or {}), sort_keys=True).encode()).hexdigest() + + +@dataclass +class Caller: + """Every model call of a program run: cached by request in ``responses/.json`` (and in memory), replayed at + $0. ``client`` is a metered client (``plane.budget.MeteredClient``); ``None`` is replay-only.""" + + client: Any + model: str + parameters: Mapping[str, Any] = field(default_factory=dict) + store: MemoryStore | None = None + calls: int = 0 + replayed: int = 0 + keys: list[str] = field(default_factory=list) + usage: dict[str, int] = field(default_factory=lambda: {"prompt_tokens": 0, "prompt_cache_hit_tokens": 0, "completion_tokens": 0}) + _memory: dict[str, dict[str, Any]] = field(default_factory=dict) + + @property + def params_digest(self) -> str: + return params_digest(self.parameters) + + def _cached(self, key: str) -> dict[str, Any] | None: + if key in self._memory: + return self._memory[key] + if self.store is not None: + data = self.store.get_blob("responses", key, verify=False) + if data is not None: + return dict(json.loads(data)) + return None + + def ask(self, messages: list[dict[str, object]], *, temperature: float, sample: str, json_object: bool = True) -> str: + key = request_key(self.model, self.params_digest, messages, temperature, sample, json_object) + self.keys.append(key) + entry = self._cached(key) + if entry is not None: + self.replayed += 1 + else: + if self.client is None: + raise ReplayMiss(f"response {key[:12]} is not cached and the run is replay-only") + before = dict(getattr(self.client, "usage", {}) or {}) + response = self.client.complete(model=self.model, messages=messages, tools=[], json_object=json_object, temperature=temperature) + after = dict(getattr(self.client, "usage", {}) or {}) + usage = {k: int(after.get(k, 0)) - int(before.get(k, 0)) for k in self.usage} + entry = {"content": response.content, "model": self.model, "usage": usage} + self.calls += 1 + if self.store is not None: + self.store.put_blob("responses", json.dumps(entry, sort_keys=True).encode("utf-8"), name=key) + if key not in self._memory: # spend counts once per distinct request, the same on a replay + for name in self.usage: + self.usage[name] += int((entry.get("usage") or {}).get(name, 0)) + self._memory[key] = entry + return str(entry.get("content") or "") + + def usd(self) -> float | None: + """Spend of every response used, priced from its recorded usage: identical on a replay.""" + prices = prices_from(self.parameters) if self.parameters else None + return None if prices is None else round(cost_of(self.usage, prices), 6) + + def cache_digest(self) -> str: + return hashlib.sha256("\n".join(sorted(set(self.keys))).encode()).hexdigest() + + +# --- classify ---------------------------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Decision: + verdict: str + family: str + s: float + votes: Mapping[str, int] + rationale: str + cited_lines: tuple[int, ...] + parse_failed: int + answers: tuple[Answer, ...] + retrieval: str = "signals_only" + neighbours: Mapping[str, Mapping[str, int]] = field(default_factory=dict) + + @property + def blockable(self) -> bool: + """A majority ``unsure`` (parse failures included) can never reach BLOCK.""" + return self.verdict != UNSURE + + +def aggregate(answers: Sequence[Answer]) -> Decision: + votes = {v: sum(a.verdict == v for a in answers) for v in VERDICTS} + top = max(votes.values()) + leaders = [v for v, n in votes.items() if n == top] + verdict = leaders[0] if len(leaders) == 1 else UNSURE + lead = next((a for a in answers if a.verdict == verdict and not a.parse_failed), None) + s = sum(a.p for a in answers) / len(answers) if answers else 0.5 + return Decision( + verdict, lead.family if lead else "none", round(s, 6), votes, lead.rationale if lead else "", lead.cited_lines if lead else (), + sum(a.parse_failed for a in answers), tuple(answers), + ) # fmt: skip + + +def classify(prompt: Prompt, caller: Caller, k: int = 1, *, families: Sequence[str] = ()) -> Decision: + """``k`` samples (the first at temperature 0, the rest at 0.7); one retry at temperature 0 per unparsable answer.""" + answers: list[Answer] = [] + for j in range(k): + temperature = TEMPERATURES[0] if j == 0 else TEMPERATURES[1] + answer = parse_answer(caller.ask(prompt.messages, temperature=temperature, sample=str(j)), prompt.valid_lines, families) + if answer is None: + retry = caller.ask(prompt.messages, temperature=TEMPERATURES[0], sample=f"{j}:retry") + answer = parse_answer(retry, prompt.valid_lines, families) or PARSE_FAILED + answers.append(answer) + return aggregate(answers) + + +@dataclass +class Program: + """``Retrieve -> Classify(k)`` over one memory snapshot.""" + + spec: ProgramSpec + memory: Sequence[Example] + excerpt_text: ExcerptText + vectors: Mapping[str, Sequence[float]] | None = None + families: Sequence[str] = () + + def __post_init__(self) -> None: + self.index = {e.id: e for e in self.memory} + + def prepare(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> tuple[Prompt, Retrieval]: + target = candidate.target + got = retrieve(target, self.memory, fold, vectors=self.vectors, n1=self.spec.n1, k=self.spec.k_ret, lam=self.spec.lam, seed=seed, + priors=self.spec.priors) # fmt: skip + usable = [Demo.from_dict(d) for d in demonstrations_for([d.as_dict() for d in self.spec.demos], self.index, target, fold)] + prompt = render(self.spec, candidate, got.examples, usable, self.index, self.excerpt_text) + assert_boundary([*prompt.example_ids, *prompt.demo_ids], self.index, target, fold, vectors=self.vectors) + return prompt, got + + def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision: + 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) + + +__all__ = [ + "BASELINE_INSTRUCTION", "PARSE_FAILED", "SIGNATURE", "UNSURE", "VERDICTS", "Answer", "Caller", "Candidate", "Decision", "Demo", + "Program", "ProgramSpec", "Prompt", "ReplayMiss", "aggregate", "classify", "label_word", "line_numbers", + "params_digest", "parse_answer", "render", "render_prefix", "render_roles", "render_signals", "request_key", +] # fmt: skip diff --git a/src/openultrasast/model/audit.py b/src/openultrasast/model/audit.py index b9d6abd..dbbf5d0 100644 --- a/src/openultrasast/model/audit.py +++ b/src/openultrasast/model/audit.py @@ -73,9 +73,9 @@ SPEC_OWNED = { "learned-decision-engine Req 8.1: model role classification per file chunk (model_sinks.py moved in); " "its Run template comes with the harvest (task 5)" ), - "learn.retrieve": ( - "learned-decision-engine Req 3.2/3.4: retrieval from the local memory and the evaluation boundary; " - "the program (task 6.4) and `ousast learn compile` (6.5) are its callers" + "learn.program": ( + "learned-decision-engine Req 3.1/3.3: the AI classifier program (signature, Retrieve -> Classify(k), response " + "cache); `ousast learn compile` (task 6.5) and the decide task (task 8) are its callers" ), "plane.workspaces": "ai-service-plane task 4: Workspace manifests per case pin of a population; `ousast plane workspaces`", "plane.runner": "ai-service-plane Req 4.1-4.4: the ax-task-runner entrypoint (PID 1 of the task image); delivers the output directory", diff --git a/src/openultrasast/model/endpoint.py b/src/openultrasast/model/endpoint.py index 3614909..f0c9332 100644 --- a/src/openultrasast/model/endpoint.py +++ b/src/openultrasast/model/endpoint.py @@ -106,13 +106,26 @@ class DeepSeekChatClient: timeout_seconds: int = 60, json_object: bool = False, logprobs: bool = False, + temperature: float | None = None, ) -> ChatResponse: response = self._call( - model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, logprobs=logprobs + model=model, + messages=messages, + tools=tools, + timeout_seconds=timeout_seconds, + json_object=json_object, + logprobs=logprobs, + temperature=temperature, ) if json_object and not (response.content or "").strip(): response = self._call( - model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, logprobs=logprobs + model=model, + messages=messages, + tools=tools, + timeout_seconds=timeout_seconds, + json_object=json_object, + logprobs=logprobs, + temperature=temperature, ) # documented behaviour: JSON mode occasionally returns empty content. One retry, never a loop. return response @@ -131,10 +144,13 @@ class DeepSeekChatClient: timeout_seconds: int, json_object: bool, logprobs: bool, + temperature: float | None = None, ) -> ChatResponse: extra: dict[str, object] = {} if self._disable_thinking: extra["thinking"] = {"type": "disabled"} # otherwise temperature is silently ignored + if temperature is not None: + extra["temperature"] = temperature # overrides the raw call's default 0 if json_object: extra["response_format"] = {"type": "json_object"} if logprobs: diff --git a/src/openultrasast/plane/budget.py b/src/openultrasast/plane/budget.py index 5a95ceb..4f24d78 100644 --- a/src/openultrasast/plane/budget.py +++ b/src/openultrasast/plane/budget.py @@ -121,11 +121,20 @@ class MeteredClient: tools: list[dict[str, object]], timeout_seconds: int = 60, json_object: bool = False, + temperature: float | None = None, + logprobs: bool = False, ) -> ChatResponse: + """One metered call. ``temperature`` and ``logprobs`` pass through only when set, so a wrapped client that + does not take them (a scripted one) is called as before.""" self.check() + extra: dict[str, Any] = {} + if temperature is not None: + extra["temperature"] = temperature + if logprobs: + extra["logprobs"] = True try: response = self.inner.complete( - model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object + model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, **extra ) except Exception as exc: self._absorb_rows() diff --git a/tests/learn_fixtures.py b/tests/learn_fixtures.py index 31f1a6f..b6c54d8 100644 --- a/tests/learn_fixtures.py +++ b/tests/learn_fixtures.py @@ -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)) diff --git a/tests/test_learn_program.py b/tests/test_learn_program.py new file mode 100644 index 0000000..ce32551 --- /dev/null +++ b/tests/test_learn_program.py @@ -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 == []