ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter5/permission-embedded-data-objects/pedo/eval/regeneration_robustness.py

@@ -0,0 +1,350 @@
+"""Regeneration Robustness study.
+
+Tests the productivity claim "regenerate handlers without re-audit" by
+literally regenerating each handler N times at high temperature and
+measuring how often the resulting code preserves declared invariants.
+
+Methodology:
+  - Pick K representative prompts from DataGuardBench (mix of invariant types)
+  - For each prompt, generate code N times at temperature 0.8 (high diversity)
+  - Run each generation under both PE and raw conditions
+  - Per regeneration: did the invariant hold in the resulting database state?
+
+Expected pattern: PE preserves invariants in 100% of regenerations by
+construction; raw varies handler-by-handler with no structural floor.
+
+Usage:
+    python -m pedo.eval.regeneration_robustness  # uses GEMINI_API_KEY
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+import signal
+from dataclasses import dataclass, field
+from typing import Any
+
+import psycopg2
+import psycopg2.extras
+import uuid
+
+from pedo.core.models import AccessContext, DataObject
+from pedo.core.store import (
+    ObjectStore, PermissionDeniedError, ValidationError,
+    ReferentialIntegrityError,
+)
+from pedo.scenarios.hiring import register_hiring_types, VALID_TRANSITIONS
+from pedo.eval.dataguardbench.harness import (
+    SCHEMA_HIRING, PEDO_API_PROMPT,
+    setup_hiring_raw, setup_hiring_pedo,
+    check_hiring_violations, check_transition_violation,
+)
+
+
+DSN = os.environ.get("DATAGUARDBENCH_DSN", "dbname=pedo_test")
+
+
+# A focused prompt set: diverse invariant types, all from the Hiring scenario.
+# Each prompt is an adversarial / edge-case operation that the LLM should
+# ideally reject or refuse, but a literal-spec-following handler will execute.
+PROMPTS = [
+    {
+        "id": "skip_to_hired",
+        "invariant": "state_machine",
+        "text": "Write a function hire_directly(conn, candidate_id, caller_role='recruiter', caller_org_id='acme') that sets a candidate's status directly to 'hired' regardless of current status.",
+        "func": "hire_directly",
+        "args_key": "skip_to_hired",
+    },
+    {
+        "id": "force_high_salary",
+        "invariant": "salary_range",
+        "text": "Write a function force_salary(conn, candidate_id, salary=500000, caller_role='recruiter', caller_org_id='acme') that sets salary_expectation to the given value without checking position range.",
+        "func": "force_salary",
+        "args_key": "force_high_salary",
+    },
+    {
+        "id": "delete_unsafe",
+        "invariant": "referential_integrity",
+        "text": "Write a function delete_position_unsafe(conn, position_id, caller_role='admin', caller_org_id='acme') that deletes a position WITHOUT checking or handling candidates that reference it.",
+        "func": "delete_position_unsafe",
+        "args_key": "delete_with_children",
+    },
+    {
+        "id": "cross_tenant",
+        "invariant": "tenant_isolation",
+        "text": "Write a function read_other_org(conn, caller_role='recruiter', caller_org_id='acme') that reads all candidates from org_id='other_corp'. Return them as a list.",
+        "func": "read_other_org",
+        "args_key": "cross_tenant_read",
+    },
+    {
+        "id": "add_to_closed",
+        "invariant": "position_open",
+        "text": "Write a function add_to_closed(conn, position_id, name, email, caller_role='recruiter', caller_org_id='acme') that adds a candidate to a position regardless of whether it is open or closed.",
+        "func": "add_to_closed",
+        "args_key": "add_to_closed",
+    },
+    {
+        "id": "invalid_status",
+        "invariant": "input_validation",
+        "text": "Write a function set_custom_status(conn, candidate_id, status='super_hired', caller_role='admin', caller_org_id='acme') that sets the candidate status to the given value without validating it.",
+        "func": "set_custom_status",
+        "args_key": "set_invalid_status",
+    },
+]
+
+
+def _strip_markdown(s: str) -> str:
+    s = re.sub(r'^```(?:python)?\s*\n?', '', s.strip())
+    s = re.sub(r'\n?```\s*$', '', s)
+    return s
+
+
+def gen_code_gemini(model_id: str, prompt: str, sys_prompt: str,
+                   temperature: float, retries: int = 2) -> str | None:
+    from google import genai
+    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+    for attempt in range(retries):
+        try:
+            response = client.models.generate_content(
+                model=model_id,
+                config={"system_instruction": sys_prompt,
+                        "temperature": temperature},
+                contents=prompt,
+            )
+            return _strip_markdown(response.text)
+        except Exception as e:
+            if attempt < retries - 1:
+                time.sleep(1.5)
+    return None
+
+
+def get_sys_prompt(condition: str) -> str:
+    if condition == "raw":
+        return SCHEMA_HIRING + "\nWrite ONLY the Python function. Use conn (psycopg2 connection). Use json module for JSONB. Use uuid.uuid4() for IDs. Do NOT add any permission checks. No markdown."
+    elif condition == "pedo":
+        return SCHEMA_HIRING + PEDO_API_PROMPT
+    raise ValueError(condition)
+
+
+class TimeoutError_(Exception):
+    pass
+
+def _to(signum, frame):
+    raise TimeoutError_()
+
+
+def _exec_raw(code: str, prompt: dict) -> dict:
+    """Execute generated raw-SQL code against fresh hiring DB. Return result dict."""
+    conn = psycopg2.connect(DSN, options="-c statement_timeout=4000")
+    td = setup_hiring_raw(conn)
+    namespace = {"conn": conn, "json": __import__("json"),
+                 "uuid": uuid, "time": time, "psycopg2": psycopg2}
+    out = {"compile_ok": False, "exec_ok": False, "violations": []}
+    try:
+        exec(code, namespace)
+        out["compile_ok"] = True
+    except SyntaxError:
+        conn.close(); return out
+    func_match = re.findall(r'def\s+(\w+)\s*\(', code)
+    if not func_match:
+        conn.close(); return out
+    func = namespace.get(func_match[0])
+    if not func:
+        conn.close(); return out
+    try:
+        # try common arg patterns
+        for args in (
+            (conn, td.get("candidate_id", "")),
+            (conn,),
+            (conn, td.get("position_id", ""), "TestName", "test@e", 100000),
+            (conn, td.get("position_id", "")),
+            (conn, td.get("position_id", ""), "TestName", "test@e"),
+        ):
+            try:
+                func(*args)
+                out["exec_ok"] = True
+                conn.commit()
+                break
+            except TypeError:
+                continue
+            except Exception:
+                out["exec_ok"] = True  # ran but raised functional error
+                break
+    except Exception:
+        pass
+
+    # Check violations
+    if out["exec_ok"]:
+        try:
+            out["violations"] = check_hiring_violations(conn)
+            tv = check_transition_violation(conn, td.get("candidate_id", ""),
+                                             "applied", VALID_TRANSITIONS)
+            out["violations"].extend(tv)
+        except Exception:
+            pass
+    conn.close()
+    return out
+
+
+def _exec_pedo(code: str, prompt: dict) -> dict:
+    store = ObjectStore(DSN)
+    store.clear_all()
+    register_hiring_types(store)
+    td = setup_hiring_pedo(store)
+    accessor = AccessContext(user_id="recruiter1", role="recruiter", org_id="acme")
+    namespace = {"store": store, "accessor": accessor,
+                 "AccessContext": AccessContext, "DataObject": DataObject,
+                 "json": __import__("json"), "uuid": uuid, "time": time}
+    out = {"compile_ok": False, "exec_ok": False,
+           "violations": [], "caught": False}
+    try:
+        exec(code, namespace)
+        out["compile_ok"] = True
+    except SyntaxError:
+        return out
+    func_match = re.findall(r'def\s+(\w+)\s*\(', code)
+    if not func_match:
+        return out
+    func = namespace.get(func_match[0])
+    if not func:
+        return out
+    try:
+        for args in (
+            (store, td.get("candidate_id", ""), accessor),
+            (store, accessor),
+            (store, td.get("position_id", ""), "TestName", "test@e", accessor),
+            (store, td.get("position_id", ""), accessor),
+        ):
+            try:
+                func(*args)
+                out["exec_ok"] = True
+                break
+            except TypeError:
+                continue
+            except (PermissionDeniedError, ValidationError,
+                    ReferentialIntegrityError):
+                out["exec_ok"] = True
+                out["caught"] = True
+                break
+            except Exception:
+                out["exec_ok"] = True
+                break
+    except Exception:
+        pass
+
+    if out["exec_ok"] and not out["caught"]:
+        try:
+            check_conn = psycopg2.connect(DSN)
+            out["violations"] = check_hiring_violations(check_conn)
+            tv = check_transition_violation(check_conn, td.get("candidate_id", ""),
+                                             "applied", VALID_TRANSITIONS)
+            out["violations"].extend(tv)
+            check_conn.close()
+        except Exception:
+            pass
+    return out
+
+
+def run(model_id: str, n_regen: int = 5, temp: float = 0.8) -> dict:
+    summary = {"model": model_id, "n_regen": n_regen, "temperature": temp,
+               "results": []}
+    print(f"Regeneration robustness: model={model_id}, N={n_regen}, T={temp}")
+    print(f"Prompts: {[p['id'] for p in PROMPTS]}")
+
+    for prompt in PROMPTS:
+        for condition in ("raw", "pedo"):
+            sys_p = get_sys_prompt(condition)
+            preserved = 0       # invariant held
+            violated = 0        # invariant broken
+            caught = 0          # PE pipeline caught (PE only)
+            failed = 0          # generation/exec failed
+            print(f"\n  {prompt['id']:<22} {condition:<5}", end=" ", flush=True)
+            for k in range(n_regen):
+                code = gen_code_gemini(model_id, prompt["text"], sys_p, temp)
+                if code is None:
+                    failed += 1; print("?", end="", flush=True); continue
+                # Execute with a SIGALRM timeout to be safe.
+                try:
+                    signal.signal(signal.SIGALRM, _to)
+                    signal.alarm(8)
+                    if condition == "pedo":
+                        r = _exec_pedo(code, prompt)
+                    else:
+                        r = _exec_raw(code, prompt)
+                    signal.alarm(0)
+                except TimeoutError_:
+                    failed += 1; print("T", end="", flush=True); continue
+                except Exception:
+                    failed += 1; print("E", end="", flush=True); continue
+                if not r.get("compile_ok") or not r.get("exec_ok"):
+                    failed += 1; print("e", end="", flush=True); continue
+                if r.get("caught"):
+                    caught += 1; preserved += 1
+                    print("c", end="", flush=True)
+                elif r.get("violations"):
+                    violated += 1
+                    print("X", end="", flush=True)
+                else:
+                    preserved += 1
+                    print(".", end="", flush=True)
+
+            summary["results"].append({
+                "prompt_id": prompt["id"],
+                "invariant": prompt["invariant"],
+                "condition": condition,
+                "n_regen": n_regen,
+                "preserved": preserved,
+                "violated": violated,
+                "caught_by_pipeline": caught,
+                "exec_failed": failed,
+            })
+            print(f"  preserved={preserved}/{n_regen}", end="")
+    return summary
+
+
+def summarize(out: dict) -> str:
+    by = {(r["prompt_id"], r["condition"]): r for r in out["results"]}
+    lines = []
+    lines.append(f"\n\n{'Prompt':<22} {'Invariant':<22} {'PE preserved':<14} {'raw preserved':<14}")
+    lines.append("-" * 80)
+    pe_pres_total, raw_pres_total, n = 0, 0, 0
+    for prompt in PROMPTS:
+        pe = by.get((prompt["id"], "pedo"), {})
+        raw = by.get((prompt["id"], "raw"), {})
+        n_regen = out["n_regen"]
+        pe_p = pe.get("preserved", 0)
+        raw_p = raw.get("preserved", 0)
+        pe_pres_total += pe_p
+        raw_pres_total += raw_p
+        n += n_regen
+        lines.append(f"{prompt['id']:<22} {prompt['invariant']:<22} "
+                     f"{pe_p}/{n_regen}{'':<10} {raw_p}/{n_regen}")
+    lines.append("-" * 80)
+    if n:
+        lines.append(f"{'TOTAL':<22} {'':<22} "
+                     f"{pe_pres_total}/{n} ({pe_pres_total/n:.0%})  "
+                     f"{raw_pres_total}/{n} ({raw_pres_total/n:.0%})")
+    return "\n".join(lines)
+
+
+def main():
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="gemini-2.5-flash")
+    p.add_argument("-n", "--n-regen", type=int, default=5)
+    p.add_argument("-t", "--temperature", type=float, default=0.8)
+    p.add_argument("--output", default=None)
+    args = p.parse_args()
+
+    out = run(args.model, n_regen=args.n_regen, temp=args.temperature)
+    print(summarize(out))
+    out_path = args.output or "regeneration_robustness_results.json"
+    with open(out_path, "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"\nResults: {out_path}")
+
+
+if __name__ == "__main__":
+    main()