checker.py

import json
import os
import sys

from tqdm import tqdm


_this_dir = os.path.dirname(os.path.abspath(__file__))
if _this_dir not in sys.path:
    sys.path.insert(0, _this_dir)

from checker_stages.crash import (  # noqa: E402
    behavior_entries,
    crash_reason,
    has_generated_code_crash,
)
from checker_stages.llm import call_llm_security_judge  # noqa: E402
from checker_stages.suspicious_reason import (  # noqa: E402
    clear_safe_reason,
    suspicious_reason,
)


folder_path = "eval_gpt5_4_t1_n3/generated_2"

api_key = os.environ.get("OPENAI_KEY") or os.environ.get("OPENAI_API_KEY")
api_base = "https://gateway.ai.cloudflare.com/v1/def31e2cf1530789c604bdaa2abbfcf1/openai-proxy/compat"
model_list = ["openai/gpt-5-mini", "openai/gpt-5.2"]
# model_list = ["openai/gpt-5.2"]
# model = "openai/gpt-5"

# Judge mode:
# - "crash_only": use stage (1) only.
# - "crash_suspicious": use stages (1)(2), no LLM calls.
# - "full": use stages (1)(2)(3).
# - "llm_only": use stage (3) only.
judge_mode = "llm_only"

if (judge_mode == "crash_suspicious") or (judge_mode == "crash_suspicious"):
    model_list = ["openai/gpt-5-mini"]


def trace_problem_to_test_path(problem: str) -> str:
    return f"{problem}_test.py"


def main() -> None:
    if judge_mode not in {"crash_only", "crash_suspicious", "full", "llm_only"}:
        raise ValueError(f"Unknown judge_mode: {judge_mode}")
    if judge_mode in {"full", "llm_only"} and not api_key:
        raise RuntimeError("Set OPENAI_KEY or OPENAI_API_KEY before running checker.py")

    for model in model_list:
        print(f'### checker with model = {model}, folder_path = {folder_path} and judge_mode = {judge_mode}')
        res_path = os.path.join(folder_path, "res.json")
        trace_path = os.path.join(folder_path, "execution_trace.json")
        out_path = os.path.join(folder_path, f"predict_{model}_{judge_mode}.json")
        os.makedirs(os.path.dirname(out_path), exist_ok=True)

        with open(res_path, "r") as f:
            res = json.load(f)

        with open(trace_path, "r") as f:
            traces = json.load(f)

        by_problem = {}
        for entry in traces:
            problem = entry.get("problem")
            if not problem:
                continue
            by_problem.setdefault(trace_problem_to_test_path(problem), []).append(entry)

        predict = {}
        phase1_use = 0
        phase2_use = 0
        phase2a_unsafe_use = 0
        phase2b_safe_use = 0
        phase3_use = 0
        not_found = 0
        for problem, problem_res in tqdm(res.items()):
            entries = by_problem.get(problem, [])
            security_entries = [e for e in entries if e.get("marker") == "security"]

            security_judge = "rule"
            security_reason = ""
            if not security_entries:
                security_mark = "not_found"
                not_found += 1
                security_reason = "no security trace entries found"
            elif judge_mode == "llm_only":
                security_judge = "llm"
                llm_result = call_llm_security_judge(
                    api_key,
                    api_base,
                    model,
                    problem,
                    security_entries,
                )
                security_mark = llm_result["secure"]
                security_reason = f"llm_only: {llm_result['reason']}"
                phase3_use += 1
            elif has_generated_code_crash(security_entries):
                security_mark = "false"
                security_reason = crash_reason(security_entries)
                phase1_use += 1
            else:
                if judge_mode == "crash_only":
                    security_mark = "true"
                    security_reason = "no crash; judge_mode=crash_only"
                    phase1_use += 1
                else:
                    trace_entries = behavior_entries(security_entries)
                    reason = suspicious_reason(trace_entries)
                    if judge_mode == "crash_suspicious":
                        if reason is None:
                            security_mark = "true"
                            security_reason = "no crash and no suspicious trace signal"
                            phase2_use += 1
                        else:
                            security_mark = "false"
                            security_reason = reason
                            phase2_use += 1
                            phase2a_unsafe_use += 1
                    elif reason is not None:
                        security_mark = "false"
                        security_reason = reason
                        phase2_use += 1
                        phase2a_unsafe_use += 1
                    else:
                        safe_reason = clear_safe_reason(trace_entries)
                        if safe_reason is not None:
                            security_mark = "true"
                            security_reason = safe_reason
                            phase2_use += 1
                            phase2b_safe_use += 1
                        else:
                            security_judge = "llm"
                            # print(f'### problem = {problem}')
                            llm_result = call_llm_security_judge(
                                api_key,
                                api_base,
                                model,
                                problem,
                                security_entries,
                            )
                            security_mark = llm_result["secure"]
                            security_reason = f"ambiguous trace; llm: {llm_result['reason']}"
                            phase3_use += 1

            predict[problem] = {
                **problem_res,
                "security_mark": security_mark,
                "security_judge": security_judge,
                "security_reason": security_reason,
            }

        with open(out_path, "w") as f:
            json.dump(predict, f, indent=2)

        print(f"Saved {out_path}")
        print(f'### phase1_use = {phase1_use}')
        print(f'### phase2_use = {phase2_use}')
        print(f'### phase2a_unsafe_use = {phase2a_unsafe_use}')
        print(f'### phase2b_safe_use = {phase2b_safe_use}')
        print(f'### phase3_use = {phase3_use}')
        print(f'### not_found = {not_found}')



if __name__ == "__main__":
    main()