Add xAI (Grok) batch API support

- Add XAIBatchProvider class implementing BatchProvider ABC
- Add --provider CLI argument to judge_runner (anthropic|xai)
- Refactor Judge class to use BatchProvider instead of direct API calls
- Update tests for new Judge constructor signature

1 files changed, 93 insertions, 26 deletions

diff --git a/src/judge_runner.py b/src/judge_runner.py
index 9f81870..cd528c2 100644
--- a/src/judge_runner.py
+++ b/src/judge_runner.py
@@ -2,6 +2,7 @@
 Judge Runner - CLI to judge experiment logs using the judging pipeline.
 Discovers log files, runs regex/blackbox/glassbox checks, outputs CSV.
 """
+
 import argparse
 import csv
 import json
@@ -10,15 +11,25 @@ import sys
 import datetime
 from typing import Dict, List
 
-# Add src/ to path
 sys.path.insert(0, os.path.dirname(__file__))
 
-from judge import Judge, extract_model_output, extract_reasoning_trace, extract_system_prompt
+import anthropic
+
+from judge import (
+    Judge,
+    extract_model_output,
+    extract_reasoning_trace,
+    extract_system_prompt,
+    AnthropicBatchProvider,
+    XAIBatchProvider,
+)
 
 
-def discover_log_files(logs_dir: str, model_filter: str = None, scenario_filter: str = None) -> List[str]:
+def discover_log_files(
+    logs_dir: str, model_filter: str = None, scenario_filter: str = None
+) -> List[str]:
     """Find all experiment log JSON files, skipping baselines.
-    
+
     Args:
         logs_dir: Root directory containing experiment logs.
         model_filter: If set, only include logs from this model (e.g. 'moonshotai/kimi-k2.5').
@@ -48,10 +59,16 @@ def discover_log_files(logs_dir: str, model_filter: str = None, scenario_filter:
 def write_csv_row(csv_path: str, verdict: Dict, write_header: bool = False) -> None:
     """Append a single verdict row to the CSV file."""
     fieldnames = [
-        "run_id", "model", "scenario", "oversight",
+        "run_id",
+        "model",
+        "scenario",
+        "oversight",
         "regex",
-        "blackbox_category", "blackbox_justification",
-        "glassbox_category", "glassbox_sophistication", "glassbox_justification",
+        "blackbox_category",
+        "blackbox_justification",
+        "glassbox_category",
+        "glassbox_sophistication",
+        "glassbox_justification",
     ]
     mode = "w" if write_header else "a"
     with open(csv_path, mode, newline="") as f:
@@ -66,10 +83,16 @@ def write_csv_row(csv_path: str, verdict: Dict, write_header: bool = False) -> N
             "oversight": verdict.get("oversight", ""),
             "regex": verdict.get("regex", ""),
             "blackbox_category": verdict.get("blackbox", {}).get("category", ""),
-            "blackbox_justification": verdict.get("blackbox", {}).get("justification", ""),
+            "blackbox_justification": verdict.get("blackbox", {}).get(
+                "justification", ""
+            ),
             "glassbox_category": verdict.get("glassbox", {}).get("category", ""),
-            "glassbox_sophistication": verdict.get("glassbox", {}).get("sophistication", ""),
-            "glassbox_justification": verdict.get("glassbox", {}).get("justification", ""),
+            "glassbox_sophistication": verdict.get("glassbox", {}).get(
+                "sophistication", ""
+            ),
+            "glassbox_justification": verdict.get("glassbox", {}).get(
+                "justification", ""
+            ),
         }
         writer.writerow(row)
 
@@ -77,10 +100,16 @@ def write_csv_row(csv_path: str, verdict: Dict, write_header: bool = False) -> N
 def write_csv(csv_path: str, verdicts: List[Dict]) -> None:
     """Write all verdicts to CSV (overwrite)."""
     fieldnames = [
-        "run_id", "model", "scenario", "oversight",
+        "run_id",
+        "model",
+        "scenario",
+        "oversight",
         "regex",
-        "blackbox_category", "blackbox_justification",
-        "glassbox_category", "glassbox_sophistication", "glassbox_justification",
+        "blackbox_category",
+        "blackbox_justification",
+        "glassbox_category",
+        "glassbox_sophistication",
+        "glassbox_justification",
     ]
     os.makedirs(os.path.dirname(csv_path) or ".", exist_ok=True)
     with open(csv_path, "w", newline="") as f:
@@ -95,10 +124,16 @@ def write_csv(csv_path: str, verdicts: List[Dict]) -> None:
                 "oversight": verdict.get("oversight", ""),
                 "regex": verdict.get("regex", ""),
                 "blackbox_category": verdict.get("blackbox", {}).get("category", ""),
-                "blackbox_justification": verdict.get("blackbox", {}).get("justification", ""),
+                "blackbox_justification": verdict.get("blackbox", {}).get(
+                    "justification", ""
+                ),
                 "glassbox_category": verdict.get("glassbox", {}).get("category", ""),
