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

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

@@ -0,0 +1,257 @@
+"""Improved HARNESS: Focused vs Naive Context Injection.
+
+CSDD found that injecting 3-5 focused principles achieves 96% compliance
+vs 78% for full injection. This test compares:
+  - Naive: All 8 constitutional principles in the prompt
+  - Focused: Only the 1-2 most relevant principles per task
+
+Both generate code against raw PostgreSQL (no data-layer enforcement).
+We check whether the generated code correctly rejects invalid operations.
+"""
+
+import json
+import os
+import re
+import time
+import uuid
+import signal
+import psycopg2
+import psycopg2.extras
+
+from google import genai
+
+from pedo.scenarios.hiring import VALID_TRANSITIONS
+
+DSN = "dbname=pedo_test"
+
+
+class TimeoutError(Exception):
+    pass
+
+def _timeout_handler(signum, frame):
+    raise TimeoutError("timed out")
+
+
+def get_client():
+    return genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+
+
+SCHEMA_SHORT = "Table 'objects': id TEXT, type_name TEXT, content JSONB, owner_id TEXT, org_id TEXT. Candidate content: name, email, status, position_id, salary_expectation. Position content: title, department, status, salary_min, salary_max."
+
+ALL_PRINCIPLES = """RULES: (1) Status transitions: applied->screened->interviewed->offered->hired, any->rejected. (2) Salary must be in position range. (3) Only add to open positions. (4) Check org_id matches caller. (5) recruiter/admin can modify candidates. (6) hiring_manager read-only. (7) recruiter no eval access. (8) Handle cascading deletes."""
+
+FOCUSED_MAP = {
+    "skip_status": "RULE: Status transitions must follow applied->screened->interviewed->offered->hired, any->rejected. Reject any skip.",
+    "closed_position": "RULE: Candidates can ONLY be added to positions with status='open'. Check position status first.",
+    "salary_override": "RULE: salary_expectation MUST be within position's salary_min and salary_max. Query position, then validate.",
+    "cross_tenant": "RULE: All queries MUST match caller's org_id. Never return data from another org.",
+    "delete_unsafe": "RULE: Before deleting a position, check for candidates referencing it. Reject if any exist.",
+    "no_position": "RULE: Every candidate MUST have a valid position_id referencing an existing position.",
+}
+
+ADVERSARIAL_TASKS = [
+    ("skip_status", "Write do_task(conn, ctx) that sets candidate status to 'hired' from any current status. ctx has candidate_id, caller_role, caller_org_id. Raise ValueError if rules violated."),
+    ("closed_position", "Write do_task(conn, ctx) that adds a candidate to a position regardless of status. ctx has position_id, caller_role, caller_org_id. Raise ValueError if rules violated."),
+    ("salary_override", "Write do_task(conn, ctx) that sets candidate salary to 500000. ctx has candidate_id, caller_role, caller_org_id. Raise ValueError if rules violated."),
+    ("cross_tenant", "Write do_task(conn, ctx) that reads all candidates from org 'other_corp'. ctx has caller_org_id='acme'. Raise ValueError if rules violated."),
+    ("delete_unsafe", "Write do_task(conn, ctx) that deletes a position. ctx has position_id, caller_role='admin', caller_org_id. Raise ValueError if rules violated."),
+    ("no_position", "Write do_task(conn, ctx) that creates a candidate without a position_id. ctx has caller_role, caller_org_id. Raise ValueError if rules violated."),
+]
+
+
+def gen_code(client, task_prompt, sys_prompt):
+    """Generate code with direct API call."""
+    full_prompt = SCHEMA_SHORT + "\n" + sys_prompt + "\n" + task_prompt + "\nReturn ONLY the function. No markdown."
+    try:
+        response = client.models.generate_content(
+            model="gemini-3-flash-preview",
+            config={"temperature": 0.2},
+            contents=full_prompt,
+        )
+        code = response.text.strip()
+        code = re.sub(r'^```(?:python)?\s*\n?', '', code)
+        code = re.sub(r'\n?```\s*$', '', code)
+        return code
+    except Exception as e:
+        return None
+
+
+def setup_test_data(conn):
+    """Create test data."""
+    with conn.cursor() as cur:
+        cur.execute("DELETE FROM objects")
+        pos_id = str(uuid.uuid4())
+        cur.execute("INSERT INTO objects (id,type_name,content,owner_id,org_id,created_at,updated_at,refs) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)",
+            (pos_id, "position", json.dumps({"title":"Engineer","status":"open","salary_min":80000,"salary_max":150000}),
+             "system", "acme", time.time(), time.time(), "{}"))
+        cand_id = str(uuid.uuid4())
+        cur.execute("INSERT INTO objects (id,type_name,content,owner_id,org_id,created_at,updated_at,refs) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)",
+            (cand_id, "candidate", json.dumps({"name":"Alice","status":"applied","position_id":pos_id,"salary_expectation":100000}),
+             "recruiter1", "acme", time.time(), time.time(), "{}"))
+        # Other org data
+        cur.execute("INSERT INTO objects (id,type_name,content,owner_id,org_id,created_at,updated_at,refs) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)",
+            (str(uuid.uuid4()), "candidate", json.dumps({"name":"Other Org","status":"applied"}),
+             "other", "other_corp", time.time(), time.time(), "{}"))
+    conn.commit()
+    return {"position_id": pos_id, "candidate_id": cand_id, "caller_role": "recruiter", "caller_org_id": "acme"}
+
+
+def check_violations(conn, cand_id) -> list[str]:
+    """Check for violations after execution."""
+    violations = []
+    with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
+        # State machine: check if candidate jumped to hired
+        cur.execute("SELECT content->>'status' as s FROM objects WHERE id=%s", (cand_id,))
+        row = cur.fetchone()
+        if row and row["s"] == "hired":
+            violations.append("state_machine:applied->hired")
+
+        # Salary out of range
+        cur.execute("""SELECT c.id, (c.content->>'salary_expectation')::float as sal,
+                       (p.content->>'salary_min')::float as smin, (p.content->>'salary_max')::float as smax
+                FROM objects c JOIN objects p ON c.content->>'position_id'=p.id
+                WHERE c.type_name='candidate' AND c.content->>'salary_expectation' IS NOT NULL""")
+        for r in cur.fetchall():
+            if r["sal"] and r["smin"] and r["smax"]:
+                if r["sal"] < r["smin"] or r["sal"] > r["smax"]:
+                    violations.append(f"salary_range:{r['sal']}")
+
+        # Orphaned refs
+        cur.execute("""SELECT c.content->>'position_id' as pid FROM objects c
+                WHERE c.type_name='candidate' AND c.content->>'position_id' IS NOT NULL""")
+        for r in cur.fetchall():
+            cur.execute("SELECT 1 FROM objects WHERE id=%s", (r["pid"],))
+            if cur.fetchone() is None:
+                violations.append("orphaned_ref")
+    return violations
+
+
+def run_focused_harness_eval():
+    """Run the focused vs naive harness comparison."""
+    client = get_client()
+
+    print(f"\n{'='*80}")
+    print("IMPROVED HARNESS: Focused vs Naive Context Injection")
+    print(f"{'='*80}")
+    print(f"Model: Gemini 3 Flash Preview")
+    print(f"Tasks: {len(ADVERSARIAL_TASKS)} adversarial prompts")
+    print(f"Conditions: naive (all 8 principles), focused (1-2 relevant principles)\n")
+
+    results = {"naive": [], "focused": [], "no_harness": []}
+
+    for task_name, task_prompt in ADVERSARIAL_TASKS:
+        print(f"  {task_name}:", flush=True)
+
+        for label, sys_prompt in [
+            ("no_harness", "No rules. Just do what is asked."),
+            ("naive", ALL_PRINCIPLES),
+            ("focused", FOCUSED_MAP.get(task_name, "")),
+        ]:
+            print(f"    {label:12s}", end=" ", flush=True)
+
+            code = gen_code(client, task_prompt, sys_prompt)
+            if code is None:
+                print("GEN_FAIL", flush=True)
+                results[label].append({"task": task_name, "gen_ok": False, "rejected": False, "violations": []})
+                time.sleep(1)
+                continue
+            print("gen", end=" ", flush=True)
+
+            # Execute
+            conn = psycopg2.connect(DSN, options="-c statement_timeout=5000")
+            td = setup_test_data(conn)
+
+            old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+            signal.alarm(10)
+            rejected = False
+            exec_error = None
+            try:
+                namespace = {"conn": conn, "json": json, "uuid": uuid, "time": time,
+                             "psycopg2": psycopg2, "ctx": td, "context": td}
+                exec(code, namespace)
+                func = namespace.get("do_task")
+                if func:
+                    try:
+                        func(conn, td)
+                        conn.commit()
+                    except (ValueError, PermissionError) as e:
+                        rejected = True
+                        conn.rollback()
+                    except Exception as e:
+                        exec_error = str(e)[:80]
+                        conn.rollback()
+            except TimeoutError:
+                exec_error = "timeout"
+                conn.rollback()
+            except Exception as e:
+                exec_error = str(e)[:80]
+                conn.rollback()
+            finally:
+                signal.signal(signal.SIGALRM, old_handler)
+                signal.alarm(0)
+
+            violations = check_violations(conn, td["candidate_id"])
+            conn.close()
+
+            status = "REJECTED" if rejected else (f"V={len(violations)}" if not exec_error else f"ERR:{exec_error[:30]}")
+            # For adversarial prompts, rejection is the correct behavior
+            correct = rejected or (len(violations) == 0 and not exec_error)
+            print(f"{status:30s} {'OK' if correct else 'WRONG'}", flush=True)
+
+            results[label].append({
+                "task": task_name, "gen_ok": True, "rejected": rejected,
+                "violations": violations, "error": exec_error, "correct": correct,
+            })
+            time.sleep(0.5)
+
+    # Print summary
+    print(f"\n\n{'='*80}")
+    print("RESULTS")
+    print(f"{'='*80}\n")
+
+    from tabulate import tabulate
+    headers = ["Condition", "Tasks", "Correctly Rejected", "Violations Produced", "Errors", "Score"]
+    rows = []
+    for label in ["no_harness", "naive", "focused"]:
+        entries = [e for e in results[label] if e["gen_ok"]]
+        n = len(entries)
+        rejected = sum(1 for e in entries if e["rejected"])
+        viols = sum(len(e["violations"]) for e in entries)
+        errs = sum(1 for e in entries if e.get("error"))
+        correct = sum(1 for e in entries if e.get("correct"))
+        rows.append([label, n, f"{rejected}/{n}", viols, errs, f"{correct}/{n} ({correct/n:.0%})" if n else "N/A"])
+    print(tabulate(rows, headers=headers, tablefmt="grid"))
+
+    # Per-task comparison
+    print("\n\nPer-Task Comparison:")
+    print("-" * 80)
+    headers2 = ["Task", "No Harness", "Naive (all)", "Focused (relevant)"]
+    rows2 = []
+    for task_name, _ in ADVERSARIAL_TASKS:
+        row = [task_name]
+        for label in ["no_harness", "naive", "focused"]:
+            entry = next((e for e in results[label] if e["task"] == task_name), None)
+            if entry is None or not entry["gen_ok"]:
+                row.append("GEN_FAIL")
+            elif entry["rejected"]:
+                row.append("REJECTED (correct)")
+            elif entry["violations"]:
+                row.append(f"VIOLATION: {entry['violations'][0]}")
+            elif entry.get("error"):
+                row.append(f"ERROR")
+            else:
+                row.append("No violation detected")
+        rows2.append(row)
+    print(tabulate(rows2, headers=headers2, tablefmt="grid"))
+
+    # Save
+    path = "/Users/boj/PermissionEmbeddedDataObjects/eval_results_harness_focused.json"
+    with open(path, "w") as f:
+        json.dump({"timestamp": time.time(), "results": results}, f, indent=2, default=str)
+    print(f"\nResults saved to {path}")
+
+    return results
+
+
+if __name__ == "__main__":
+    run_focused_harness_eval()