Run the PR risk classifier with pi instead of Claude Code

.github/workflows/pr-risk.yml

@@ -5,8 +5,8 @@ name: PR Risk
 #
 # A PR is skipped, with no comment, when it is a draft, comes from a fork, was opened by someone
 # without write access (bots included), changes more than 50 lines, or touches a deny-listed
-# path (defaults in scripts/pr_risk.py plus the caller's `deny-paths`). Otherwise Claude
-# classifies the diff and may deny it on its own judgment.
+# path (defaults in scripts/pr_risk.py plus the caller's `deny-paths`). Otherwise a model, run
+# through pi, classifies the diff and may deny it on its own judgment.
 # If a PR that already has a comment becomes ineligible on a later push, the comment is
 # rewritten so it never shows a stale "yes".
 #
@@ -27,15 +27,23 @@ on:
         type: string
         required: false
         default: main
+      model:
+        description: Pi model as provider/id. Defaults to the PI_PREFERRED_MODEL variable.
+        type: string
+        required: false
+        default: ${{ vars.PI_PREFERRED_MODEL }}
     # Named rather than inherited: `secrets: inherit` would hand this cross-repo workflow
     # every secret the caller holds.
     secrets:
       ADMIN_APP_ID:
         required: true
       ADMIN_APP_PRIVATE_KEY:
         required: true
+      # Pass the key for the model's provider. Both are optional so either provider works.
       ANTHROPIC_API_KEY:
-        required: true
+        required: false
+      OPENAI_API_KEY:
+        required: false
 
 jobs:
   classify:
@@ -69,27 +77,39 @@ jobs:
           DENY_PATHS: ${{ inputs.deny-paths }}
         run: python3 /tmp/pr_risk.py gather --repo "$REPO" --number "$NUMBER" --dir "$WORKDIR"
 
-      - name: Install Claude Code CLI
+      - name: Setup Node
         if: steps.gather.outputs.eligible == 'true'
+        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
+        with:
+          node-version: '22'
+
+      - name: Install Pi
+        if: steps.gather.outputs.eligible == 'true'
+        env:
+          PI_TELEMETRY: '0'
         run: |
