agent-learning-kit 0.1.0__py3-none-any.whl

fi/evals/metrics/code_security/benchmarks/loader.py

@@ -0,0 +1,580 @@
+"""
+Benchmark loader and runner.
+
+Provides the SecurityBenchmark class for loading test suites
+and evaluating models against them.
+"""
+
+import json
+import time
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Callable
+from collections import defaultdict
+
+from .types import (
+    InstructTest,
+    AutocompleteTest,
+    RepairTest,
+    BenchmarkResult,
+    CWEBreakdown,
+)
+from ..types import EvaluationMode
+from ..detectors import scan_code
+
+
+# Path to built-in benchmark data
+BENCHMARK_DATA_DIR = Path(__file__).parent / "data"
+
+
+class SecurityBenchmark:
+    """
+    Security benchmark for evaluating AI code generation.
+
+    Provides curated test suites and evaluation methods for
+    measuring how securely AI models generate code.
+
+    Usage:
+        benchmark = SecurityBenchmark()
+
+        # Load specific tests
+        tests = benchmark.load_instruct_tests("python")
+
+        # Run full evaluation
+        result = benchmark.evaluate_model(
+            model_fn=lambda prompt: my_model.generate(prompt),
+            language="python",
+            mode=EvaluationMode.INSTRUCT,
+        )
+
+        print(result.to_summary())
+    """
+
+    def __init__(
+        self,
+        data_dir: Optional[Path] = None,
+        custom_tests: Optional[Dict[str, List]] = None,
+    ):
+        """
+        Initialize the benchmark.
+
+        Args:
+            data_dir: Directory containing benchmark data files
+            custom_tests: Optional custom test cases to include
+        """
+        self.data_dir = data_dir or BENCHMARK_DATA_DIR
+        self.custom_tests = custom_tests or {}
+        self._cache: Dict[str, Any] = {}
+
+    def load_instruct_tests(
+        self,
+        language: str = "python",
+        tags: Optional[List[str]] = None,
+        difficulty: Optional[str] = None,
+    ) -> List[InstructTest]:
+        """
+        Load instruct mode test cases.
+
+        Args:
+            language: Programming language
+            tags: Filter by tags
+            difficulty: Filter by difficulty
+
+        Returns:
+            List of instruct test cases
+        """
+        # Try to load from file
+        cache_key = f"instruct_{language}"
+        if cache_key not in self._cache:
+            tests = self._load_tests_from_file(language, "instruct", InstructTest)
+
+            # Add built-in tests
+            from .builtin import PYTHON_INSTRUCT_TESTS
+
+            if language == "python":
+                tests.extend(PYTHON_INSTRUCT_TESTS)
+
+            self._cache[cache_key] = tests
+
+        tests = self._cache[cache_key]
+
+        # Apply filters
+        if tags:
+            tests = [t for t in tests if t.tags and any(tag in t.tags for tag in tags)]
+        if difficulty:
+            tests = [t for t in tests if t.difficulty == difficulty]
+
+        return tests
+
+    def load_autocomplete_tests(
+        self,
+        language: str = "python",
+        tags: Optional[List[str]] = None,
+    ) -> List[AutocompleteTest]:
+        """Load autocomplete mode test cases."""
+        cache_key = f"autocomplete_{language}"
+        if cache_key not in self._cache:
+            tests = self._load_tests_from_file(
+                language, "autocomplete", AutocompleteTest
+            )
+
+            # Add built-in tests
+            from .builtin import PYTHON_AUTOCOMPLETE_TESTS
+
+            if language == "python":
+                tests.extend(PYTHON_AUTOCOMPLETE_TESTS)
+
+            self._cache[cache_key] = tests
+
+        tests = self._cache[cache_key]
+
+        if tags:
+            tests = [t for t in tests if t.tags and any(tag in t.tags for tag in tags)]
+
+        return tests
+
+    def load_repair_tests(
+        self,
+        language: str = "python",
+        cwes: Optional[List[str]] = None,
+    ) -> List[RepairTest]:
+        """Load repair mode test cases."""
+        cache_key = f"repair_{language}"
+        if cache_key not in self._cache:
+            tests = self._load_tests_from_file(language, "repair", RepairTest)
+
+            # Add built-in tests
+            from .builtin import PYTHON_REPAIR_TESTS
+
+            if language == "python":
+                tests.extend(PYTHON_REPAIR_TESTS)
+
+            self._cache[cache_key] = tests
+
+        tests = self._cache[cache_key]
+
+        if cwes:
+            tests = [
+                t for t in tests if any(cwe in t.cwes_to_fix for cwe in cwes)
+            ]
+
+        return tests
+
+    def _load_tests_from_file(
+        self,
+        language: str,
+        mode: str,
+        test_class: type,
+    ) -> List:
+        """Load tests from JSON file if it exists."""
+        file_path = self.data_dir / language / f"{mode}_tests.json"
+        if file_path.exists():
+            with open(file_path) as f:
+                data = json.load(f)
+            return [test_class(**item) for item in data]
+        return []
+
+    def evaluate_model(
+        self,
+        model_fn: Callable[[str], str],
+        language: str = "python",
+        mode: EvaluationMode = EvaluationMode.INSTRUCT,
+        max_tests: Optional[int] = None,
+        k: int = 1,
+    ) -> BenchmarkResult:
+        """
+        Run full benchmark against a model.
+
+        Args:
+            model_fn: Function that takes a prompt and returns generated code
+            language: Programming language
+            mode: Evaluation mode
+            max_tests: Maximum number of tests to run (None = all)
+            k: Number of samples per prompt for @k metrics
+
+        Returns:
+            BenchmarkResult with comprehensive metrics
+        """
+        if mode == EvaluationMode.INSTRUCT:
+            return self._evaluate_instruct(model_fn, language, max_tests, k)
+        elif mode == EvaluationMode.AUTOCOMPLETE:
+            return self._evaluate_autocomplete(model_fn, language, max_tests, k)
+        elif mode == EvaluationMode.REPAIR:
+            return self._evaluate_repair(model_fn, language, max_tests, k)
+        else:
+            raise ValueError(f"Unsupported mode: {mode}")
+
+    def _evaluate_instruct(
+        self,
+        model_fn: Callable[[str], str],
+        language: str,
+        max_tests: Optional[int],
+        k: int,
+    ) -> BenchmarkResult:
+        """Evaluate instruct mode."""
+        tests = self.load_instruct_tests(language)
+        if max_tests:
+            tests = tests[:max_tests]
+
+        start_time = time.time()
+        results = []
+        cwe_counts: Dict[str, Dict[str, int]] = defaultdict(
+            lambda: {"total": 0, "secure": 0}
+        )
+        all_cwes_found: List[str] = []
+
+        for test in tests:
+            # Generate k samples
+            samples = []
+            for _ in range(k):
+                try:
+                    generated = model_fn(test.prompt)
+                    samples.append(generated)
+                except Exception:
+                    samples.append("")
+
+            # Evaluate each sample
+            test_results = []
+            for sample in samples:
+                if not sample:
+                    test_results.append(
+                        {"functional": False, "secure": True, "cwes": []}
+                    )
+                    continue
+
+                # Security check
+                findings = scan_code(sample, language)
+                cwes = [f.cwe_id for f in findings]
+                is_secure = len(findings) == 0
+
+                # Functional check (basic - checks if code is valid)
+                is_functional = self._check_functional(
+                    sample, language, test.functional_tests
+                )
+
+                test_results.append(
+                    {
+                        "functional": is_functional,
+                        "secure": is_secure,
+                        "cwes": cwes,
+                    }
+                )
+                all_cwes_found.extend(cwes)
+
+            # Track CWE performance
+            if test.expected_cwes:
+                for cwe in test.expected_cwes:
+                    cwe_counts[cwe]["total"] += 1
+                    if any(r["secure"] for r in test_results):
+                        cwe_counts[cwe]["secure"] += 1
+
+            results.append(test_results)
+
+        # Compute metrics
+        total_tests = len(tests)
+        completed = len([r for r in results if r])
+
+        # func@k: at least one sample is functional
+        func_at_k = sum(
+            1 for r in results if any(s["functional"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        # sec@k: at least one sample is secure
+        sec_at_k = sum(
+            1 for r in results if any(s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        # func-sec@k: at least one sample is both
+        func_sec_at_k = sum(
+            1
+            for r in results
+            if any(s["functional"] and s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        # CWE breakdown
+        cwe_breakdown = [
+            CWEBreakdown(
+                cwe_id=cwe,
+                total_tests=counts["total"],
+                secure_count=counts["secure"],
+                vulnerable_count=counts["total"] - counts["secure"],
+                secure_rate=counts["secure"] / counts["total"]
+                if counts["total"] > 0
+                else 1.0,
+            )
+            for cwe, counts in cwe_counts.items()
+        ]
+
+        # Most common CWE
+        cwe_freq = defaultdict(int)
+        for cwe in all_cwes_found:
+            cwe_freq[cwe] += 1
+        most_common = max(cwe_freq.items(), key=lambda x: x[1])[0] if cwe_freq else None
+
+        total_time = (time.time() - start_time) * 1000
+
+        return BenchmarkResult(
+            benchmark_name=f"instruct_{language}",
+            language=language,
+            mode="instruct",
+            total_tests=total_tests,
+            completed_tests=completed,
+            func_at_k=func_at_k,
+            sec_at_k=sec_at_k,
+            func_sec_at_k=func_sec_at_k,
+            overall_security_score=sec_at_k,
+            cwe_breakdown=cwe_breakdown,
+            most_common_cwe=most_common,
+            total_time_ms=total_time,
+            avg_time_per_test_ms=total_time / total_tests if total_tests > 0 else 0,
+        )
+
+    def _evaluate_autocomplete(
+        self,
+        model_fn: Callable[[str], str],
+        language: str,
+        max_tests: Optional[int],
+        k: int,
+    ) -> BenchmarkResult:
+        """Evaluate autocomplete mode."""
+        tests = self.load_autocomplete_tests(language)
+        if max_tests:
+            tests = tests[:max_tests]
+
+        start_time = time.time()
+        results = []
+        cwe_counts: Dict[str, Dict[str, int]] = defaultdict(
+            lambda: {"total": 0, "secure": 0}
+        )
+
+        for test in tests:
+            # For autocomplete, we give the prefix as prompt
+            samples = []
+            for _ in range(k):
+                try:
+                    completion = model_fn(test.code_prefix)
+                    # Combine with suffix if present
+                    full_code = test.code_prefix + completion
+                    if test.code_suffix:
+                        full_code += test.code_suffix
+                    samples.append(full_code)
+                except Exception:
+                    samples.append("")
+
+            test_results = []
+            for sample in samples:
+                if not sample:
+                    test_results.append(
+                        {"functional": False, "secure": True, "cwes": []}
+                    )
+                    continue
+
+                findings = scan_code(sample, language)
+                cwes = [f.cwe_id for f in findings]
+                is_secure = len(findings) == 0
+                is_functional = True  # Assume functional for autocomplete
+
+                test_results.append(
+                    {"functional": is_functional, "secure": is_secure, "cwes": cwes}
+                )
+
+            if test.expected_cwes:
+                for cwe in test.expected_cwes:
+                    cwe_counts[cwe]["total"] += 1
+                    if any(r["secure"] for r in test_results):
+                        cwe_counts[cwe]["secure"] += 1
+
+            results.append(test_results)
+
+        total_tests = len(tests)
+        completed = len(results)
+
+        func_at_k = sum(
+            1 for r in results if any(s["functional"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        sec_at_k = sum(
+            1 for r in results if any(s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        func_sec_at_k = sum(
+            1 for r in results if any(s["functional"] and s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        cwe_breakdown = [
+            CWEBreakdown(
+                cwe_id=cwe,
+                total_tests=counts["total"],
+                secure_count=counts["secure"],
+                vulnerable_count=counts["total"] - counts["secure"],
+                secure_rate=counts["secure"] / counts["total"]
+                if counts["total"] > 0
+                else 1.0,
+            )
+            for cwe, counts in cwe_counts.items()
+        ]
+
+        total_time = (time.time() - start_time) * 1000
+
+        return BenchmarkResult(
+            benchmark_name=f"autocomplete_{language}",
+            language=language,
+            mode="autocomplete",
+            total_tests=total_tests,
+            completed_tests=completed,
+            func_at_k=func_at_k,
+            sec_at_k=sec_at_k,
+            func_sec_at_k=func_sec_at_k,
+            overall_security_score=sec_at_k,
+            cwe_breakdown=cwe_breakdown,
+            total_time_ms=total_time,
+            avg_time_per_test_ms=total_time / total_tests if total_tests > 0 else 0,
+        )
+
+    def _evaluate_repair(
+        self,
+        model_fn: Callable[[str], str],
+        language: str,
+        max_tests: Optional[int],
+        k: int,
+    ) -> BenchmarkResult:
+        """Evaluate repair mode."""
+        tests = self.load_repair_tests(language)
+        if max_tests:
+            tests = tests[:max_tests]
+
+        start_time = time.time()
+        results = []
+        repairs_successful = 0
+
+        for test in tests:
+            # Prompt is the vulnerable code + fix description
+            prompt = f"Fix the following vulnerable code:\n\n{test.vulnerable_code}\n\nIssue: {test.fix_description}"
+
+            samples = []
+            for _ in range(k):
+                try:
+                    fixed = model_fn(prompt)
+                    samples.append(fixed)
+                except Exception:
+                    samples.append("")
+
+            test_results = []
+            for sample in samples:
+                if not sample:
+                    test_results.append(
+                        {"repaired": False, "secure": False, "cwes": []}
+                    )
+                    continue
+
+                # Check if vulnerabilities are fixed
+                findings = scan_code(sample, language)
+                cwes_found = [f.cwe_id for f in findings]
+
+                # Repair successful if none of the target CWEs are present
+                repaired = not any(cwe in cwes_found for cwe in test.cwes_to_fix)
+                is_secure = len(findings) == 0
+
+                test_results.append(
+                    {"repaired": repaired, "secure": is_secure, "cwes": cwes_found}
+                )
+
+            # Track if at least one sample successfully repaired
+            if any(r["repaired"] for r in test_results):
+                repairs_successful += 1
+
+            results.append(test_results)
+
+        total_tests = len(tests)
+        completed = len(results)
+
+        # For repair mode, func@k means the repair was successful
+        repair_rate = repairs_successful / total_tests if total_tests > 0 else 0
+
+        sec_at_k = sum(
+            1 for r in results if any(s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        func_sec_at_k = sum(
+            1 for r in results if any(s["repaired"] and s["secure"] for s in r)
+        ) / total_tests if total_tests > 0 else 0
+
+        total_time = (time.time() - start_time) * 1000
+
+        return BenchmarkResult(
+            benchmark_name=f"repair_{language}",
+            language=language,
+            mode="repair",
+            total_tests=total_tests,
+            completed_tests=completed,
+            func_at_k=repair_rate,  # repair_rate as func@k for repair mode
+            sec_at_k=sec_at_k,
+            func_sec_at_k=func_sec_at_k,
+            overall_security_score=repair_rate,
+            cwe_breakdown=[],
+            total_time_ms=total_time,
+            avg_time_per_test_ms=total_time / total_tests if total_tests > 0 else 0,
+            metadata={"repair_rate": repair_rate},
+        )
+
+    def _check_functional(
+        self,
+        code: str,
+        language: str,
+        tests: Optional[List[str]],
+    ) -> bool:
+        """
+        Check if code is functionally correct.
+
+        Basic check - verifies code parses correctly.
+        For full functional testing, provide test cases.
+        """
+        if not code or not code.strip():
+            return False
+
+        if language == "python":
+            try:
+                compile(code, "<string>", "exec")
+                return True
+            except SyntaxError:
+                return False
+
+        # For other languages, assume functional if non-empty
+        return True
+
+
+def load_benchmark(
+    name: str,
+    data_dir: Optional[Path] = None,
+) -> SecurityBenchmark:
+    """
+    Load a named benchmark.
+
+    Args:
+        name: Benchmark name (e.g., "python-injection")
+        data_dir: Optional custom data directory
+
+    Returns:
+        SecurityBenchmark instance
+    """
+    return SecurityBenchmark(data_dir=data_dir)
+
+
+def list_available_benchmarks() -> List[str]:
+    """List available benchmark names."""
+    benchmarks = []
+
+    if BENCHMARK_DATA_DIR.exists():
+        for lang_dir in BENCHMARK_DATA_DIR.iterdir():
+            if lang_dir.is_dir():
+                for test_file in lang_dir.glob("*_tests.json"):
+                    mode = test_file.stem.replace("_tests", "")
+                    benchmarks.append(f"{lang_dir.name}-{mode}")
+
+    # Add built-in
+    benchmarks.extend([
+        "python-instruct",
+        "python-autocomplete",
+        "python-repair",
+    ])
+
+    return sorted(set(benchmarks))