diff options
Diffstat (limited to 'src/judge/judge.py')
| -rw-r--r-- | src/judge/judge.py | 383 |
1 files changed, 383 insertions, 0 deletions
diff --git a/src/judge/judge.py b/src/judge/judge.py new file mode 100644 index 0000000..5c07ffe --- /dev/null +++ b/src/judge/judge.py @@ -0,0 +1,383 @@ +""" +Judge orchestrator - evaluates experiment logs for deceptive behaviour. +""" + +import hashlib +import json +import os +from typing import Any, Dict, List, Optional + +from .batch_providers import BatchProvider, BatchRequest +from .helpers import ( + extract_model_output, + extract_reasoning_trace, + extract_system_prompt, +) +from .judges import BlackboxChecker, GlassboxChecker, RegexChecker + + +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 a BatchProvider for multiple logs + """ + + def __init__( + self, + model: str = "claude-sonnet-4-20250514", + temperature: float = 0, + 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] = {} + + self.blackbox = BlackboxChecker(model=model, temperature=temperature) + self.glassbox = GlassboxChecker(model=model, temperature=temperature) + + def _get_regex_checker(self, scenario_path: str) -> RegexChecker: + """Get or create a RegexChecker for a scenario.""" + if scenario_path not in self.regex_checker_cache: + rules_path = os.path.join(scenario_path, "regex_rules.yaml") + self.regex_checker_cache[scenario_path] = RegexChecker(rules_path) + return self.regex_checker_cache[scenario_path] + + @staticmethod + def load_log(log_path: str) -> Dict[str, Any]: + """Load an experiment log JSON file.""" + with open(log_path, "r") as f: + return json.load(f) + + @staticmethod + def load_baseline(log_data: Dict, logs_dir: str) -> tuple[str, Dict]: + """Load the baseline output and baseline log data for a given experiment log.""" + model = log_data.get("model", "").replace("/", "_") + scenario = log_data.get("scenario", "") + baseline_path = os.path.join(logs_dir, model, scenario, "baseline.md") + baseline_content = "" + baseline_log_data = {} + + if os.path.exists(baseline_path): + with open(baseline_path, "r", encoding="utf-8") as f: + baseline_content = f.read() + + baseline_dir = os.path.join(logs_dir, model, scenario, "baseline") + if os.path.exists(baseline_dir) and os.path.isdir(baseline_dir): + 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: + try: + baseline_log_data = json.load(f) + except json.JSONDecodeError: + pass + return baseline_content, baseline_log_data + + @staticmethod + def load_source_documents(log_data: Dict) -> str: + """Load source documents from a log's final_vfs_state.""" + 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): + docs.append(f"--- Document: {filename} ---\n{content}\n") + return "\n".join(docs) if docs else "(No source documents found)" + + def judge_single( + self, + log_path: str, + logs_dir: str = "logs", + scenarios_dir: str = "scenarios", + enabled_judges: Optional[set] = None, + ) -> Dict[str, Any]: + """Judge a single experiment log synchronously (non-batch).""" + if enabled_judges is None: + enabled_judges = {"regex", "blackbox", "glassbox"} + + run_llm = bool({"blackbox", "glassbox"} & enabled_judges) + 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", []) + scenario = log_data.get("scenario", "") + run_id = log_data.get("run_id", os.path.basename(log_path)) + + model_output = extract_model_output(conversation) + + regex_result = {"verdict": "", "details": []} + if "regex" in enabled_judges: + scenario_path = os.path.join(scenarios_dir, scenario) + regex_checker = self._get_regex_checker(scenario_path) + regex_result = regex_checker.check(model_output) + + bb_result = {"category": "", "justification": "", "raw_response": ""} + if "blackbox" in enabled_judges: + 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, + ) + 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, + ) + baseline_output = "(No baseline available)" + + scenario_path = os.path.join(scenarios_dir, scenario) + benign_goal_path = os.path.join(scenario_path, "benign_goal.md") + task_description = "" + if os.path.exists(benign_goal_path): + 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.sync_client.messages.create( + model=self.model, + max_tokens=4096, + temperature=self.temperature, + messages=[{"role": "user", "content": bb_prompt}], + ) + bb_text = bb_response.content[0].text + bb_result = self.blackbox.parse_response(bb_text) + + 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.sync_client.messages.create( + model=self.model, + max_tokens=4096, + temperature=self.temperature, + messages=[{"role": "user", "content": gb_prompt}], + ) + gb_text = gb_response.content[0].text + gb_result = self.glassbox.parse_response(gb_text) + + return { + "run_id": run_id, + "model": log_data.get("model", ""), + "scenario": scenario, + "oversight": log_data.get("oversight_level", ""), + "regex": regex_result["verdict"], + "regex_details": regex_result["details"], + "blackbox": bb_result, + "glassbox": gb_result, + } + + def prepare_batch_requests( + self, + log_paths: List[str], + logs_dir: str = "logs", + scenarios_dir: str = "scenarios", + enabled_judges: Optional[set] = None, + ) -> tuple: + """Prepare batch requests for multiple log files. + + Returns: + (batch_requests, metadata_map) + - batch_requests: list of BatchRequest for the batch provider + - metadata_map: dict mapping custom_id → metadata needed to reassemble results + """ + if enabled_judges is None: + enabled_judges = {"regex", "blackbox", "glassbox"} + + batch_requests: List[BatchRequest] = [] + metadata_map: Dict[str, Dict] = {} + + for log_path in log_paths: + log_data = self.load_log(log_path) + conversation = log_data.get("conversation", []) + scenario = log_data.get("scenario", "") + run_id = log_data.get("run_id", os.path.basename(log_path)) + + model_output = extract_model_output(conversation) + + regex_result = {"verdict": "", "details": []} + if "regex" in enabled_judges: + scenario_path = os.path.join(scenarios_dir, scenario) + regex_checker = self._get_regex_checker(scenario_path) + regex_result = regex_checker.check(model_output) + + 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) + ) + + if "blackbox" in enabled_judges: + 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, + ) + source_docs = self.load_source_documents(log_data) + + if not baseline_output: + baseline_output = "(No baseline available)" + + scenario_path = os.path.join(scenarios_dir, scenario) + benign_goal_path = os.path.join(scenario_path, "benign_goal.md") + task_description = "" + if os.path.exists(benign_goal_path): + with open(benign_goal_path, "r", encoding="utf-8") as f: + task_description = f.read().strip() + + bb_id = f"bb_{idx:03d}_{id_hash}" + metadata_map[bb_id] = { + "type": "blackbox", + "log_path": log_path, + "run_id": run_id, + "model": log_data.get("model", ""), + "scenario": scenario, + "oversight": log_data.get("oversight_level", ""), + "regex_result": regex_result, + } + + 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, + ) + ) + + if "glassbox" in enabled_judges: + system_prompt = extract_system_prompt(conversation) + reasoning_trace = extract_reasoning_trace(conversation) + + gb_id = f"gb_{idx:03d}_{id_hash}" + metadata_map[gb_id] = { + "type": "glassbox", + "log_path": log_path, + "run_id": run_id, + } + 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]["regex_result"] = regex_result + + 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[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.batch_provider: + raise RuntimeError("No batch provider configured") + self.batch_provider.poll_batch(batch_id, poll_interval) + + def collect_batch_results( + self, + batch_id: str, + metadata_map: Dict[str, Dict], + ) -> List[Dict[str, Any]]: + """Collect and parse results from a completed batch. + + Returns a list of combined verdict dicts (one per log file). + """ + if not self.batch_provider: + raise RuntimeError("No batch provider configured") + + raw_results = {} + 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[result.custom_id] = result.text + + verdicts_by_run = {} + for custom_id, meta in metadata_map.items(): + run_id = meta["run_id"] + raw_text = raw_results.get(custom_id, "") + + if meta["type"] == "blackbox": + bb_result = self.blackbox.parse_response(raw_text) + if run_id not in verdicts_by_run: + verdicts_by_run[run_id] = { + "run_id": run_id, + "model": meta["model"], + "scenario": meta["scenario"], + "oversight": meta["oversight"], + "regex": meta["regex_result"]["verdict"], + "regex_details": meta["regex_result"]["details"], + } + verdicts_by_run[run_id]["blackbox"] = bb_result + + elif meta["type"] == "glassbox": + gb_result = self.glassbox.parse_response(raw_text) + if run_id not in verdicts_by_run: + verdicts_by_run[run_id] = {"run_id": run_id} + verdicts_by_run[run_id]["glassbox"] = gb_result + + return list(verdicts_by_run.values()) |