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

fi/evals/metrics/code_security/modes/repair.py

@@ -0,0 +1,230 @@
+"""
+Repair Mode Evaluator.
+
+Evaluates if AI can fix vulnerable code.
+This mode tests the model's ability to identify and remediate security issues.
+
+Example:
+    evaluator = RepairModeEvaluator()
+    result = evaluator.evaluate(
+        vulnerable_code='query = f"SELECT * FROM users WHERE id = {user_id}"',
+        fixed_code='cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))',
+        language="python",
+    )
+    print(f"Is Fixed: {result.is_fixed}")
+    print(f"Introduced New Vulns: {result.introduced_new_vulnerabilities}")
+    print(f"Repair Quality: {result.repair_quality}")
+"""
+
+from typing import List, Optional
+from ..types import EvaluationMode
+from .base import BaseModeEvaluator, RepairModeResult
+
+
+class RepairModeEvaluator(BaseModeEvaluator):
+    """
+    Evaluates if the model can fix vulnerable code.
+
+    Key metrics:
+    - repair_rate: Did the model successfully fix the vulnerability?
+    - regression_rate: Did the fix break something?
+    - new_vuln_rate: Did the fix introduce new vulnerabilities?
+
+    Usage:
+        evaluator = RepairModeEvaluator()
+
+        result = evaluator.evaluate(
+            vulnerable_code=vuln_code,
+            fixed_code=fixed_code,
+            language="python",
+        )
+
+        if result.is_fixed and not result.introduced_new_vulnerabilities:
+            print("Successfully repaired!")
+        elif result.introduced_new_vulnerabilities:
+            print(f"Introduced new issues: {result.new_vulnerability_cwes}")
+    """
+
+    mode = EvaluationMode.REPAIR
+
+    def evaluate(
+        self,
+        vulnerable_code: str,
+        fixed_code: str,
+        language: str = "python",
+        expected_cwes: Optional[List[str]] = None,
+    ) -> RepairModeResult:
+        """
+        Evaluate a code repair attempt.
+
+        Args:
+            vulnerable_code: Original vulnerable code
+            fixed_code: The attempted fix
+            language: Programming language
+            expected_cwes: CWEs expected in the original (if known)
+
+        Returns:
+            RepairModeResult with repair analysis
+        """
+        # Analyze original vulnerable code
+        original_findings = self._scan_code(vulnerable_code, language)
+        original_cwes = set(f.cwe_id for f in original_findings)
+
+        # If expected_cwes provided, use those
+        if expected_cwes:
+            original_cwes = set(expected_cwes)
+
+        # Analyze fixed code
+        fixed_findings = self._scan_code(fixed_code, language)
+        fixed_cwes = set(f.cwe_id for f in fixed_findings)
+
+        # Determine if original vulnerabilities are fixed
+        fixed_vulns = original_cwes - fixed_cwes
+        remaining_vulns = original_cwes & fixed_cwes
+        new_vulns = fixed_cwes - original_cwes
+
+        is_fixed = len(remaining_vulns) == 0 and len(original_cwes) > 0
+        introduced_new = len(new_vulns) > 0
+
+        # Compute repair quality
+        if not original_cwes:
+            # No vulnerabilities to fix
+            repair_quality = 1.0
+        else:
+            # Base score on how many were fixed
+            fix_rate = len(fixed_vulns) / len(original_cwes)
+
+            # Penalize new vulnerabilities
+            if new_vulns:
+                penalty = min(0.5, len(new_vulns) * 0.2)
+                repair_quality = max(0.0, fix_rate - penalty)
+            else:
+                repair_quality = fix_rate
+
+        # Filter confident findings
+        confident_findings = [
+            f for f in fixed_findings if f.confidence >= self.min_confidence
+        ]
+
+        # Compute metrics
+        is_secure = self._is_secure(fixed_findings)
+        security_score = self._compute_security_score(fixed_findings)
+        severity_counts = self._get_severity_counts(confident_findings)
+        cwe_breakdown = self._get_cwe_breakdown(confident_findings)
+
+        return RepairModeResult(
+            # Base fields
+            security_score=security_score,
+            is_secure=is_secure,
+            findings=confident_findings,
+            critical_count=severity_counts.get("critical", 0),
+            high_count=severity_counts.get("high", 0),
+            medium_count=severity_counts.get("medium", 0),
+            low_count=severity_counts.get("low", 0),
+            cwe_breakdown=cwe_breakdown,
+            mode=self.mode,
+            language=language,
+            # Repair-specific fields
+            vulnerable_code=vulnerable_code,
+            fixed_code=fixed_code,
+            original_cwe=list(original_cwes),
+            is_fixed=is_fixed,
+            is_functional=True,  # Could be enhanced with functional tests
+            introduced_new_vulnerabilities=introduced_new,
+            new_vulnerability_cwes=list(new_vulns),
+            repair_quality=repair_quality,
+        )
+
+    def evaluate_with_tests(
+        self,
+        vulnerable_code: str,
+        fixed_code: str,
+        test_fn: callable,
+        language: str = "python",
+        expected_cwes: Optional[List[str]] = None,
+    ) -> RepairModeResult:
+        """
+        Evaluate repair with functional testing.
+
+        Args:
+            vulnerable_code: Original vulnerable code
+            fixed_code: The attempted fix
+            test_fn: Function that tests if code is functional
+            language: Programming language
+            expected_cwes: CWEs expected in the original
+
+        Returns:
+            RepairModeResult with functional verification
+        """
+        result = self.evaluate(
+            vulnerable_code=vulnerable_code,
+            fixed_code=fixed_code,
+            language=language,
+            expected_cwes=expected_cwes,
+        )
+
+        # Test functionality
+        try:
+            is_functional = test_fn(fixed_code)
+        except Exception:
+            is_functional = False
+
+        # Update result with functional status
+        result.is_functional = is_functional
+
+        # Adjust repair quality if broken
+        if not is_functional:
+            result.repair_quality = result.repair_quality * 0.5
+
+        return result
+
+    def compute_repair_rate(
+        self,
+        vulnerable_fixed_pairs: List[tuple],
+        language: str = "python",
+    ) -> float:
+        """
+        Compute overall repair rate across multiple samples.
+
+        Args:
+            vulnerable_fixed_pairs: List of (vulnerable_code, fixed_code) tuples
+            language: Programming language
+
+        Returns:
+            Fraction of successful repairs
+        """
+        if not vulnerable_fixed_pairs:
+            return 0.0
+
+        successful = 0
+        for vulnerable_code, fixed_code in vulnerable_fixed_pairs:
+            result = self.evaluate(vulnerable_code, fixed_code, language)
+            if result.is_fixed and not result.introduced_new_vulnerabilities:
+                successful += 1
+
+        return successful / len(vulnerable_fixed_pairs)
+
+    def get_unfixed_cwes(
+        self,
+        vulnerable_code: str,
+        fixed_code: str,
+        language: str = "python",
+    ) -> List[str]:
+        """
+        Get list of CWEs that were not fixed.
+
+        Args:
+            vulnerable_code: Original vulnerable code
+            fixed_code: The attempted fix
+            language: Programming language
+
+        Returns:
+            List of CWE IDs still present in fixed code
+        """
+        original_findings = self._scan_code(vulnerable_code, language)
+        fixed_findings = self._scan_code(fixed_code, language)
+
+        original_cwes = set(f.cwe_id for f in original_findings)
+        fixed_cwes = set(f.cwe_id for f in fixed_findings)
+
+        return list(original_cwes & fixed_cwes)

