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
2026-10-01 10:13:31 +02:00
co-authored by Claude Opus 5.5
parent f5180a14ec
commit ea61dc1bb3
8 changed files with 764 additions and 13 deletions
+1 -1
View File
@@ -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/<sha>.json` cache
@@ -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"
]
+473
View File
@@ -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/<sha>.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": <family or\n'
'"none">, "confidence": <probability in [0, 1] that your verdict is right>, "rationale": <at most 3 sentences>,\n'
'"cited_lines": [<line numbers of the code that support the verdict>]}. 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; ``<instrument>: 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/<sha>.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
+3 -3
View File
@@ -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",
+18 -2
View File
@@ -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:
+10 -1
View File
@@ -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()
+47 -1
View File
@@ -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))
+190
View File
@@ -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 == []