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

fi/evals/metrics/code_security/types.py

@@ -0,0 +1,534 @@
+"""
+Code Security Evaluation Types.
+
+Foundational types for AI-generated code security evaluation.
+Supports multiple evaluation modes, CWE-based vulnerability classification,
+and joint functional-security metrics.
+"""
+
+from enum import Enum
+from typing import Optional, List, Dict, Any
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class Severity(str, Enum):
+    """Vulnerability severity levels (CVSS-aligned)."""
+    CRITICAL = "critical"  # CVSS 9.0-10.0
+    HIGH = "high"          # CVSS 7.0-8.9
+    MEDIUM = "medium"      # CVSS 4.0-6.9
+    LOW = "low"            # CVSS 0.1-3.9
+    INFO = "info"          # Informational only
+
+
+class EvaluationMode(str, Enum):
+    """
+    AI code security evaluation modes.
+
+    Each mode tests a different aspect of AI code generation:
+    - INSTRUCT: Code generated from natural language instructions
+    - AUTOCOMPLETE: Code completion given partial context
+    - REPAIR: Fixing vulnerable code
+    - ADVERSARIAL: Resistance to security-bypassing prompts
+    """
+    INSTRUCT = "instruct"
+    AUTOCOMPLETE = "autocomplete"
+    REPAIR = "repair"
+    ADVERSARIAL = "adversarial"
+
+
+class VulnerabilityCategory(str, Enum):
+    """High-level vulnerability categories."""
+    INJECTION = "injection"
+    AUTHENTICATION = "authentication"
+    CRYPTOGRAPHY = "cryptography"
+    INPUT_VALIDATION = "input_validation"
+    SECRETS = "secrets"
+    MEMORY = "memory"
+    RESOURCE = "resource"
+    INFORMATION = "information"
+    SERIALIZATION = "serialization"
+    ACCESS_CONTROL = "access_control"
+
+
+# CWE to Category mapping
+CWE_CATEGORIES: Dict[str, VulnerabilityCategory] = {
+    # Injection
+    "CWE-89": VulnerabilityCategory.INJECTION,    # SQL Injection
+    "CWE-78": VulnerabilityCategory.INJECTION,    # OS Command Injection
+    "CWE-79": VulnerabilityCategory.INJECTION,    # XSS
+    "CWE-94": VulnerabilityCategory.INJECTION,    # Code Injection
+    "CWE-611": VulnerabilityCategory.INJECTION,   # XXE
+    "CWE-918": VulnerabilityCategory.INJECTION,   # SSRF
+
+    # Authentication
+    "CWE-287": VulnerabilityCategory.AUTHENTICATION,  # Improper Auth
+    "CWE-306": VulnerabilityCategory.AUTHENTICATION,  # Missing Auth
+    "CWE-352": VulnerabilityCategory.AUTHENTICATION,  # CSRF
+    "CWE-384": VulnerabilityCategory.AUTHENTICATION,  # Session Fixation
+
+    # Cryptography
+    "CWE-327": VulnerabilityCategory.CRYPTOGRAPHY,  # Weak Crypto
+    "CWE-328": VulnerabilityCategory.CRYPTOGRAPHY,  # Weak Hash
+    "CWE-330": VulnerabilityCategory.CRYPTOGRAPHY,  # Insufficient Randomness
+    "CWE-326": VulnerabilityCategory.CRYPTOGRAPHY,  # Inadequate Encryption
+
+    # Input Validation
+    "CWE-20": VulnerabilityCategory.INPUT_VALIDATION,   # Improper Input Validation
+    "CWE-22": VulnerabilityCategory.INPUT_VALIDATION,   # Path Traversal
+    "CWE-434": VulnerabilityCategory.INPUT_VALIDATION,  # Unrestricted Upload
+    "CWE-129": VulnerabilityCategory.INPUT_VALIDATION,  # Improper Array Index
+
+    # Secrets
+    "CWE-798": VulnerabilityCategory.SECRETS,  # Hardcoded Credentials
+    "CWE-259": VulnerabilityCategory.SECRETS,  # Hardcoded Password
+    "CWE-321": VulnerabilityCategory.SECRETS,  # Hardcoded Crypto Key
+    "CWE-532": VulnerabilityCategory.SECRETS,  # Log Sensitive Info
+
+    # Memory
+    "CWE-120": VulnerabilityCategory.MEMORY,  # Buffer Overflow
+    "CWE-476": VulnerabilityCategory.MEMORY,  # NULL Pointer Dereference
+    "CWE-190": VulnerabilityCategory.MEMORY,  # Integer Overflow
+    "CWE-787": VulnerabilityCategory.MEMORY,  # Out-of-bounds Write
+
+    # Resource
+    "CWE-400": VulnerabilityCategory.RESOURCE,  # Resource Exhaustion
+    "CWE-404": VulnerabilityCategory.RESOURCE,  # Improper Resource Shutdown
+    "CWE-772": VulnerabilityCategory.RESOURCE,  # Missing Resource Release
+
+    # Information
+    "CWE-209": VulnerabilityCategory.INFORMATION,  # Error Message Info Leak
+    "CWE-200": VulnerabilityCategory.INFORMATION,  # Sensitive Data Exposure
+
+    # Serialization
+    "CWE-502": VulnerabilityCategory.SERIALIZATION,  # Unsafe Deserialization
+}
+
+
+class CodeLocation(BaseModel):
+    """Location of a vulnerability in source code."""
+    model_config = ConfigDict(extra="allow")
+
+    line: int = Field(..., description="Line number (1-indexed)")
+    column: Optional[int] = Field(default=None, description="Column number")
+    end_line: Optional[int] = Field(default=None, description="End line for multi-line")
+    end_column: Optional[int] = Field(default=None, description="End column")
+    function: Optional[str] = Field(default=None, description="Enclosing function name")
+    snippet: Optional[str] = Field(default=None, description="Code snippet")
+
+
+class SecurityFinding(BaseModel):
+    """A single security vulnerability finding."""
+    model_config = ConfigDict(extra="allow")
+
+    cwe_id: str = Field(..., description="CWE identifier (e.g., 'CWE-89')")
+    vulnerability_type: str = Field(..., description="Human-readable type (e.g., 'SQL Injection')")
+    category: VulnerabilityCategory = Field(..., description="High-level category")
+    severity: Severity = Field(..., description="Severity level")
+    confidence: float = Field(
+        ...,
+        ge=0.0,
+        le=1.0,
+        description="Detection confidence (0.0-1.0)"
+    )
+    description: str = Field(..., description="Detailed description of the issue")
+    location: Optional[CodeLocation] = Field(default=None, description="Code location")
+    suggested_fix: Optional[str] = Field(default=None, description="Recommended fix")
+    references: Optional[List[str]] = Field(default=None, description="Reference URLs")
+
+    @property
+    def category_from_cwe(self) -> Optional[VulnerabilityCategory]:
+        """Get category from CWE mapping."""
+        return CWE_CATEGORIES.get(self.cwe_id)
+
+
+class FunctionalTestCase(BaseModel):
+    """A functional test case for func@k evaluation."""
+    model_config = ConfigDict(extra="allow")
+
+    input: Any = Field(..., description="Test input")
+    expected_output: Any = Field(..., description="Expected output")
+    description: Optional[str] = Field(default=None, description="Test description")
+
+
+# Alias for backward compatibility
+TestCase = FunctionalTestCase
+
+
+class CodeSecurityInput(BaseModel):
+    """
+    Input for code security evaluation.
+
+    Supports all evaluation modes:
+    - INSTRUCT: Provide instruction + generated code
+    - AUTOCOMPLETE: Provide code_prefix + generated completion
+    - REPAIR: Provide vulnerable code + fixed code
+    - ADVERSARIAL: Provide base + adversarial instruction + generated code
+    """
+    model_config = ConfigDict(extra="allow")
+
+    # Core fields
+    response: str = Field(..., description="The code to evaluate (generated code)")
+    language: str = Field(default="python", description="Programming language")
+
+    # Evaluation mode
+    mode: EvaluationMode = Field(
+        default=EvaluationMode.INSTRUCT,
+        description="Evaluation mode"
+    )
+
+    # Instruct mode fields
+    instruction: Optional[str] = Field(
+        default=None,
+        description="Natural language instruction (for INSTRUCT mode)"
+    )
+
+    # Autocomplete mode fields
+    code_prefix: Optional[str] = Field(
+        default=None,
+        description="Code before cursor (for AUTOCOMPLETE mode)"
+    )
+    code_suffix: Optional[str] = Field(
+        default=None,
+        description="Code after cursor (for AUTOCOMPLETE mode)"
+    )
+    cursor_line: Optional[int] = Field(
+        default=None,
+        description="Cursor line number (for AUTOCOMPLETE mode)"
+    )
+
+    # Repair mode fields
+    vulnerable_code: Optional[str] = Field(
+        default=None,
+        description="Original vulnerable code (for REPAIR mode)"
+    )
+    original_vulnerability: Optional[SecurityFinding] = Field(
+        default=None,
+        description="The vulnerability to fix (for REPAIR mode)"
+    )
+
+    # Adversarial mode fields
+    base_instruction: Optional[str] = Field(
+        default=None,
+        description="Normal instruction (for ADVERSARIAL mode)"
+    )
+    adversarial_instruction: Optional[str] = Field(
+        default=None,
+        description="Adversarial instruction (for ADVERSARIAL mode)"
+    )
+
+    # Functional testing (for func@k)
+    test_cases: Optional[List[TestCase]] = Field(
+        default=None,
+        description="Test cases for functional correctness evaluation"
+    )
+
+    # Reference implementation
+    expected_response: Optional[str] = Field(
+        default=None,
+        description="Reference secure implementation"
+    )
+
+    # Filtering options
+    include_categories: Optional[List[VulnerabilityCategory]] = Field(
+        default=None,
+        description="Only check these vulnerability categories"
+    )
+    exclude_cwes: Optional[List[str]] = Field(
+        default=None,
+        description="Skip these specific CWEs"
+    )
+    min_severity: Severity = Field(
+        default=Severity.LOW,
+        description="Minimum severity to report"
+    )
+    min_confidence: float = Field(
+        default=0.5,
+        ge=0.0,
+        le=1.0,
+        description="Minimum confidence to report"
+    )
+
+
+class CodeSecurityOutput(BaseModel):
+    """Output from code security evaluation."""
+    model_config = ConfigDict(extra="allow")
+
+    # Core scores
+    score: float = Field(
+        ...,
+        ge=0.0,
+        le=1.0,
+        description="Overall security score (1.0 = secure)"
+    )
+    passed: bool = Field(..., description="Whether code passes security threshold")
+
+    # Joint metrics (when test_cases provided)
+    func_score: Optional[float] = Field(
+        default=None,
+        description="Functional correctness score (func@k)"
+    )
+    sec_score: Optional[float] = Field(
+        default=None,
+        description="Security score (sec@k)"
+    )
+    func_sec_score: Optional[float] = Field(
+        default=None,
+        description="Joint score - both correct AND secure (func-sec@k)"
+    )
+
+    # Findings
+    findings: List[SecurityFinding] = Field(
+        default_factory=list,
+        description="List of vulnerability findings"
+    )
+    total_findings: int = Field(
+        default=0,
+        description="Total number of findings"
+    )
+    critical_count: int = Field(
+        default=0,
+        description="Number of critical-severity findings"
+    )
+    high_count: int = Field(
+        default=0,
+        description="Number of high-severity findings"
+    )
+
+    # Breakdowns
+    severity_counts: Dict[str, int] = Field(
+        default_factory=dict,
+        description="Count of findings per severity"
+    )
+    cwe_counts: Dict[str, int] = Field(
+        default_factory=dict,
+        description="Count of findings per CWE"
+    )
+    category_scores: Dict[str, float] = Field(
+        default_factory=dict,
+        description="Score per vulnerability category"
+    )
+
+    # Metadata
+    language: str = Field(default="python", description="Detected/specified language")
+    mode: EvaluationMode = Field(
+        default=EvaluationMode.INSTRUCT,
+        description="Evaluation mode used"
+    )
+    parse_errors: Optional[List[str]] = Field(
+        default=None,
+        description="Any parsing errors encountered"
+    )
+    detectors_run: List[str] = Field(
+        default_factory=list,
+        description="Names of detectors that were run"
+    )
+    runtime_ms: Optional[float] = Field(
+        default=None,
+        description="Evaluation runtime in milliseconds"
+    )
+
+    # Repair mode specific
+    is_fixed: Optional[bool] = Field(
+        default=None,
+        description="Was the vulnerability fixed? (REPAIR mode)"
+    )
+    new_vulnerabilities: Optional[List[SecurityFinding]] = Field(
+        default=None,
+        description="New vulnerabilities introduced (REPAIR mode)"
+    )
+
+    # Adversarial mode specific
+    resisted: Optional[bool] = Field(
+        default=None,
+        description="Did model resist adversarial prompt? (ADVERSARIAL mode)"
+    )
+    security_delta: Optional[float] = Field(
+        default=None,
+        description="Change in security vs normal prompt (ADVERSARIAL mode)"
+    )
+
+
+# Convenience type aliases
+Finding = SecurityFinding
+Location = CodeLocation
+Input = CodeSecurityInput
+Output = CodeSecurityOutput
+
+
+# Severity weights for scoring
+SEVERITY_WEIGHTS: Dict[Severity, float] = {
+    Severity.CRITICAL: 1.0,
+    Severity.HIGH: 0.8,
+    Severity.MEDIUM: 0.5,
+    Severity.LOW: 0.2,
+    Severity.INFO: 0.0,
+}
+
+
+# CWE metadata
+CWE_METADATA: Dict[str, Dict[str, Any]] = {
+    "CWE-89": {
+        "name": "SQL Injection",
+        "description": "Improper neutralization of special elements in SQL commands",
+        "default_severity": Severity.HIGH,
+        "owasp": "A03:2021",
+    },
+    "CWE-78": {
+        "name": "OS Command Injection",
+        "description": "Improper neutralization of special elements in OS commands",
+        "default_severity": Severity.CRITICAL,
+        "owasp": "A03:2021",
+    },
+    "CWE-79": {
+        "name": "Cross-Site Scripting (XSS)",
+        "description": "Improper neutralization of input during web page generation",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A03:2021",
+    },
+    "CWE-798": {
+        "name": "Hardcoded Credentials",
+        "description": "Use of hard-coded credentials",
+        "default_severity": Severity.HIGH,
+        "owasp": "A07:2021",
+    },
+    "CWE-22": {
+        "name": "Path Traversal",
+        "description": "Improper limitation of pathname to a restricted directory",
+        "default_severity": Severity.HIGH,
+        "owasp": "A01:2021",
+    },
+    "CWE-502": {
+        "name": "Deserialization of Untrusted Data",
+        "description": "Deserialization of untrusted data can result in code execution",
+        "default_severity": Severity.CRITICAL,
+        "owasp": "A08:2021",
+    },
+    "CWE-327": {
+        "name": "Broken Cryptographic Algorithm",
+        "description": "Use of a broken or risky cryptographic algorithm",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A02:2021",
+    },
+    "CWE-330": {
+        "name": "Insufficient Randomness",
+        "description": "Use of insufficiently random values",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A02:2021",
+    },
+    "CWE-287": {
+        "name": "Improper Authentication",
+        "description": "Improper authentication allowing unauthorized access",
+        "default_severity": Severity.HIGH,
+        "owasp": "A07:2021",
+    },
+    "CWE-306": {
+        "name": "Missing Authentication",
+        "description": "Missing authentication for critical function",
+        "default_severity": Severity.HIGH,
+        "owasp": "A07:2021",
+    },
+    "CWE-352": {
+        "name": "Cross-Site Request Forgery (CSRF)",
+        "description": "Cross-site request forgery vulnerability",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A01:2021",
+    },
+    "CWE-611": {
+        "name": "XML External Entity (XXE)",
+        "description": "Improper restriction of XML external entity reference",
+        "default_severity": Severity.HIGH,
+        "owasp": "A05:2021",
+    },
+    "CWE-918": {
+        "name": "Server-Side Request Forgery (SSRF)",
+        "description": "Server-side request forgery",
+        "default_severity": Severity.HIGH,
+        "owasp": "A10:2021",
+    },
+    "CWE-434": {
+        "name": "Unrestricted File Upload",
+        "description": "Unrestricted upload of file with dangerous type",
+        "default_severity": Severity.HIGH,
+        "owasp": "A04:2021",
+    },
+    "CWE-20": {
+        "name": "Improper Input Validation",
+        "description": "Improper input validation",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A03:2021",
+    },
+    "CWE-120": {
+        "name": "Buffer Overflow",
+        "description": "Buffer copy without checking size of input",
+        "default_severity": Severity.CRITICAL,
+        "owasp": "A06:2021",
+    },
+    "CWE-476": {
+        "name": "NULL Pointer Dereference",
+        "description": "NULL pointer dereference",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A06:2021",
+    },
+    "CWE-190": {
+        "name": "Integer Overflow",
+        "description": "Integer overflow or wraparound",
+        "default_severity": Severity.HIGH,
+        "owasp": "A06:2021",
+    },
+    "CWE-209": {
+        "name": "Information Exposure Through Error Message",
+        "description": "Generation of error message containing sensitive information",
+        "default_severity": Severity.LOW,
+        "owasp": "A04:2021",
+    },
+    "CWE-400": {
+        "name": "Resource Exhaustion",
+        "description": "Uncontrolled resource consumption",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A05:2021",
+    },
+    "CWE-404": {
+        "name": "Improper Resource Shutdown",
+        "description": "Improper resource shutdown or release",
+        "default_severity": Severity.LOW,
+        "owasp": "A05:2021",
+    },
+    "CWE-259": {
+        "name": "Hardcoded Password",
+        "description": "Use of hard-coded password",
+        "default_severity": Severity.HIGH,
+        "owasp": "A07:2021",
+    },
+    "CWE-321": {
+        "name": "Hardcoded Cryptographic Key",
+        "description": "Use of hard-coded cryptographic key",
+        "default_severity": Severity.HIGH,
+        "owasp": "A02:2021",
+    },
+    "CWE-532": {
+        "name": "Information Exposure Through Log Files",
+        "description": "Insertion of sensitive information into log file",
+        "default_severity": Severity.MEDIUM,
+        "owasp": "A09:2021",
+    },
+}
+
+
+def get_cwe_metadata(cwe_id: str) -> Dict[str, Any]:
+    """Get metadata for a CWE ID."""
+    return CWE_METADATA.get(cwe_id, {
+        "name": f"Unknown ({cwe_id})",
+        "description": "Unknown vulnerability type",
+        "default_severity": Severity.MEDIUM,
+    })
+
+
+def get_cwe_severity(cwe_id: str) -> Severity:
+    """Get default severity for a CWE ID."""
+    metadata = get_cwe_metadata(cwe_id)
+    return metadata.get("default_severity", Severity.MEDIUM)
+
+
+def get_cwe_category(cwe_id: str) -> VulnerabilityCategory:
+    """Get category for a CWE ID."""
+    return CWE_CATEGORIES.get(cwe_id, VulnerabilityCategory.INPUT_VALIDATION)

fi/evals/metrics/function_calling/__init__.py

@@ -0,0 +1,34 @@
+"""
+Function Calling Evaluation Metrics.
+
+AST-based evaluation of LLM function/tool calling capabilities.
+Provides deterministic, fast evaluation without LLM-as-judge.
+
+Based on BFCL (Berkeley Function Calling Leaderboard) methodology.
+"""
+
+from .types import (
+    FunctionCallInput,
+    FunctionCall,
+    FunctionDefinition,
+    ParameterSpec,
+)
+from .metrics import (
+    FunctionNameMatch,
+    ParameterValidation,
+    FunctionCallAccuracy,
+    FunctionCallExactMatch,
+)
+
+__all__ = [
+    # Types
+    "FunctionCallInput",
+    "FunctionCall",
+    "FunctionDefinition",
+    "ParameterSpec",
+    # Metrics
+    "FunctionNameMatch",
+    "ParameterValidation",
+    "FunctionCallAccuracy",
+    "FunctionCallExactMatch",
+]