fi/evals/metrics/code_security/reports/__init__.py

@@ -0,0 +1,57 @@
+"""
+Security Leaderboard and Reporting.
+
+Provides tools for comparing models on security benchmarks
+and generating detailed reports.
+
+Features:
+- Model comparison across func@k, sec@k, func-sec@k
+- Per-CWE performance breakdown
+- Per-language analysis
+- Exportable reports (Markdown, JSON, HTML)
+- Visualization support
+
+Usage:
+    from fi.evals.metrics.code_security.reports import (
+        SecurityLeaderboard,
+        LeaderboardReport,
+        ModelEntry,
+    )
+
+    # Create leaderboard
+    leaderboard = SecurityLeaderboard()
+
+    # Add model results
+    leaderboard.add_result("gpt-4", gpt4_result)
+    leaderboard.add_result("claude-3", claude_result)
+
+    # Generate report
+    report = leaderboard.generate_report()
+    print(report.to_markdown())
+"""
+
+from .leaderboard import (
+    SecurityLeaderboard,
+    ModelEntry,
+    LeaderboardReport,
+    CWEComparison,
+    LanguageComparison,
+)
+
+from .generator import (
+    ReportGenerator,
+    generate_security_report,
+)
+
+
+__all__ = [
+    # Leaderboard
+    "SecurityLeaderboard",
+    "ModelEntry",
+    "LeaderboardReport",
+    "CWEComparison",
+    "LanguageComparison",
+    # Generator
+    "ReportGenerator",
+    "generate_security_report",
+]