-          curl -fsSL https://claude.ai/install.sh | bash
-          echo "$HOME/.local/bin" >> "$GITHUB_PATH"
+          npm install -g --ignore-scripts @earendil-works/pi-coding-agent
+          pi --version
 
       - name: Classify
         if: steps.gather.outputs.eligible == 'true'
         env:
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
-          MODEL: ${{ vars.CLAUDE_CODE_PREFERRED_MODEL }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          MODEL: ${{ inputs.model }}
+          PI_SKIP_VERSION_CHECK: '1'
+          PI_TELEMETRY: '0'
         run: |
           set -euo pipefail
-          # An empty directory, so no CLAUDE.md or project settings are picked up.
+          : "${MODEL:?Set the PI_PREFERRED_MODEL organization variable or the model input}"
+          # An empty directory, and no context files or project settings, so nothing but the
+          # prompt and the PR data reaches the model.
           mkdir -p "$WORKDIR/cwd" && cd "$WORKDIR/cwd"
-          claude -p ${MODEL:+--model "$MODEL"} \
-            --tools "" --strict-mcp-config \
-            --system-prompt-file /tmp/pr-risk-prompt.md \
-            --json-schema "$(python3 /tmp/pr_risk.py schema)" \
-            --output-format json --max-budget-usd 1 \
-            < "$WORKDIR/input.md" > "$WORKDIR/result.json"
+          pi --mode json --no-session --no-approve --model "$MODEL" --thinking high \
+            --no-tools --no-mcp --no-extensions --no-skills --no-prompt-templates --no-context-files \
+            --system-prompt /tmp/pr-risk-prompt.md \
+            < "$WORKDIR/input.md" > "$WORKDIR/pi-events.jsonl"
+          python3 /tmp/pr_risk.py verdict --dir "$WORKDIR"
 
       - name: Generate app token
         id: app-token

.github/workflows/pr-risk/classify-prompt.md

@@ -25,8 +25,10 @@ Set `deny` to true, whatever the tier, if the change touches any of these, even
 
 ## Output
 
+Reply with one JSON object and nothing else, with exactly these keys:
+
 - `tier`: one of the tiers above.
 - `deny`: as defined above.
 - `category`: a short kebab-case label for the kind of change, such as `flag-allowlist`, `flag-rollout`, `docs`, `test-only`, `config`, `copy`, `bug-fix` or `refactor`.
-- `confidence`: how sure you are of the tier. Use `low` when the diff depends on context you were not given.
+- `confidence`: `high`, `medium` or `low`: how sure you are of the tier. Use `low` when the diff depends on context you were not given.
 - `reason`: one or two plain sentences, at most 40 words, naming what the change does and why that is or is not safe to approve without review. If `deny` is true, name the denied area. No preamble and no restating the tier.

README.md

@@ -74,8 +74,8 @@ approves.
 These PRs are skipped with no comment: drafts, fork PRs, PRs whose author lacks write
 access to the repo (including bots), PRs over 50 changed lines, and PRs that touch a
 deny-listed path. The defaults live in `scripts/pr_risk.py`, and callers can add more
-through `deny-paths`. Claude classifies everything else and can deny a PR on its own
-judgment. The comment stores its verdict in a hidden `pr-risk:v1` JSON marker for later analysis.
+through `deny-paths`. A model run through pi classifies everything else and can deny a PR on its
+own judgment. The comment stores its verdict in a hidden `pr-risk:v1` JSON marker for later analysis.
 
 ```yaml
 # In your repo's .github/workflows/pr-risk.yml
@@ -96,8 +96,12 @@ jobs:
       ADMIN_APP_ID: ${{ secrets.ADMIN_APP_ID }}
       ADMIN_APP_PRIVATE_KEY: ${{ secrets.ADMIN_APP_PRIVATE_KEY }}
       ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
 ```
 
+The model defaults to the `PI_PREFERRED_MODEL` variable and can be overridden with the `model`
+input. Pass the API key for that model's provider; both keys are optional secrets.
+
 ## Enrollment
 
 Enroll repositories that process, store, transmit, or can materially affect production customer data or production operations. Examples include application services, data pipelines, infrastructure-as-code, deployment tooling, internal admin tools, and customer-facing dashboards.

scripts/pr_risk.py

@@ -2,6 +2,7 @@
 """Shadow-mode risk classifier for small pull requests, run by pr-risk.yml.
 
 `gather` decides whether a PR is eligible and builds the model input using a read-only token.
+`verdict` turns pi's JSON event stream into result.json, the verdict envelope `comment` reads.
 `comment` turns the model's verdict into one sticky PR comment using the kernel-internal token.
 Nothing here approves anything: the comment only records whether the PR would have been
 auto-approved, so the call can be checked against what reviewers actually did.
@@ -318,6 +319,56 @@ def gather(gh: GitHub, number: int, extra_deny: str, out: Path) -> dict[str, Any
     return gate
 
 
+def schema_error(value: Any) -> str | None:
+    if not isinstance(value, dict):
+        return "reply is not a JSON object"
+    props = VERDICT_SCHEMA["properties"]
+    extra = sorted(set(value) - set(props))
+    if extra:
+        return f"unexpected keys {extra}"
+    for key in VERDICT_SCHEMA["required"]:
+        if key not in value:
+            return f"missing {key}"
+    for key, spec in props.items():
+        expected = {"string": str, "boolean": bool}[spec["type"]]
+        if not isinstance(value[key], expected):
+            return f"{key} must be a {spec['type']}"
+        if "enum" in spec and value[key] not in spec["enum"]:
+            return f"{key} must be one of {spec['enum']}"
+    return None
+
+
+def verdict_from_events(events: str) -> dict[str, Any]:
+    """Build the result envelope `parse_verdict` reads from `pi --mode json` output.
+
+    pi has no schema-constrained output, so the model replies with one JSON object and it is
+    checked here against VERDICT_SCHEMA. Anything else is reported as a classifier error.
+    """
+    reply = None
+    for line in events.split("\n"):
+        if not line.strip():
+            continue
+        event = json.loads(line)
+        message = event.get("message") or {}
+        if event.get("type") == "message_end" and message.get("role") == "assistant":
+            reply = message
+    if reply is None:
+        return {"is_error": True, "result": "pi produced no assistant reply"}
+    if reply.get("stopReason") != "stop":
+        detail = reply.get("errorMessage") or "no error message"
+        return {"is_error": True, "result": f"pi stopped with {reply.get('stopReason')}: {detail}"}
+    text = "".join(c["text"] for c in reply["content"] if c.get("type") == "text").strip()
+    fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", text, re.S)
+    try:
+        value = json.loads(fenced.group(1) if fenced else text)
+    except json.JSONDecodeError:
+        value = None
+    error = schema_error(value) if value is not None else "reply is not JSON"
+    if error:
+        return {"is_error": True, "result": f"classifier reply does not match the verdict schema: {error}"}
+    return {"is_error": False, "result": text, "structured_output": value}
+
+
 def parse_verdict(result: dict[str, Any]) -> dict[str, Any]:
     if result.get("is_error"):
         raise ValueError(f"classifier errored: {result.get('result')}")
@@ -431,6 +482,7 @@ def main() -> int:
     parser = argparse.ArgumentParser()
     sub = parser.add_subparsers(dest="cmd", required=True)
     sub.add_parser("schema")
+    sub.add_parser("verdict").add_argument("--dir", type=Path, required=True)
     for name in ("gather", "comment"):
         p = sub.add_parser(name)
         p.add_argument("--repo", required=True)
@@ -441,6 +493,10 @@ def main() -> int:
     if args.cmd == "schema":
         print(json.dumps(VERDICT_SCHEMA))
         return 0
+    if args.cmd == "verdict":
+        result = verdict_from_events((args.dir / "pi-events.jsonl").read_text())
+        (args.dir / "result.json").write_text(json.dumps(result, indent=2))
+        return 0
     gh = GitHub(args.repo)
     if args.cmd == "gather":
         gather(gh, args.number, os.environ.get("DENY_PATHS", ""), args.dir)

scripts/test_pr_risk.py

@@ -2,6 +2,7 @@
 import os
 from pathlib import Path
 import re
+import subprocess
 import sys
 import tempfile
 import unittest
@@ -37,6 +38,15 @@
     print("{}")
 '''
 
+# Stands in for `pi`: records its arguments, working directory and stdin, and replays canned events.
+FAKE_PI = r'''#!/usr/bin/env python3
+import json, os, sys
+from pathlib import Path
+Path(os.environ["FAKE_PI_LOG"]).write_text(json.dumps({
+    "args": sys.argv[1:], "cwd": os.getcwd(), "stdin": sys.stdin.read()}))
+sys.stdout.write(Path(os.environ["FAKE_PI_EVENTS"]).read_text())
+'''
+
 PATCH = "@@ -10,3 +10,4 @@ flags:\n   - a\n   - b\n+  - c\n   - d"
 
 
@@ -102,6 +112,26 @@ def visible_words(body):
            "reason": "Adds one org to an existing beta allowlist."}
 
 
+def pi_events(text=None, stop="stop", error=None):
+    """A `pi --mode json` stream whose final assistant reply is `text`."""
+    reply = {"role": "assistant", "content": [{"type": "thinking", "thinking": "{}"}],
+             "stopReason": stop}
+    if text is not None:
+        reply["content"].append({"type": "text", "text": text})
+    if error:
+        reply["errorMessage"] = error
+    events = [
+        {"type": "session", "version": 3, "id": "s", "cwd": "/tmp"},
+        {"type": "agent_start"},
+        {"type": "message_end", "message": {"role": "user", "content": [{"type": "text", "text": "PR"}]}},
+        {"type": "message_update", "assistantMessageEvent": {"type": "text_delta", "delta": "{"}},
+        {"type": "message_end", "message": reply},
+        {"type": "agent_end", "messages": [reply], "willRetry": False},
+        {"type": "agent_settled", "aborted": False},
+    ]
+    return "".join(json.dumps(e) + "\n" for e in events)
+
+
 class PRRiskTests(unittest.TestCase):
     def run_pr(self, gh, verdict=VERDICT, result=None, extra_deny=""):
         with tempfile.TemporaryDirectory() as tmp:
@@ -285,20 +315,101 @@ def test_github_client_calls_gh(self):
         self.assertEqual(json.loads(calls[3]["stdin"]), {"body": "hello"})
         self.assertEqual(calls[4]["args"], ["api", "/repos/kernel/infra/collaborators/kernel-internal%5Bbot%5D/permission"])
 
+    def test_verdict_from_events(self):
+        reply = json.dumps(VERDICT)
+        self.assertEqual(pr_risk.verdict_from_events(pi_events(reply)),
+                         {"is_error": False, "result": reply, "structured_output": VERDICT})
+        fenced = pr_risk.verdict_from_events(pi_events(f"```json\n{reply}\n```"))
+        self.assertEqual(fenced["structured_output"], VERDICT)
+        failures = {
+            "pi stopped with error: 401 bad key": pi_events(None, stop="error", error="401 bad key"),
+            "pi stopped with length": pi_events(reply[:20], stop="length"),
+            "pi produced no assistant reply": pi_events(reply).replace('"assistant"', '"user"'),
+            "reply is not JSON": pi_events(f"Here is the verdict: {reply}"),
+            "reply is not a JSON object": pi_events("[]"),
+            "unexpected keys ['approve']": pi_events(json.dumps({**VERDICT, "approve": True})),
+            "missing reason": pi_events(json.dumps({k: v for k, v in VERDICT.items() if k != "reason"})),
+            "deny must be a boolean": pi_events(json.dumps({**VERDICT, "deny": "false"})),
+            "tier must be one of": pi_events(json.dumps({**VERDICT, "tier": "none"})),
+            "confidence must be one of": pi_events(json.dumps({**VERDICT, "confidence": "certain"})),
+        }
+        for expected, events in failures.items():
+            with self.subTest(expected=expected):
+                result = pr_risk.verdict_from_events(events)
+                self.assertTrue(result["is_error"])
+                self.assertNotIn("structured_output", result)
+                self.assertIn(expected, result["result"])
+                with self.assertRaisesRegex(ValueError, "classifier errored"):
+                    pr_risk.parse_verdict(result)
+
+    def test_verdict_from_events_feeds_comment(self):
+        gh = FakeGitHub(pr(), [file()])
+        _, code, _ = self.run_pr(gh, result=pr_risk.verdict_from_events(pi_events(json.dumps(VERDICT))))
+        self.assertEqual(code, 0)
+        self.assertIn("**Would auto-approve: yes** · low risk", gh.created[0])
+
+        existing = {"id": 3, "user": {"type": "Bot"}, "body": pr_risk.MARKER}
+        gh = FakeGitHub(pr(), [file()], comments=[existing])
+        _, code, _ = self.run_pr(gh, result=pr_risk.verdict_from_events(pi_events(None, stop="error", error="x")))
+        self.assertEqual(code, 1)
+        self.assertIn("**Not assessed**", gh.updated[0][1])
+
     def test_workflow_isolates_the_model(self):
         workflow = yaml.safe_load(WORKFLOW.read_text())
         job = workflow["jobs"]["classify"]
         self.assertEqual(job["permissions"], {"contents": "read", "pull-requests": "read"})
         steps = {s["name"]: s for s in job["steps"]}
         classify = steps["Classify"]
-        self.assertEqual(set(classify["env"]), {"ANTHROPIC_API_KEY", "MODEL"})
-        self.assertIn('--tools ""', classify["run"])
+        self.assertEqual(set(classify["env"]), {"ANTHROPIC_API_KEY", "OPENAI_API_KEY", "MODEL",
+                                                "PI_SKIP_VERSION_CHECK", "PI_TELEMETRY"})
+        self.assertEqual(classify["env"]["MODEL"], "${{ inputs.model }}")
+        for flag in ("--no-tools", "--no-mcp", "--no-extensions", "--no-skills", "--no-prompt-templates",
+                     "--no-context-files", "--no-approve", "--no-session"):
+            self.assertIn(flag, classify["run"])
+        self.assertNotIn("--approve ", classify["run"])
+        self.assertNotIn("GH_TOKEN", json.dumps({k: steps[k] for k in ("Install Pi", "Classify")}))
         names = list(steps)
         self.assertLess(names.index("Classify"), names.index("Generate app token"))
         self.assertEqual(steps["Generate app token"]["with"]["permission-pull-requests"], "write")
         self.assertEqual(steps["Gather"]["env"]["GH_TOKEN"], "${{ github.token }}")
         triggers = workflow[True]["workflow_call"]
-        self.assertEqual(set(triggers["secrets"]), {"ADMIN_APP_ID", "ADMIN_APP_PRIVATE_KEY", "ANTHROPIC_API_KEY"})
+        self.assertEqual(triggers["inputs"]["model"]["default"], "${{ vars.PI_PREFERRED_MODEL }}")
+        self.assertEqual(set(triggers["secrets"]),
+                         {"ADMIN_APP_ID", "ADMIN_APP_PRIVATE_KEY", "ANTHROPIC_API_KEY", "OPENAI_API_KEY"})
+        self.assertFalse(triggers["secrets"]["ANTHROPIC_API_KEY"]["required"])
+        self.assertFalse(triggers["secrets"]["OPENAI_API_KEY"]["required"])
+
+    def test_classify_step_runs_pi(self):
+        script = yaml.safe_load(WORKFLOW.read_text())["jobs"]["classify"]["steps"]
+        script = next(s for s in script if s["name"] == "Classify")["run"]
+        repo = Path(__file__).resolve().parents[1]
+        script = (script.replace("/tmp/pr_risk.py", str(repo / "scripts/pr_risk.py"))
+                        .replace("/tmp/pr-risk-prompt.md", str(repo / ".github/workflows/pr-risk/classify-prompt.md")))
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            work = root / "work"
+            work.mkdir()
+            (work / "input.md").write_text("# Pull request\n")
+            (root / "bin").mkdir()
+            (root / "bin/pi").write_text(FAKE_PI)
+            (root / "bin/pi").chmod(0o755)
+            (root / "events.jsonl").write_text(pi_events(json.dumps(VERDICT)))
+            env = {"PATH": f"{root / 'bin'}:{os.environ['PATH']}", "WORKDIR": str(work),
+                   "MODEL": "openai/some-model", "OPENAI_API_KEY": "sk-test",
+                   "FAKE_PI_LOG": str(root / "pi.json"), "FAKE_PI_EVENTS": str(root / "events.jsonl")}
+            subprocess.run(["bash", "-c", script], env=env, check=True)
+            call = json.loads((root / "pi.json").read_text())
+            result = json.loads((work / "result.json").read_text())
+
+            empty = subprocess.run(["bash", "-c", script], env={**env, "MODEL": ""}, capture_output=True, text=True)
+        self.assertNotEqual(empty.returncode, 0)
+        self.assertIn("PI_PREFERRED_MODEL", empty.stderr)
+        self.assertEqual(result, {"is_error": False, "result": json.dumps(VERDICT), "structured_output": VERDICT})
+        self.assertEqual(call["stdin"], "# Pull request\n")
+        self.assertEqual(call["cwd"], str(work / "cwd"))
+        args = call["args"]
+        self.assertEqual(args[args.index("--model") + 1], "openai/some-model")
+        self.assertTrue(args[args.index("--system-prompt") + 1].endswith("classify-prompt.md"))
 
 
 if __name__ == "__main__":