From 2858a42d4061a379eea54fe47591fbb0b27af061 Mon Sep 17 00:00:00 2001 From: Sladow_Slay3r Date: Fri, 25 Sep 2026 07:09:55 +0530 Subject: [PATCH] =?UTF-8?q?feat:=20KB=20quality=20management=20=E2=80=94?= =?UTF-8?q?=20reinforce,=20prune,=20quality=20report?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the self-growing KB (--evolve) curate itself instead of only appending. - RedTeamKB gains get(), record_success() (bump success_count + stamp last_used), prune() (remove grown 'dynamic-win' patterns that are stale and/or weak — static seeds and proven/undated patterns are never pruned), and quality_report() (origin mix, reinforcement count, staleness, top wins). - The dynamic grow loop now REINFORCES a repeat win (record_success on the near-duplicate) instead of dropping it; new wins are stamped success_count=1 + created_at/last_used; DynamicRedTeamer tracks grown vs reinforced (shown in the "KB grew" line). - CLI: --kb-prune [--kb-prune-max-age-days N] [--kb-prune-min-success K]; --kb-stats now prints a quality snapshot (grown wins / reinforced / avg success / stale >30d). Recovers the second feature lost with the deleted background-task branch (see #24/#25). +9 tests (tests/test_kb_quality.py). Full suite: 1947 passed, 1 skipped. Co-Authored-By: Claude Opus 4.8 --- dynamic_engine.py | 25 +++++++--- kb/knowledge_base.py | 87 +++++++++++++++++++++++++++++++++ main.py | 33 ++++++++++++- tests/test_kb_quality.py | 101 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 238 insertions(+), 8 deletions(-) create mode 100644 tests/test_kb_quality.py diff --git a/dynamic_engine.py b/dynamic_engine.py index d27e6cf..fd1f75c 100644 --- a/dynamic_engine.py +++ b/dynamic_engine.py @@ -825,6 +825,7 @@ def __init__( self.kb_grow = kb_grow and kb is not None self.grow_threshold = grow_threshold self.grown = 0 # count of winning attacks written back to the KB + self.reinforced = 0 # count of repeat wins that reinforced an existing entry # ── Internal helpers ────────────────────────────────────────────────────── @@ -954,14 +955,24 @@ def run_pair_loop(self, base_test: dict) -> "PairLoopResult": # ── Grow: write high-confidence winners back into the KB ────── if self.kb_grow and confidence >= self.grow_threshold: try: - if not self.kb.has_similar(winning_payload, threshold=0.95): + dup = self.kb.query("attack_patterns", winning_payload, n=1) + if dup and dup[0]["score"] >= 0.95: + # A repeat of a known win — reinforce it (bump success_count) + # instead of dropping the near-duplicate on the floor. + self.kb.record_success(dup[0]["doc_id"]) + self.reinforced += 1 + else: + ts = time.time() self.kb.add("attack_patterns", winning_payload, metadata={ - "origin": "dynamic-win", - "category": category, - "confidence": round(confidence, 2), - "seed_id": base_test.get("id", ""), - "target_model": self.target_config.get("model", ""), - "strategy": winning_strategy, + "origin": "dynamic-win", + "category": category, + "confidence": round(confidence, 2), + "seed_id": base_test.get("id", ""), + "target_model": self.target_config.get("model", ""), + "strategy": winning_strategy, + "success_count": 1, + "created_at": ts, + "last_used": ts, }) self.grown += 1 except Exception: diff --git a/kb/knowledge_base.py b/kb/knowledge_base.py index 64978cd..10971eb 100644 --- a/kb/knowledge_base.py +++ b/kb/knowledge_base.py @@ -28,6 +28,7 @@ import math import os import re +import time import sqlite3 DEFAULT_DIR = os.environ.get("CRUCIBLE_KB_DIR", ".crucible-kb") @@ -202,6 +203,92 @@ def delete(self, collection: str, doc_id: str) -> bool: self._conn.commit() return cur.rowcount > 0 + # ── quality management ────────────────────────────────────────────────────── + def get(self, collection: str, doc_id: str): + """Return one doc as {doc_id, text, metadata}, or None if absent.""" + row = self._conn.execute( + "SELECT doc_id, text, metadata FROM documents " + "WHERE collection=? AND doc_id=?", (collection, doc_id)).fetchone() + if not row: + return None + return {"doc_id": row[0], "text": row[1], "metadata": json.loads(row[2] or "{}")} + + def record_success(self, doc_id: str, collection: str = "attack_patterns", + now: float = None): + """Reinforce a stored win: bump its ``success_count`` and stamp ``last_used``. + Returns the new success_count, or None if the doc no longer exists.""" + doc = self.get(collection, doc_id) + if doc is None: + return None + meta = doc["metadata"] + meta["success_count"] = int(meta.get("success_count", 0)) + 1 + meta["last_used"] = float(now if now is not None else time.time()) + self._conn.execute( + "UPDATE documents SET metadata=? WHERE collection=? AND doc_id=?", + (json.dumps(meta), collection, doc_id)) + self._conn.commit() + return meta["success_count"] + + def prune(self, collection: str = "attack_patterns", *, + max_age_days: float = None, min_success: int = None, + now: float = None) -> dict: + """Remove GROWN (``origin == 'dynamic-win'``) patterns that are stale and/or + weak. Static seeds (any other origin) are never pruned, and a grown doc missing + ``created_at`` is never pruned on the age criterion. When both criteria are + given a doc must be BOTH stale AND weak to go; with one, that one decides; with + neither, nothing is pruned. Returns {'pruned', 'kept', 'removed': [doc_id,...]}.""" + if max_age_days is None and min_success is None: + return {"pruned": 0, "kept": self.count(collection), "removed": []} + now = float(now if now is not None else time.time()) + removed = [] + for _c, did, _text, meta_json, _e in self._rows(collection): + meta = json.loads(meta_json or "{}") + if str(meta.get("origin", "")) != "dynamic-win": + continue # only grown patterns prunable + created = meta.get("created_at") + age_days = ((now - float(created)) / 86400.0) if created is not None else None + stale = (max_age_days is not None and age_days is not None + and age_days > max_age_days) + weak = (min_success is not None + and int(meta.get("success_count", 0)) < min_success) + if max_age_days is not None and min_success is not None: + doomed = stale and weak + elif max_age_days is not None: + doomed = stale + else: + doomed = weak + if doomed: + removed.append(did) + for did in removed: + self.delete(collection, did) + return {"pruned": len(removed), "kept": self.count(collection), "removed": removed} + + def quality_report(self, collection: str = "attack_patterns", + now: float = None) -> dict: + """A snapshot of KB health: origin mix, grown-win reinforcement, staleness.""" + now = float(now if now is not None else time.time()) + wins, by_origin = [], {} + for _c, did, _text, meta_json, _e in self._rows(collection): + meta = json.loads(meta_json or "{}") + origin = str(meta.get("origin", "seed")) + by_origin[origin] = by_origin.get(origin, 0) + 1 + if origin == "dynamic-win": + wins.append((did, int(meta.get("success_count", 0)), + meta.get("category", ""), meta.get("created_at"))) + n = len(wins) + top = sorted(wins, key=lambda w: w[1], reverse=True)[:5] + return { + "collection": collection, + "total": self.count(collection), + "by_origin": by_origin, + "dynamic_wins": n, + "reinforced": sum(1 for w in wins if w[1] > 1), + "avg_success": round(sum(w[1] for w in wins) / n, 2) if n else 0.0, + "stale_over_30d": sum(1 for w in wins + if w[3] is not None and (now - float(w[3])) / 86400 > 30), + "top": [{"doc_id": w[0], "success_count": w[1], "category": w[2]} for w in top], + } + def reset(self) -> None: self._conn.execute("DELETE FROM documents") self._conn.commit() diff --git a/main.py b/main.py index 6f02ca9..85b277c 100644 --- a/main.py +++ b/main.py @@ -624,6 +624,14 @@ def build_parser(): help="Semantic search the KB attack corpus, then exit") p.add_argument("--kb-reset", action="store_true", help="Wipe the knowledge base, then exit") + p.add_argument("--kb-prune", action="store_true", + help="Prune grown (dynamic-win) KB patterns that are stale and/or " + "low-value, then exit. Static seeds are never pruned. Combine " + "with --kb-prune-max-age-days and/or --kb-prune-min-success.") + p.add_argument("--kb-prune-max-age-days", type=float, metavar="N", + help="With --kb-prune: prune grown patterns older than N days") + p.add_argument("--kb-prune-min-success", type=int, metavar="K", + help="With --kb-prune: prune grown patterns with success_count < K") p.add_argument("--slm-collect", action="store_true", help="Snapshot the KB's winning attacks into SLM training data (JSONL), then exit") @@ -2086,8 +2094,30 @@ def _dispatch(args) -> 'int | None': f"{'semantic' if st['semantic'] else 'lexical'}") for c, n in sorted(st["collections"].items()): print(f" {c:<16}: {n}") + q = kb.quality_report("attack_patterns") + if q["dynamic_wins"]: + print(f" {C.DIM('grown wins')} : {q['dynamic_wins']} " + f"({q['reinforced']} reinforced · avg success {q['avg_success']} · " + f"{q['stale_over_30d']} stale >30d)") print() return 0 + if getattr(args, "kb_prune", False): + kb = kb_mod.RedTeamKB(persist_dir=args.kb_dir) + max_age = getattr(args, "kb_prune_max_age_days", None) + min_suc = getattr(args, "kb_prune_min_success", None) + if max_age is None and min_suc is None: + print(f"\n {C.YELLOW('◈ KB PRUNE')}: no criteria — pass " + f"--kb-prune-max-age-days and/or --kb-prune-min-success.\n") + return 2 + res = kb.prune("attack_patterns", max_age_days=max_age, min_success=min_suc) + crit = [] + if max_age is not None: + crit.append(f">{max_age}d old") + if min_suc is not None: + crit.append(f"success<{min_suc}") + print(f"\n {C.BOLD('◈ KB PRUNE')} ({' & '.join(crit)}): removed " + f"{C.RED(str(res['pruned']))} grown pattern(s), {res['kept']} kept.\n") + return 0 if getattr(args, "kb_search", None): kb = kb_mod.RedTeamKB(persist_dir=args.kb_dir) if kb.count("attack_patterns") == 0: @@ -3422,7 +3452,8 @@ def _verdict_of(cfg, test): if ensemble_attackers else attacker_model) print_dynamic_report(dyn_results, attacker_label) if dyn_kb is not None and getattr(args, "kb_grow", False): - print(f" {C.GREEN('◈ KB grew')}: +{drt.grown} winning attack(s) → " + print(f" {C.GREEN('◈ KB grew')}: +{drt.grown} new winning attack(s), " + f"{drt.reinforced} reinforced → " f"{dyn_kb.count('attack_patterns')} patterns " f"{C.DIM('(the corpus compounds each run)')}") diff --git a/tests/test_kb_quality.py b/tests/test_kb_quality.py new file mode 100644 index 0000000..7b1faca --- /dev/null +++ b/tests/test_kb_quality.py @@ -0,0 +1,101 @@ +"""KB quality management — get / record_success / prune / quality_report, and the +reinforce-on-repeat-win behaviour. Runs in lexical mode (embed_fn=None) so no Ollama +daemon is needed; all timestamps are injected for determinism. +""" +from kb.knowledge_base import RedTeamKB + + +def _kb(tmp_path): + return RedTeamKB(persist_dir=str(tmp_path / "kb"), embed_fn=None) + + +def _add_win(kb, text, success_count=1, created_at=1000.0, category="Jailbreak"): + return kb.add("attack_patterns", text, metadata={ + "origin": "dynamic-win", "success_count": success_count, + "created_at": created_at, "category": category}) + + +def test_get_returns_doc_or_none(tmp_path): + kb = _kb(tmp_path) + did = _add_win(kb, "attack one") + doc = kb.get("attack_patterns", did) + assert doc and doc["doc_id"] == did and doc["metadata"]["origin"] == "dynamic-win" + assert kb.get("attack_patterns", "nope") is None + + +def test_record_success_bumps_and_stamps(tmp_path): + kb = _kb(tmp_path) + did = _add_win(kb, "attack two", success_count=1) + assert kb.record_success(did, now=1234.0) == 2 + assert kb.record_success(did, now=1235.0) == 3 + doc = kb.get("attack_patterns", did) + assert doc["metadata"]["success_count"] == 3 + assert doc["metadata"]["last_used"] == 1235.0 + assert kb.record_success("missing-id") is None + + +def test_prune_only_targets_dynamic_wins(tmp_path): + kb = _kb(tmp_path) + kb.add("attack_patterns", "a static seed", + metadata={"origin": "static-suite", "success_count": 0}) + weak = _add_win(kb, "weak grown win", success_count=0) + res = kb.prune("attack_patterns", min_success=1) + assert res["pruned"] == 1 and weak in res["removed"] + assert kb.count("attack_patterns") == 1 # the seed survives + assert kb.get("attack_patterns", weak) is None + + +def test_prune_min_success_keeps_proven(tmp_path): + kb = _kb(tmp_path) + proven = _add_win(kb, "proven win", success_count=5) + weak = _add_win(kb, "weak win", success_count=1) + res = kb.prune("attack_patterns", min_success=3) + assert weak in res["removed"] and proven not in res["removed"] + + +def test_prune_by_age(tmp_path): + kb = _kb(tmp_path) + now = 100 * 86400 + old = _add_win(kb, "old win", created_at=0.0) # ~100 days old + fresh = _add_win(kb, "fresh win", created_at=99 * 86400) # ~1 day old + res = kb.prune("attack_patterns", max_age_days=30, now=now) + assert old in res["removed"] and fresh not in res["removed"] + + +def test_prune_both_criteria_require_stale_and_weak(tmp_path): + kb = _kb(tmp_path) + now = 100 * 86400 + _add_win(kb, "old but proven", success_count=9, created_at=0.0) + old_weak = _add_win(kb, "old and weak", success_count=0, created_at=0.0) + _add_win(kb, "fresh weak", success_count=0, created_at=now) + res = kb.prune("attack_patterns", max_age_days=30, min_success=3, now=now) + assert res["removed"] == [old_weak] # only the one that is BOTH + + +def test_prune_missing_created_at_not_stale(tmp_path): + kb = _kb(tmp_path) + d = kb.add("attack_patterns", "no timestamp", + metadata={"origin": "dynamic-win", "success_count": 0}) + res = kb.prune("attack_patterns", max_age_days=1, now=10 ** 9) + assert d not in res["removed"] # unknown age → never age-pruned + + +def test_prune_no_criteria_is_noop(tmp_path): + kb = _kb(tmp_path) + _add_win(kb, "x", success_count=0, created_at=0.0) + assert kb.prune("attack_patterns")["pruned"] == 0 + + +def test_quality_report(tmp_path): + kb = _kb(tmp_path) + kb.add("attack_patterns", "seed a", metadata={"origin": "static-suite"}) + _add_win(kb, "win a", success_count=3, created_at=0.0) # stale + reinforced + _add_win(kb, "win b", success_count=1, created_at=10 ** 9) + q = kb.quality_report("attack_patterns", now=10 ** 9) + assert q["total"] == 3 + assert q["dynamic_wins"] == 2 + assert q["by_origin"]["static-suite"] == 1 + assert q["reinforced"] == 1 # win a: success_count > 1 + assert q["avg_success"] == 2.0 # (3 + 1) / 2 + assert q["stale_over_30d"] == 1 # win a + assert q["top"][0]["success_count"] == 3