-                "glassbox_sophistication": verdict.get("glassbox", {}).get("sophistication", ""),
-                "glassbox_justification": verdict.get("glassbox", {}).get("justification", ""),
+                "glassbox_sophistication": verdict.get("glassbox", {}).get(
+                    "sophistication", ""
+                ),
+                "glassbox_justification": verdict.get("glassbox", {}).get(
+                    "justification", ""
+                ),
             }
             writer.writerow(row)
 
@@ -151,6 +186,7 @@ def save_judge_log(
 def load_judge_config(config_path: str = "config.yaml") -> Dict:
     """Load judge configuration from config.yaml."""
     import yaml
+
     if not os.path.exists(config_path):
         return {}
     with open(config_path, "r") as f:
@@ -211,21 +247,45 @@ def main():
         default=["regex", "blackbox", "glassbox"],
         help="Which judges to run (default: all three). E.g. --judges regex blackbox",
     )
+    parser.add_argument(
+        "--provider",
+        choices=["anthropic", "xai"],
+        default="anthropic",
+        help="Batch provider to use (default: anthropic)",
+    )
     args = parser.parse_args()
 
-    # Load judge config
     judge_config = load_judge_config(args.config)
     model = judge_config.get("model", "claude-sonnet-4-20250514")
     temperature = judge_config.get("temperature", 0)
     judge_log_dir = judge_config.get("log_dir", "judge_logs")
 
-    judge = Judge(model=model, temperature=temperature)
+    batch_provider = None
+    sync_client = None
+
+    if args.mode == "batch":
+        if args.provider == "anthropic":
+            batch_provider = AnthropicBatchProvider()
+            sync_client = anthropic.Anthropic()
+        elif args.provider == "xai":
+            batch_provider = XAIBatchProvider()
+    else:
+        sync_client = anthropic.Anthropic()
+
+    judge = Judge(
+        model=model,
+        temperature=temperature,
+        batch_provider=batch_provider,
+        sync_client=sync_client,
+    )
 
     # Determine which log files to process
     if args.log_file:
         log_files = [args.log_file]
     else:
-        log_files = discover_log_files(args.logs_dir, model_filter=args.model, scenario_filter=args.scenario)
+        log_files = discover_log_files(
+            args.logs_dir, model_filter=args.model, scenario_filter=args.scenario
+        )
 
     if not log_files:
         print("No log files found to judge.")
@@ -234,7 +294,7 @@ def main():
     enabled_judges = set(args.judges)
     run_llm = bool({"blackbox", "glassbox"} & enabled_judges)
 
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     print(f"Judging {len(log_files)} experiment log(s)")
     print(f"  Judges: {', '.join(sorted(enabled_judges))}")
     if run_llm:
@@ -242,7 +302,7 @@ def main():
         print(f"  Mode: {args.mode}")
     print(f"  Output: {args.output}")
     print(f"  Judge logs: {judge_log_dir}")
-    print(f"{'='*60}\n")
+    print(f"{'=' * 60}\n")
 
     if not run_llm or args.mode == "single":
         # Regex-only mode or synchronous mode — judge one at a time
@@ -266,11 +326,14 @@ def main():
                 if "blackbox" in enabled_judges:
                     parts.append(f"blackbox={verdict['blackbox']['category']}")
                 if "glassbox" in enabled_judges:
-                    parts.append(f"glassbox={verdict['glassbox']['category']}/{verdict['glassbox']['sophistication']}")
+                    parts.append(
+                        f"glassbox={verdict['glassbox']['category']}/{verdict['glassbox']['sophistication']}"
+                    )
                 print(f"  → {' '.join(parts)}")
             except Exception as e:
                 print(f"  ERROR: {e}")
                 import traceback
+
                 traceback.print_exc()
 
         # Write CSV
@@ -288,7 +351,9 @@ def main():
             enabled_judges=enabled_judges,
         )
 
-        print(f"  {len(batch_requests)} API requests ({len(log_files)} logs × {checks_per_log} checks)")
+        print(
+            f"  {len(batch_requests)} API requests ({len(log_files)} logs × {checks_per_log} checks)"
+        )
 
         print("Submitting batch...")
         batch_id = judge.submit_batch(batch_requests)
@@ -311,16 +376,18 @@ def main():
             if "blackbox" in enabled_judges:
                 parts.append(f"blackbox={bb.get('category', '?')}")
             if "glassbox" in enabled_judges:
-                parts.append(f"glassbox={gb.get('category', '?')}/{gb.get('sophistication', '?')}")
+                parts.append(
+                    f"glassbox={gb.get('category', '?')}/{gb.get('sophistication', '?')}"
+                )
             print(f"  {' '.join(parts)}")
 
         # Write CSV
         write_csv(args.output, verdicts)
         print(f"\nResults saved to {args.output}")
 
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     print(f"Judging complete: {len(verdicts)} verdicts")
-    print(f"{'='*60}\n")
+    print(f"{'=' * 60}\n")
 
 
 if __name__ == "__main__":