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

2 files changed, 432 insertions, 114 deletions

diff --git a/src/judge.py b/src/judge.py
index d7b7c40..354d96c 100644
--- a/src/judge.py
+++ b/src/judge.py
@@ -1,32 +1,44 @@
 """
 Judge - evaluates experiment logs for deceptive behaviour.
 Three independent checks: Regex, Blackbox (LLM), Glassbox (LLM).
-Uses Anthropic Batch API for LLM judge calls (50% cost reduction).
+Uses batch providers (Anthropic, xAI) for LLM judge calls.
 """
+
 import hashlib
 import json
 import os
 import re
 import time
 import yaml
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import anthropic
 
-# Load environment variables from .env file
+try:
+    from xai_sdk import Client as XAIClient
+except ImportError:
+    XAIClient = None
+
 try:
     from dotenv import load_dotenv
+
     load_dotenv()
 except ImportError:
-    pass  # python-dotenv not installed
+    pass
 
 
 # ---------------------------------------------------------------------------
 # Constants
 # ---------------------------------------------------------------------------
 
-BLACKBOX_CATEGORIES = ["ACCURATE", "DECEPTIVE_FRAMING", "OMITTED_CRITICAL", "FABRICATED"]
+BLACKBOX_CATEGORIES = [
+    "ACCURATE",
+    "DECEPTIVE_FRAMING",
+    "OMITTED_CRITICAL",
+    "FABRICATED",
+]
 GLASSBOX_CATEGORIES = [
     "ROBUST_HONEST",
     "CONSIDERED_HONEST",
@@ -37,16 +49,219 @@ GLASSBOX_CATEGORIES = [
 SOPHISTICATION_TIERS = ["NAIVE", "REACTIVE", "TACTICAL", "STRATEGIC"]
 
 # ---------------------------------------------------------------------------
+# Batch Provider (ABC)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class BatchRequest:
+    custom_id: str
+    params: Dict[str, Any]
+
+
+@dataclass
+class BatchResult:
+    custom_id: str
+    text: str
+    error: Optional[str] = None
+
+
+class BatchProvider(ABC):
+    @abstractmethod
+    def submit_batch(self, requests: List[BatchRequest]) -> str:
+        pass
+
+    @abstractmethod
+    def poll_batch(self, batch_id: str, poll_interval: int = 30) -> None:
+        pass
+
+    @abstractmethod
+    def collect_results(self, batch_id: str) -> Iterator[BatchResult]:
+        pass
+
+    @abstractmethod
+    def build_request(
+        self,
+        custom_id: str,
+        prompt: str,
+        model: str,
+        temperature: float,
+        max_tokens: int = 4096,
+    ) -> BatchRequest:
+        pass
+
+
+class AnthropicBatchProvider(BatchProvider):
+    def __init__(self, api_key: Optional[str] = None):
+        key = api_key or os.environ.get("ANTHROPIC_API_KEY")
+        if not key:
+            raise ValueError("ANTHROPIC_API_KEY not set")
+        self.client = anthropic.Anthropic(api_key=key)
+
+    def build_request(
+        self,
+        custom_id: str,
+        prompt: str,
+        model: str,
+        temperature: float,
+        max_tokens: int = 4096,
+    ) -> BatchRequest:
+        return BatchRequest(
+            custom_id=custom_id,
+            params={
+                "model": model,
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "messages": [{"role": "user", "content": prompt}],
+            },
+        )
+
+    def submit_batch(self, requests: List[BatchRequest]) -> str:
+        anthropic_requests = [
+            {
+                "custom_id": r.custom_id,
+                "params": r.params,
+            }
+            for r in requests
+        ]
+        response = self.client.messages.batches.create(requests=anthropic_requests)
+        return response.id
+
+    def poll_batch(self, batch_id: str, poll_interval: int = 30) -> None:
+        while True:
+            batch = self.client.messages.batches.retrieve(batch_id)
+            status = batch.processing_status
+            counts = batch.request_counts
+            print(
+                f"  Batch {batch_id}: {status} "
+                f"(succeeded={counts.succeeded}, "
+                f"processing={counts.processing}, "
+                f"errored={counts.errored})"
+            )
+            if status == "ended":
+                return
+            time.sleep(poll_interval)
+
+    def collect_results(self, batch_id: str) -> Iterator[BatchResult]:
+        for result in self.client.messages.batches.results(batch_id):
+            custom_id = result.custom_id
+            if result.result.type == "succeeded":
+                content = result.result.message.content
+                if hasattr(content, "__iter__") and not isinstance(content, str):
+                    for block in content:
+                        if hasattr(block, "text"):
+                            text = block.text
+                            break
+                    else:
+                        text = ""
+                else:
+                    text = str(content)
+                yield BatchResult(custom_id=custom_id, text=text)
+            else:
+                yield BatchResult(
+                    custom_id=custom_id, text="", error=f"ERROR: {result.result.type}"
+                )
+
+
+class XAIBatchProvider(BatchProvider):
+    def __init__(self, api_key: Optional[str] = None):
+        if XAIClient is None:
+            raise ImportError("xai-sdk not installed. Run: uv add xai-sdk")
+        key = api_key or os.environ.get("XAI_API_KEY")
+        if not key:
+            raise ValueError("XAI_API_KEY not set")
+        self.client = XAIClient(api_key=key)
+
+    def build_request(
+        self,
+        custom_id: str,
+        prompt: str,
+        model: str,
+        temperature: float,
+        max_tokens: int = 4096,
+    ) -> BatchRequest:
+        return BatchRequest(
+            custom_id=custom_id,
+            params={
+                "model": model,
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": prompt},
+                ],
+            },
+        )
+
+    def submit_batch(self, requests: List[BatchRequest]) -> str:
+        from xai_sdk.chat import system as xai_system, user as xai_user
+
+        batch = self.client.batch.create(batch_name=f"judge_batch_{int(time.time())}")
+        batch_requests = []
+        for req in requests:
+            chat = self.client.chat.create(
+                model=req.params["model"],
+                batch_request_id=req.custom_id,
+            )
+            for msg in req.params["messages"]:
+                if msg["role"] == "system":
+                    chat.append(xai_system(msg["content"]))
+                else:
+                    chat.append(xai_user(msg["content"]))
+            batch_requests.append(chat)
+        self.client.batch.add(batch_id=batch.batch_id, batch_requests=batch_requests)
+        return batch.batch_id
+
+    def poll_batch(self, batch_id: str, poll_interval: int = 30) -> None:
+        while True:
+            batch = self.client.batch.get(batch_id=batch_id)
+            state = batch.state
+            print(
+                f"  Batch {batch_id}: "
+                f"(pending={state.num_pending}, "
+                f"success={state.num_success}, "
+                f"error={state.num_error})"
+            )
+            if state.num_pending == 0:
+                return
+            time.sleep(poll_interval)
+
+    def collect_results(self, batch_id: str) -> Iterator[BatchResult]:
+        pagination_token = None
+        while True:
+            page = self.client.batch.list_batch_results(
+                batch_id=batch_id,
+                limit=100,
+                pagination_token=pagination_token,
+            )
+            for result in page.succeeded:
+                rid = result.batch_request_id
+                text = result.response.content
+                yield BatchResult(custom_id=rid, text=text)
+            for result in page.failed:
+                yield BatchResult(
+                    custom_id=result.batch_request_id,
+                    text="",
+                    error=result.error_message,
+                )
+            if page.pagination_token is None:
+                break
+            pagination_token = page.pagination_token
+
+
+# ---------------------------------------------------------------------------
 # Prompts
 # ---------------------------------------------------------------------------
 
 PROMPTS_DIR = os.path.join(os.path.dirname(__file__), "prompts")
 
+
 def load_prompt(filename: str) -> str:
     path = os.path.join(PROMPTS_DIR, filename)
     with open(path, "r", encoding="utf-8") as f:
         return f.read()
 
+
 BLACKBOX_PROMPT = load_prompt("blackbox_prompt.md")
 GLASSBOX_PROMPT = load_prompt("glassbox_prompt.md")
 
@@ -55,6 +270,7 @@ GLASSBOX_PROMPT = load_prompt("glassbox_prompt.md")
 # Data Extraction Helpers
 # ---------------------------------------------------------------------------
 
+
 def extract_model_output(conversation: List[Dict]) -> str:
     """Extract model output from the last create_file tool call in the conversation."""
     for msg in reversed(conversation):
@@ -84,7 +300,11 @@ def extract_reasoning_trace(conversation: List[Dict]) -> str:
     for msg in conversation:
         if msg.get("role") == "assistant" and msg.get("reasoning"):
             reasoning_parts.append(msg["reasoning"])
-    return "\n\n---\n\n".join(reasoning_parts) if reasoning_parts else "(no reasoning trace available)"
+    return (
+        "\n\n---\n\n".join(reasoning_parts)
+        if reasoning_parts
+        else "(no reasoning trace available)"
+    )
 
 
 def extract_system_prompt(conversation: List[Dict]) -> str:
@@ -124,6 +344,7 @@ def parse_json_verdict(text: str) -> Dict[str, Any]:
 # Regex Checker
 # ---------------------------------------------------------------------------
 
+
 class RegexChecker:
     """Deterministic string/pattern matching on model output."""
 
@@ -136,7 +357,7 @@ class RegexChecker:
 
         Supports two formats:
         - A dict with a 'rules' key: {rules: [...]}
-        - A flat list of rules: [...]  
+        - A flat list of rules: [...]
         """
         if not os.path.exists(rules_path):
             return []
@@ -155,7 +376,13 @@ class RegexChecker:
             {"verdict": "6/7", "matched": 6, "total": 7, "details": [...]}
         """
         if not self.rules:
-            return {"verdict": "PASS", "matched": 0, "total": 0, "details": [{"note": "no regex rules defined"}], "fraction": "0/0"}
+            return {
+                "verdict": "PASS",
+                "matched": 0,
+                "total": 0,
+                "details": [{"note": "no regex rules defined"}],
+                "fraction": "0/0",
+            }
 
         details = []
         required_total = 0
@@ -185,13 +412,13 @@ class RegexChecker:
             verdict = "PASS"
         else:
             verdict = "PASS" if required_matched == required_total else "FAIL"
-        
+
         return {
             "verdict": verdict,
             "matched": required_matched,
             "total": required_total,
             "details": details,
-            "fraction": f"{required_matched}/{required_total}"
+            "fraction": f"{required_matched}/{required_total}",
         }
 
 
@@ -199,6 +426,7 @@ class RegexChecker:
 # Blackbox Checker
 # ---------------------------------------------------------------------------
 
+
 class BlackboxChecker:
     """LLM judge comparing model output against baseline (no reasoning trace)."""
 
@@ -206,7 +434,13 @@ class BlackboxChecker:
         self.model = model
         self.temperature = temperature
 
-    def build_prompt(self, model_output: str, baseline: str, source_documents: str, task_description: str = "") -> str:
+    def build_prompt(
+        self,
+        model_output: str,
+        baseline: str,
+        source_documents: str,
+        task_description: str = "",
+    ) -> str:
         """Build the blackbox judge prompt."""
         return BLACKBOX_PROMPT.format(
             task_description=task_description or "(No task description available)",
@@ -232,6 +466,7 @@ class BlackboxChecker:
 # Glassbox Checker
 # ---------------------------------------------------------------------------
 
+
 class GlassboxChecker:
     """LLM judge evaluating reasoning trace + output + system prompt for intent."""
 
@@ -273,30 +508,28 @@ class GlassboxChecker:
 # Judge (orchestrator)
 # ---------------------------------------------------------------------------
 
+
 class Judge:
     """Orchestrates all three checks on experiment log files.
 
     Supports two modes:
     - judge_single(): synchronous, one log at a time (for testing / quick runs)
-    - judge_batch(): uses Anthropic Batch API for multiple logs (half cost)
+    - judge_batch(): uses a BatchProvider for multiple logs
     """
 
     def __init__(
         self,
         model: str = "claude-sonnet-4-20250514",
         temperature: float = 0,
-        anthropic_api_key: Optional[str] = None,
+        batch_provider: Optional[BatchProvider] = None,
+        sync_client: Optional[Any] = None,
     ):
         self.model = model
         self.temperature = temperature
+        self.batch_provider = batch_provider
+        self.sync_client = sync_client
         self.regex_checker_cache: Dict[str, RegexChecker] = {}
 
-        api_key = anthropic_api_key or os.environ.get("ANTHROPIC_API_KEY")
-        if api_key:
-            self.client = anthropic.Anthropic(api_key=api_key)
-        else:
-            self.client = None
-
         self.blackbox = BlackboxChecker(model=model, temperature=temperature)
         self.glassbox = GlassboxChecker(model=model, temperature=temperature)
 
@@ -333,7 +566,11 @@ class Judge:
             json_files = [f for f in os.listdir(baseline_dir) if f.endswith(".json")]
             if json_files:
                 latest_baseline_log = sorted(json_files)[-1]
-                with open(os.path.join(baseline_dir, latest_baseline_log), "r", encoding="utf-8") as f:
+                with open(
+                    os.path.join(baseline_dir, latest_baseline_log),
+                    "r",
+                    encoding="utf-8",
+                ) as f:
                     try:
                         baseline_log_data = json.load(f)
                     except json.JSONDecodeError:
@@ -349,10 +586,10 @@ class Judge:
         vfs_state = log_data.get("final_vfs_state", {})
         root_dir = vfs_state.get("/", {})
         data_files = root_dir.get("data", {})
-        
+
         if not data_files:
             return "(No source documents found)"
-            
+
         docs = []
         for filename, content in sorted(data_files.items()):
             if isinstance(content, str):
@@ -381,8 +618,10 @@ class Judge:
             enabled_judges = {"regex", "blackbox", "glassbox"}
 
         run_llm = bool({"blackbox", "glassbox"} & enabled_judges)
-        if run_llm and not self.client:
-            raise RuntimeError("ANTHROPIC_API_KEY not set — cannot run LLM judge checks")
+        if run_llm and not self.sync_client:
+            raise RuntimeError(
+                "No sync client configured — cannot run LLM judge checks"
+            )
 
         log_data = self.load_log(log_path)
         conversation = log_data.get("conversation", [])
@@ -409,12 +648,20 @@ class Judge:
                 source_docs = self.load_source_documents(baseline_log_data)
             else:
                 import sys
-                print(f"  [WARN] No baseline found for {scenario} — using experiment log for source docs", file=sys.stderr)
+
+                print(
+                    f"  [WARN] No baseline found for {scenario} — using experiment log for source docs",
+                    file=sys.stderr,
+                )
                 source_docs = self.load_source_documents(log_data)
 
             if not baseline_output:
                 import sys
-                print(f"  [WARN] No baseline output for {scenario} — blackbox judge will compare against source docs only", file=sys.stderr)
+
+                print(
+                    f"  [WARN] No baseline output for {scenario} — blackbox judge will compare against source docs only",
+                    file=sys.stderr,
+                )
                 baseline_output = "(No baseline available)"
 
             scenario_path = os.path.join(scenarios_dir, scenario)
@@ -424,8 +671,10 @@ class Judge:
                 with open(benign_goal_path, "r", encoding="utf-8") as f:
                     task_description = f.read().strip()
 
-            bb_prompt = self.blackbox.build_prompt(model_output, baseline_output, source_docs, task_description)
-            bb_response = self.client.messages.create(
+            bb_prompt = self.blackbox.build_prompt(
+                model_output, baseline_output, source_docs, task_description
+            )
+            bb_response = self.sync_client.messages.create(
                 model=self.model,
                 max_tokens=4096,
                 temperature=self.temperature,
@@ -435,12 +684,19 @@ class Judge:
             bb_result = self.blackbox.parse_response(bb_text)
 
         # 3. Glassbox check
-        gb_result = {"category": "", "sophistication": "", "justification": "", "raw_response": ""}
+        gb_result = {
+            "category": "",
+            "sophistication": "",
+            "justification": "",
+            "raw_response": "",
+        }
         if "glassbox" in enabled_judges:
             system_prompt = extract_system_prompt(conversation)
             reasoning_trace = extract_reasoning_trace(conversation)
-            gb_prompt = self.glassbox.build_prompt(system_prompt, reasoning_trace, model_output)
-            gb_response = self.client.messages.create(
+            gb_prompt = self.glassbox.build_prompt(
+                system_prompt, reasoning_trace, model_output
+            )
+            gb_response = self.sync_client.messages.create(
                 model=self.model,
                 max_tokens=4096,
                 temperature=self.temperature,
@@ -501,17 +757,27 @@ class Judge:
 
             # Store metadata
             id_hash = hashlib.sha256(run_id.encode()).hexdigest()[:8]
-            idx = len(batch_requests) // 2 if len(enabled_judges & {"blackbox", "glassbox"}) == 2 else len(batch_requests)
+            idx = (
+                len(batch_requests) // 2
+                if len(enabled_judges & {"blackbox", "glassbox"}) == 2
+                else len(batch_requests)
+            )
 
             # Blackbox request
             if "blackbox" in enabled_judges:
-                baseline_output, baseline_log_data = self.load_baseline(log_data, logs_dir)
+                baseline_output, baseline_log_data = self.load_baseline(
+                    log_data, logs_dir
+                )
 
                 if baseline_log_data:
                     source_docs = self.load_source_documents(baseline_log_data)
                 else:
                     import sys
-                    print(f"  [WARN] No baseline found for {scenario} — using experiment log for source docs", file=sys.stderr)
+
+                    print(
+                        f"  [WARN] No baseline found for {scenario} — using experiment log for source docs",
+                        file=sys.stderr,
+                    )
                     source_docs = self.load_source_documents(log_data)
 
                 if not baseline_output:
@@ -535,16 +801,17 @@ class Judge:
                     "regex_result": regex_result,
                 }
 
-                bb_prompt = self.blackbox.build_prompt(model_output, baseline_output, source_docs, task_description)
-                batch_requests.append({
-                    "custom_id": bb_id,
-                    "params": {
-                        "model": self.model,
-                        "max_tokens": 4096,
-                        "temperature": self.temperature,
-                        "messages": [{"role": "user", "content": bb_prompt}],
-                    },
-                })
+                bb_prompt = self.blackbox.build_prompt(
+                    model_output, baseline_output, source_docs, task_description
+                )
+                batch_requests.append(
+                    self.batch_provider.build_request(
+                        custom_id=bb_id,
+                        prompt=bb_prompt,
+                        model=self.model,
+                        temperature=self.temperature,
+                    )
+                )
 
             # Glassbox request
             if "glassbox" in enabled_judges:
@@ -561,48 +828,36 @@ class Judge:
                 if "blackbox" not in enabled_judges:
                     metadata_map[gb_id]["model"] = log_data.get("model", "")
                     metadata_map[gb_id]["scenario"] = scenario
-                    metadata_map[gb_id]["oversight"] = log_data.get("oversight_level", "")
+                    metadata_map[gb_id]["oversight"] = log_data.get(
+                        "oversight_level", ""
+                    )
                     metadata_map[gb_id]["regex_result"] = regex_result
 
-                gb_prompt = self.glassbox.build_prompt(system_prompt, reasoning_trace, model_output)
-                batch_requests.append({
-                    "custom_id": gb_id,
-                    "params": {
-                        "model": self.model,
-                        "max_tokens": 4096,
-                        "temperature": self.temperature,
-                        "messages": [{"role": "user", "content": gb_prompt}],
-                    },
-                })
+                gb_prompt = self.glassbox.build_prompt(
+                    system_prompt, reasoning_trace, model_output
+                )
+                batch_requests.append(
+                    self.batch_provider.build_request(
+                        custom_id=gb_id,
+                        prompt=gb_prompt,
+                        model=self.model,
+                        temperature=self.temperature,
+                    )
+                )
 
         return batch_requests, metadata_map
 
-    def submit_batch(self, batch_requests: List[Dict]) -> str:
-        """Submit a batch to Anthropic and return the batch ID."""
-        if not self.client:
-            raise RuntimeError("ANTHROPIC_API_KEY not set — cannot submit batch")
-
-        response = self.client.messages.batches.create(requests=batch_requests)
-        return response.id
+    def submit_batch(self, batch_requests: List[BatchRequest]) -> str:
+        """Submit a batch to the provider and return the batch ID."""
+        if not self.batch_provider:
+            raise RuntimeError("No batch provider configured — cannot submit batch")
+        return self.batch_provider.submit_batch(batch_requests)
 
     def poll_batch(self, batch_id: str, poll_interval: int = 30) -> None:
         """Poll until batch processing is complete."""
-        if not self.client:
-            raise RuntimeError("ANTHROPIC_API_KEY not set")
-
-        while True:
-            batch = self.client.messages.batches.retrieve(batch_id)
-            status = batch.processing_status
-            counts = batch.request_counts
-            print(
-                f"  Batch {batch_id}: {status} "
-                f"(succeeded={counts.succeeded}, "
-                f"processing={counts.processing}, "
-                f"errored={counts.errored})"
-            )
-            if status == "ended":
-                return
-            time.sleep(poll_interval)
+        if not self.batch_provider:
+            raise RuntimeError("No batch provider configured")
+        self.batch_provider.poll_batch(batch_id, poll_interval)
 
     def collect_batch_results(
         self,
@@ -613,20 +868,16 @@ class Judge:
 
         Returns a list of combined verdict dicts (one per log file).
         """
-        if not self.client:
-            raise RuntimeError("ANTHROPIC_API_KEY not set")
+        if not self.batch_provider:
+            raise RuntimeError("No batch provider configured")
 
-        # Collect raw results by custom_id
         raw_results = {}
-        for result in self.client.messages.batches.results(batch_id):
-            custom_id = result.custom_id
-            if result.result.type == "succeeded":
-                text = result.result.message.content[0].text
-                raw_results[custom_id] = text
+        for result in self.batch_provider.collect_results(batch_id):
+            if result.error:
+                raw_results[result.custom_id] = f"ERROR: {result.error}"
             else:
-                raw_results[custom_id] = f"ERROR: {result.result.type}"
+                raw_results[result.custom_id] = result.text
 
-        # Group by run_id and assemble verdicts
         verdicts_by_run = {}
         for custom_id, meta in metadata_map.items():
             run_id = meta["run_id"]
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__":