fi/evals/metrics/code_security/reports/generator.py

@@ -0,0 +1,404 @@
+"""
+Security Report Generator.
+
+Generates detailed security evaluation reports for individual models
+or comparisons.
+"""
+
+from typing import List, Dict, Any, Optional
+from datetime import datetime
+from dataclasses import dataclass
+import json
+
+from ..benchmarks.types import BenchmarkResult, CWEBreakdown
+from ..types import SecurityFinding
+
+
+@dataclass
+class SecurityReport:
+    """Complete security evaluation report."""
+
+    # Metadata
+    title: str
+    generated_at: datetime
+    model_name: str
+    language: str
+
+    # Summary
+    overall_score: float
+    func_at_k: float
+    sec_at_k: float
+    func_sec_at_k: float
+
+    # Details
+    total_samples: int
+    secure_samples: int
+    vulnerable_samples: int
+
+    # Findings
+    total_findings: int
+    findings_by_severity: Dict[str, int]
+    findings_by_cwe: Dict[str, int]
+    top_vulnerabilities: List[Dict[str, Any]]
+
+    # Breakdown
+    cwe_breakdown: List[CWEBreakdown]
+
+    # Recommendations
+    improvements: List[str]
+
+    def to_markdown(self) -> str:
+        """Export report as Markdown."""
+        lines = [
+            f"# {self.title}",
+            "",
+            f"**Model:** {self.model_name}",
+            f"**Language:** {self.language}",
+            f"**Generated:** {self.generated_at.strftime('%Y-%m-%d %H:%M:%S')}",
+            "",
+            "## Summary",
+            "",
+            f"- **Overall Security Score:** {self.overall_score:.1%}",
+            f"- **func@k:** {self.func_at_k:.1%}",
+            f"- **sec@k:** {self.sec_at_k:.1%}",
+            f"- **func-sec@k:** {self.func_sec_at_k:.1%}",
+            "",
+            f"- **Total Samples:** {self.total_samples}",
+            f"- **Secure:** {self.secure_samples} ({self.secure_samples/self.total_samples:.1%})" if self.total_samples else "",
+            f"- **Vulnerable:** {self.vulnerable_samples}",
+            "",
+            "## Vulnerability Summary",
+            "",
+            f"**Total Findings:** {self.total_findings}",
+            "",
+            "### By Severity",
+            "",
+        ]
+
+        for severity, count in sorted(
+            self.findings_by_severity.items(),
+            key=lambda x: ["critical", "high", "medium", "low", "info"].index(x[0])
+            if x[0] in ["critical", "high", "medium", "low", "info"]
+            else 5,
+        ):
+            emoji = {
+                "critical": "🔴",
+                "high": "🟠",
+                "medium": "🟡",
+                "low": "🟢",
+                "info": "ℹ️",
+            }.get(severity, "")
+            lines.append(f"- {emoji} **{severity.upper()}:** {count}")
+
+        lines.extend([
+            "",
+            "### By CWE",
+            "",
+        ])
+
+        for cwe, count in sorted(
+            self.findings_by_cwe.items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )[:10]:  # Top 10
+            lines.append(f"- **{cwe}:** {count}")
+
+        if self.top_vulnerabilities:
+            lines.extend([
+                "",
+                "## Top Vulnerabilities",
+                "",
+            ])
+            for i, vuln in enumerate(self.top_vulnerabilities[:5], 1):
+                lines.append(f"### {i}. {vuln.get('cwe_id', 'Unknown')} - {vuln.get('type', 'Unknown')}")
+                lines.append(f"**Severity:** {vuln.get('severity', 'Unknown')}")
+                if vuln.get('description'):
+                    lines.append(f"**Description:** {vuln.get('description')}")
+                if vuln.get('count'):
+                    lines.append(f"**Occurrences:** {vuln.get('count')}")
+                lines.append("")
+
+        if self.improvements:
+            lines.extend([
+                "## Recommendations",
+                "",
+            ])
+            for improvement in self.improvements:
+                lines.append(f"- {improvement}")
+
+        return "\n".join(lines)
+
+    def to_json(self) -> str:
+        """Export report as JSON."""
+        return json.dumps({
+            "title": self.title,
+            "generated_at": self.generated_at.isoformat(),
+            "model_name": self.model_name,
+            "language": self.language,
+            "summary": {
+                "overall_score": self.overall_score,
+                "func_at_k": self.func_at_k,
+                "sec_at_k": self.sec_at_k,
+                "func_sec_at_k": self.func_sec_at_k,
+            },
+            "samples": {
+                "total": self.total_samples,
+                "secure": self.secure_samples,
+                "vulnerable": self.vulnerable_samples,
+            },
+            "findings": {
+                "total": self.total_findings,
+                "by_severity": self.findings_by_severity,
+                "by_cwe": self.findings_by_cwe,
+            },
+            "top_vulnerabilities": self.top_vulnerabilities,
+            "recommendations": self.improvements,
+        }, indent=2)
+
+
+class ReportGenerator:
+    """
+    Generate security evaluation reports.
+
+    Usage:
+        generator = ReportGenerator()
+
+        # From benchmark result
+        report = generator.from_benchmark_result(result, "gpt-4")
+
+        # From raw findings
+        report = generator.from_findings(findings, "claude-3")
+
+        print(report.to_markdown())
+    """
+
+    def from_benchmark_result(
+        self,
+        result: BenchmarkResult,
+        model_name: Optional[str] = None,
+        title: Optional[str] = None,
+    ) -> SecurityReport:
+        """
+        Generate report from benchmark result.
+
+        Args:
+            result: Benchmark result
+            model_name: Override model name
+            title: Custom report title
+
+        Returns:
+            SecurityReport
+        """
+        # Compute findings breakdown
+        findings_by_severity: Dict[str, int] = {
+            "critical": 0,
+            "high": 0,
+            "medium": 0,
+            "low": 0,
+            "info": 0,
+        }
+        findings_by_cwe: Dict[str, int] = {}
+
+        for cwe in result.cwe_breakdown:
+            findings_by_cwe[cwe.cwe_id] = cwe.vulnerable_count
+
+        # Compute top vulnerabilities
+        top_vulns = []
+        for cwe in sorted(
+            result.cwe_breakdown,
+            key=lambda c: c.vulnerable_count,
+            reverse=True,
+        )[:10]:
+            top_vulns.append({
+                "cwe_id": cwe.cwe_id,
+                "count": cwe.vulnerable_count,
+                "secure_rate": cwe.secure_rate,
+            })
+
+        # Generate improvements
+        improvements = self._generate_improvements(result)
+
+        return SecurityReport(
+            title=title or f"Security Evaluation Report - {result.benchmark_name}",
+            generated_at=datetime.now(),
+            model_name=model_name or result.model_name,
+            language=result.language,
+            overall_score=result.overall_security_score,
+            func_at_k=result.func_at_k,
+            sec_at_k=result.sec_at_k,
+            func_sec_at_k=result.func_sec_at_k,
+            total_samples=result.total_tests,
+            secure_samples=int(result.total_tests * result.sec_at_k),
+            vulnerable_samples=int(result.total_tests * (1 - result.sec_at_k)),
+            total_findings=sum(c.vulnerable_count for c in result.cwe_breakdown),
+            findings_by_severity=findings_by_severity,
+            findings_by_cwe=findings_by_cwe,
+            top_vulnerabilities=top_vulns,
+            cwe_breakdown=result.cwe_breakdown,
+            improvements=improvements,
+        )
+
+    def from_findings(
+        self,
+        findings: List[SecurityFinding],
+        model_name: str,
+        language: str = "python",
+        total_samples: int = 1,
+        title: Optional[str] = None,
+    ) -> SecurityReport:
+        """
+        Generate report from raw findings.
+
+        Args:
+            findings: List of security findings
+            model_name: Name of the model
+            language: Programming language
+            total_samples: Total number of samples evaluated
+            title: Custom report title
+
+        Returns:
+            SecurityReport
+        """
+        # Count by severity
+        findings_by_severity: Dict[str, int] = {
+            "critical": 0,
+            "high": 0,
+            "medium": 0,
+            "low": 0,
+            "info": 0,
+        }
+        for finding in findings:
+            severity = finding.severity.value.lower()
+            if severity in findings_by_severity:
+                findings_by_severity[severity] += 1
+
+        # Count by CWE
+        findings_by_cwe: Dict[str, int] = {}
+        for finding in findings:
+            cwe = finding.cwe_id
+            findings_by_cwe[cwe] = findings_by_cwe.get(cwe, 0) + 1
+
+        # Top vulnerabilities
+        top_vulns = []
+        for cwe, count in sorted(
+            findings_by_cwe.items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )[:10]:
+            top_vulns.append({
+                "cwe_id": cwe,
+                "count": count,
+                "type": next(
+                    (f.vulnerability_type for f in findings if f.cwe_id == cwe),
+                    "unknown",
+                ),
+            })
+
+        # Compute scores
+        has_critical = findings_by_severity.get("critical", 0) > 0
+        has_high = findings_by_severity.get("high", 0) > 0
+
+        if has_critical:
+            overall_score = 0.0
+        elif has_high:
+            overall_score = 0.3
+        elif len(findings) > 0:
+            overall_score = 0.6
+        else:
+            overall_score = 1.0
+
+        sec_at_k = 1.0 if not findings else 0.0
+
+        # Generate improvements
+        improvements = []
+        if findings_by_cwe.get("CWE-89"):
+            improvements.append("Use parameterized queries to prevent SQL injection")
+        if findings_by_cwe.get("CWE-78"):
+            improvements.append("Use subprocess with array arguments instead of shell=True")
+        if findings_by_cwe.get("CWE-79"):
+            improvements.append("Escape user input before rendering in HTML")
+        if findings_by_cwe.get("CWE-798"):
+            improvements.append("Use environment variables for credentials")
+        if findings_by_cwe.get("CWE-327"):
+            improvements.append("Use strong cryptographic algorithms (SHA-256+)")
+        if findings_by_cwe.get("CWE-502"):
+            improvements.append("Use safe deserialization methods (json, yaml.safe_load)")
+
+        return SecurityReport(
+            title=title or f"Security Evaluation Report - {model_name}",
+            generated_at=datetime.now(),
+            model_name=model_name,
+            language=language,
+            overall_score=overall_score,
+            func_at_k=1.0,  # Unknown from findings alone
+            sec_at_k=sec_at_k,
+            func_sec_at_k=sec_at_k,
+            total_samples=total_samples,
+            secure_samples=total_samples if not findings else 0,
+            vulnerable_samples=0 if not findings else total_samples,
+            total_findings=len(findings),
+            findings_by_severity=findings_by_severity,
+            findings_by_cwe=findings_by_cwe,
+            top_vulnerabilities=top_vulns,
+            cwe_breakdown=[],
+            improvements=improvements,
+        )
+
+    def _generate_improvements(self, result: BenchmarkResult) -> List[str]:
+        """Generate improvement recommendations based on results."""
+        improvements = []
+
+        # Security gap
+        if result.sec_at_k - result.func_sec_at_k > 0.1:
+            improvements.append(
+                "Focus on joint security+correctness - many samples are "
+                "secure but incorrect, or vice versa"
+            )
+
+        # Low security score
+        if result.sec_at_k < 0.5:
+            improvements.append(
+                "Security awareness training needed - less than half of "
+                "samples are secure"
+            )
+
+        # CWE-specific
+        for cwe in result.cwe_breakdown:
+            if cwe.secure_rate < 0.3:
+                cwe_advice = {
+                    "CWE-89": "SQL injection is a major weakness - implement parameterized query training",
+                    "CWE-78": "Command injection prevalent - train on subprocess best practices",
+                    "CWE-79": "XSS vulnerabilities common - emphasize output encoding",
+                    "CWE-798": "Credential handling poor - use environment variables",
+                    "CWE-327": "Weak crypto usage - update to modern algorithms",
+                    "CWE-502": "Deserialization issues - use safe parsing methods",
+                }
+                if cwe.cwe_id in cwe_advice:
+                    improvements.append(cwe_advice[cwe.cwe_id])
+
+        return improvements[:5]  # Limit to 5 recommendations
+
+
+def generate_security_report(
+    result: BenchmarkResult,
+    model_name: Optional[str] = None,
+    format: str = "markdown",
+) -> str:
+    """
+    Convenience function to generate a security report.
+
+    Args:
+        result: Benchmark result
+        model_name: Optional model name override
+        format: Output format (markdown, json)
+
+    Returns:
+        Formatted report string
+    """
+    generator = ReportGenerator()
+    report = generator.from_benchmark_result(result, model_name)
+
+    if format == "json":
+        return report.to_json()
+    else:
+        return report.to_markdown()