raxe 0.4.6__py3-none-any.whl

raxe/domain/ml/voting/head_voters.py

@@ -0,0 +1,378 @@
+"""Individual head voting logic - Pure functions for per-head voting.
+
+This module contains ONLY pure functions - no I/O operations.
+Each function takes head output data and configuration, returns a vote.
+
+Head Voters:
+- vote_binary: Binary (is_threat) head voting
+- vote_family: Family head voting
+- vote_severity: Severity head voting
+- vote_technique: Technique head voting
+- vote_harm: Harm types head voting
+"""
+
+from raxe.domain.ml.voting.config import (
+    BinaryHeadThresholds,
+    FamilyHeadThresholds,
+    HarmHeadThresholds,
+    HeadWeights,
+    SeverityHeadThresholds,
+    TechniqueHeadThresholds,
+)
+from raxe.domain.ml.voting.models import HeadVoteDetail, Vote
+
+
+def vote_binary(
+    threat_probability: float,
+    safe_probability: float,
+    thresholds: BinaryHeadThresholds,
+    weight: float,
+) -> HeadVoteDetail:
+    """Cast vote for the binary (is_threat) classifier head.
+
+    Voting rules:
+    - THREAT if threat_probability >= threat_threshold
+    - SAFE if threat_probability < safe_threshold
+    - ABSTAIN otherwise (gray zone)
+
+    Args:
+        threat_probability: Model output probability of threat (0.0 to 1.0)
+        safe_probability: Model output probability of safe (0.0 to 1.0)
+        thresholds: Threshold configuration for this head
+        weight: Vote weight for this head
+
+    Returns:
+        HeadVoteDetail with the vote, confidence, and rationale
+
+    Examples:
+        >>> detail = vote_binary(0.80, 0.20, BinaryHeadThresholds(), 1.0)
+        >>> detail.vote
+        <Vote.THREAT: 'threat'>
+        >>> detail.confidence
+        0.8
+    """
+    if threat_probability >= thresholds.threat_threshold:
+        vote = Vote.THREAT
+        confidence = threat_probability
+        threshold_used = thresholds.threat_threshold
+        rationale = (
+            f"threat_probability ({threat_probability:.2%}) >= "
+            f"threat_threshold ({thresholds.threat_threshold:.2%})"
+        )
+    elif threat_probability < thresholds.safe_threshold:
+        vote = Vote.SAFE
+        confidence = safe_probability
+        threshold_used = thresholds.safe_threshold
+        rationale = (
+            f"threat_probability ({threat_probability:.2%}) < "
+            f"safe_threshold ({thresholds.safe_threshold:.2%})"
+        )
+    else:
+        vote = Vote.ABSTAIN
+        # Confidence for abstain is how close to the decision boundary
+        confidence = 1.0 - abs(
+            threat_probability - (thresholds.threat_threshold + thresholds.safe_threshold) / 2
+        ) / ((thresholds.threat_threshold - thresholds.safe_threshold) / 2)
+        confidence = max(0.0, min(1.0, confidence))
+        threshold_used = (thresholds.threat_threshold + thresholds.safe_threshold) / 2
+        rationale = (
+            f"threat_probability ({threat_probability:.2%}) in gray zone "
+            f"[{thresholds.safe_threshold:.2%}, {thresholds.threat_threshold:.2%})"
+        )
+
+    return HeadVoteDetail(
+        head_name="binary",
+        vote=vote,
+        confidence=confidence,
+        weight=weight,
+        raw_probability=threat_probability,
+        threshold_used=threshold_used,
+        prediction="threat" if threat_probability >= 0.5 else "safe",
+        rationale=rationale,
+    )
+
+
+def vote_family(
+    family_prediction: str,
+    family_confidence: float,
+    thresholds: FamilyHeadThresholds,
+    weight: float,
+) -> HeadVoteDetail:
+    """Cast vote for the family classifier head.
+
+    Voting rules:
+    - THREAT if family != benign AND confidence >= threat_confidence
+    - SAFE if family == benign OR confidence < safe_confidence
+    - ABSTAIN otherwise
+
+    Args:
+        family_prediction: Predicted threat family (e.g., "benign", "jailbreak")
+        family_confidence: Model confidence in the prediction (0.0 to 1.0)
+        thresholds: Threshold configuration for this head
+        weight: Vote weight for this head
+
+    Returns:
+        HeadVoteDetail with the vote, confidence, and rationale
+
+    Examples:
+        >>> detail = vote_family("jailbreak", 0.75, FamilyHeadThresholds(), 1.2)
+        >>> detail.vote
+        <Vote.THREAT: 'threat'>
+    """
+    family_is_benign = family_prediction.lower() == "benign"
+
+    if family_is_benign:
+        vote = Vote.SAFE
+        confidence = family_confidence
+        threshold_used = 0.0  # No threshold for benign
+        rationale = f"family={family_prediction} is benign"
+    elif family_confidence >= thresholds.threat_confidence:
+        vote = Vote.THREAT
+        confidence = family_confidence
+        threshold_used = thresholds.threat_confidence
+        rationale = (
+            f"family={family_prediction} with confidence ({family_confidence:.2%}) >= "
+            f"threat_confidence ({thresholds.threat_confidence:.2%})"
+        )
+    elif family_confidence < thresholds.safe_confidence:
+        vote = Vote.SAFE
+        confidence = 1.0 - family_confidence  # Low confidence in threat = high confidence in safe
+        threshold_used = thresholds.safe_confidence
+        rationale = (
+            f"family={family_prediction} with confidence ({family_confidence:.2%}) < "
+            f"safe_confidence ({thresholds.safe_confidence:.2%})"
+        )
+    else:
+        vote = Vote.ABSTAIN
+        # Confidence for abstain reflects uncertainty
+        confidence = 0.5
+        threshold_used = (thresholds.threat_confidence + thresholds.safe_confidence) / 2
+        rationale = (
+            f"family={family_prediction} with confidence ({family_confidence:.2%}) in gray zone "
+            f"[{thresholds.safe_confidence:.2%}, {thresholds.threat_confidence:.2%})"
+        )
+
+    return HeadVoteDetail(
+        head_name="family",
+        vote=vote,
+        confidence=confidence,
+        weight=weight,
+        raw_probability=family_confidence,
+        threshold_used=threshold_used,
+        prediction=family_prediction,
+        rationale=rationale,
+    )
+
+
+def vote_severity(
+    severity_prediction: str,
+    severity_confidence: float,
+    thresholds: SeverityHeadThresholds,
+    weight: float,
+) -> HeadVoteDetail:
+    """Cast vote for the severity classifier head.
+
+    Voting rules:
+    - THREAT if severity in threat_severities (low, medium, high, critical)
+    - SAFE if severity in safe_severities (none)
+    - No abstain - severity always has an opinion
+
+    Note: Severity head has the highest weight (1.5) because it directly
+    indicates threat severity and has strong signal quality.
+
+    Args:
+        severity_prediction: Predicted severity (e.g., "none", "high", "critical")
+        severity_confidence: Model confidence in the prediction (0.0 to 1.0)
+        thresholds: Threshold configuration for this head
+        weight: Vote weight for this head
+
+    Returns:
+        HeadVoteDetail with the vote, confidence, and rationale
+
+    Examples:
+        >>> detail = vote_severity("high", 0.85, SeverityHeadThresholds(), 1.5)
+        >>> detail.vote
+        <Vote.THREAT: 'threat'>
+    """
+    severity_lower = severity_prediction.lower()
+
+    if severity_lower in thresholds.safe_severities:
+        vote = Vote.SAFE
+        confidence = severity_confidence
+        threshold_used = 0.0
+        rationale = f"severity={severity_prediction} in safe_severities {thresholds.safe_severities}"
+    elif severity_lower in thresholds.threat_severities:
+        vote = Vote.THREAT
+        confidence = severity_confidence
+        threshold_used = 0.0
+        rationale = f"severity={severity_prediction} in threat_severities {thresholds.threat_severities}"
+    else:
+        # Unknown severity - vote based on whether it's closer to threat or safe
+        # This shouldn't happen with valid model outputs
+        vote = Vote.ABSTAIN
+        confidence = 0.5
+        threshold_used = 0.5
+        rationale = f"severity={severity_prediction} not in known categories"
+
+    return HeadVoteDetail(
+        head_name="severity",
+        vote=vote,
+        confidence=confidence,
+        weight=weight,
+        raw_probability=severity_confidence,
+        threshold_used=threshold_used,
+        prediction=severity_prediction,
+        rationale=rationale,
+    )
+
+
+def vote_technique(
+    technique_prediction: str | None,
+    technique_confidence: float,
+    thresholds: TechniqueHeadThresholds,
+    weight: float,
+) -> HeadVoteDetail:
+    """Cast vote for the primary technique classifier head.
+
+    Voting rules:
+    - THREAT if technique not in safe_techniques AND confidence >= threat_confidence
+    - SAFE if technique in safe_techniques OR confidence < safe_confidence
+    - ABSTAIN otherwise
+
+    Args:
+        technique_prediction: Predicted attack technique (e.g., "none", "instruction_override")
+        technique_confidence: Model confidence in the prediction (0.0 to 1.0)
+        thresholds: Threshold configuration for this head
+        weight: Vote weight for this head
+
+    Returns:
+        HeadVoteDetail with the vote, confidence, and rationale
+
+    Examples:
+        >>> detail = vote_technique("instruction_override", 0.70, TechniqueHeadThresholds(), 1.0)
+        >>> detail.vote
+        <Vote.THREAT: 'threat'>
+    """
+    # Handle None technique prediction
+    if technique_prediction is None:
+        technique_prediction = "none"
+
+    technique_lower = technique_prediction.lower()
+    is_safe_technique = technique_lower in [t.lower() for t in thresholds.safe_techniques]
+
+    if is_safe_technique:
+        vote = Vote.SAFE
+        confidence = technique_confidence
+        threshold_used = 0.0
+        rationale = f"technique={technique_prediction} in safe_techniques"
+    elif technique_confidence >= thresholds.threat_confidence:
+        vote = Vote.THREAT
+        confidence = technique_confidence
+        threshold_used = thresholds.threat_confidence
+        rationale = (
+            f"technique={technique_prediction} with confidence ({technique_confidence:.2%}) >= "
+            f"threat_confidence ({thresholds.threat_confidence:.2%})"
+        )
+    elif technique_confidence < thresholds.safe_confidence:
+        vote = Vote.SAFE
+        confidence = 1.0 - technique_confidence
+        threshold_used = thresholds.safe_confidence
+        rationale = (
+            f"technique={technique_prediction} with confidence ({technique_confidence:.2%}) < "
+            f"safe_confidence ({thresholds.safe_confidence:.2%})"
+        )
+    else:
+        vote = Vote.ABSTAIN
+        confidence = 0.5
+        threshold_used = (thresholds.threat_confidence + thresholds.safe_confidence) / 2
+        rationale = (
+            f"technique={technique_prediction} with confidence ({technique_confidence:.2%}) in gray zone "
+            f"[{thresholds.safe_confidence:.2%}, {thresholds.threat_confidence:.2%})"
+        )
+
+    return HeadVoteDetail(
+        head_name="technique",
+        vote=vote,
+        confidence=confidence,
+        weight=weight,
+        raw_probability=technique_confidence,
+        threshold_used=threshold_used,
+        prediction=technique_prediction,
+        rationale=rationale,
+    )
+
+
+def vote_harm(
+    max_probability: float,
+    active_labels: list[str],
+    thresholds: HarmHeadThresholds,
+    weight: float,
+) -> HeadVoteDetail:
+    """Cast vote for the harm types (multilabel) classifier head.
+
+    Voting rules:
+    - THREAT if max_probability >= threat_threshold
+    - SAFE if max_probability < safe_threshold
+    - ABSTAIN otherwise
+
+    Note: Harm head has the lowest weight (0.8) because it can trigger
+    false positives on benign content discussing sensitive topics.
+
+    Args:
+        max_probability: Maximum probability across all harm types (0.0 to 1.0)
+        active_labels: List of active harm type labels
+        thresholds: Threshold configuration for this head
+        weight: Vote weight for this head
+
+    Returns:
+        HeadVoteDetail with the vote, confidence, and rationale
+
+    Examples:
+        >>> detail = vote_harm(0.95, ["violence_or_physical_harm"], HarmHeadThresholds(), 0.8)
+        >>> detail.vote
+        <Vote.THREAT: 'threat'>
+    """
+    # Build prediction string from active labels
+    if active_labels:
+        prediction = ",".join(active_labels[:3])  # Limit to top 3 for readability
+        if len(active_labels) > 3:
+            prediction += f",+{len(active_labels) - 3}"
+    else:
+        prediction = "none"
+
+    if max_probability >= thresholds.threat_threshold:
+        vote = Vote.THREAT
+        confidence = max_probability
+        threshold_used = thresholds.threat_threshold
+        rationale = (
+            f"max_probability ({max_probability:.2%}) >= "
+            f"threat_threshold ({thresholds.threat_threshold:.2%})"
+        )
+    elif max_probability < thresholds.safe_threshold:
+        vote = Vote.SAFE
+        confidence = 1.0 - max_probability
+        threshold_used = thresholds.safe_threshold
+        rationale = (
+            f"max_probability ({max_probability:.2%}) < "
+            f"safe_threshold ({thresholds.safe_threshold:.2%})"
+        )
+    else:
+        vote = Vote.ABSTAIN
+        # Confidence for abstain reflects uncertainty
+        confidence = 0.5
+        threshold_used = (thresholds.threat_threshold + thresholds.safe_threshold) / 2
+        rationale = (
+            f"max_probability ({max_probability:.2%}) in gray zone "
+            f"[{thresholds.safe_threshold:.2%}, {thresholds.threat_threshold:.2%})"
+        )
+
+    return HeadVoteDetail(
+        head_name="harm",
+        vote=vote,
+        confidence=confidence,
+        weight=weight,
+        raw_probability=max_probability,
+        threshold_used=threshold_used,
+        prediction=prediction,
+        rationale=rationale,
+    )

raxe/domain/ml/voting/models.py

@@ -0,0 +1,222 @@
+"""Voting domain models - Pure data structures for ensemble voting.
+
+This module contains ONLY pure domain models - no I/O operations.
+All classes are immutable value objects representing voting states and results.
+
+Model Hierarchy:
+- Vote: Three-way vote enum (SAFE, ABSTAIN, THREAT)
+- HeadVoteDetail: Detailed vote from a single classifier head
+- VotingResult: Complete voting engine output with full transparency
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class Vote(str, Enum):
+    """Three-way vote classification.
+
+    Each classifier head casts one of these votes:
+    - SAFE: Head is confident the input is benign
+    - ABSTAIN: Head is uncertain (in the "gray zone")
+    - THREAT: Head is confident the input is a threat
+    """
+
+    SAFE = "safe"
+    ABSTAIN = "abstain"
+    THREAT = "threat"
+
+
+class Decision(str, Enum):
+    """Final three-way classification decision.
+
+    The ensemble produces one of these decisions:
+    - SAFE: Consensus that input is benign, allow with confidence
+    - REVIEW: Uncertain, recommend human review
+    - THREAT: Consensus that input is malicious, recommend block
+    """
+
+    SAFE = "safe"
+    REVIEW = "review"
+    THREAT = "threat"
+
+
+@dataclass(frozen=True)
+class HeadVoteDetail:
+    """Detailed vote from a single classifier head.
+
+    Provides full transparency into each head's voting decision,
+    including the raw probabilities and thresholds used.
+
+    Attributes:
+        head_name: Name of the classifier head (binary, family, severity, technique, harm)
+        vote: The vote cast by this head
+        confidence: Confidence in the vote (0.0 to 1.0)
+        weight: Weight applied to this head's vote
+        raw_probability: Raw model output probability
+        threshold_used: Threshold that triggered this vote
+        prediction: The head's prediction label (e.g., "jailbreak", "high")
+        rationale: Human-readable explanation for the vote
+    """
+
+    head_name: str
+    vote: Vote
+    confidence: float
+    weight: float
+    raw_probability: float
+    threshold_used: float
+    prediction: str
+    rationale: str
+
+    def __post_init__(self) -> None:
+        """Validate head vote detail."""
+        if not 0.0 <= self.confidence <= 1.0:
+            raise ValueError(f"confidence must be 0-1, got {self.confidence}")
+        if not 0.0 <= self.raw_probability <= 1.0:
+            raise ValueError(f"raw_probability must be 0-1, got {self.raw_probability}")
+        if self.weight < 0:
+            raise ValueError(f"weight must be non-negative, got {self.weight}")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to JSON-serializable dictionary."""
+        return {
+            "head_name": self.head_name,
+            "vote": self.vote.value,
+            "confidence": self.confidence,
+            "weight": self.weight,
+            "raw_probability": self.raw_probability,
+            "threshold_used": self.threshold_used,
+            "prediction": self.prediction,
+            "rationale": self.rationale,
+        }
+
+
+@dataclass(frozen=True)
+class VotingResult:
+    """Complete result from the ensemble voting engine.
+
+    Provides full transparency into the voting process, including
+    per-head votes, aggregated scores, and the decision rationale.
+
+    Attributes:
+        decision: Final classification (SAFE, REVIEW, THREAT)
+        confidence: Overall confidence in the decision (0.0 to 1.0)
+        preset_used: Name of the voting preset that was applied
+        per_head_votes: Detailed vote breakdown for each head
+        aggregated_scores: Weighted vote totals {"safe": X, "review": Y, "threat": Z}
+        decision_rule_triggered: Which decision rule made the final call
+        threat_vote_count: Number of heads that voted THREAT
+        safe_vote_count: Number of heads that voted SAFE
+        abstain_vote_count: Number of heads that abstained
+        weighted_threat_score: Total weighted THREAT votes
+        weighted_safe_score: Total weighted SAFE votes
+    """
+
+    decision: Decision
+    confidence: float
+    preset_used: str
+    per_head_votes: dict[str, HeadVoteDetail]
+    aggregated_scores: dict[str, float]
+    decision_rule_triggered: str
+    threat_vote_count: int
+    safe_vote_count: int
+    abstain_vote_count: int
+    weighted_threat_score: float = 0.0
+    weighted_safe_score: float = 0.0
+
+    def __post_init__(self) -> None:
+        """Validate voting result."""
+        if not 0.0 <= self.confidence <= 1.0:
+            raise ValueError(f"confidence must be 0-1, got {self.confidence}")
+        if self.threat_vote_count < 0:
+            raise ValueError(f"threat_vote_count must be non-negative")
+        if self.safe_vote_count < 0:
+            raise ValueError(f"safe_vote_count must be non-negative")
+        if self.abstain_vote_count < 0:
+            raise ValueError(f"abstain_vote_count must be non-negative")
+
+    @property
+    def is_threat(self) -> bool:
+        """True if the decision is THREAT."""
+        return self.decision == Decision.THREAT
+
+    @property
+    def is_safe(self) -> bool:
+        """True if the decision is SAFE."""
+        return self.decision == Decision.SAFE
+
+    @property
+    def is_review(self) -> bool:
+        """True if the decision is REVIEW (uncertain)."""
+        return self.decision == Decision.REVIEW
+
+    @property
+    def total_votes(self) -> int:
+        """Total number of votes cast (excluding abstentions for ratio calc)."""
+        return self.threat_vote_count + self.safe_vote_count
+
+    @property
+    def threat_ratio(self) -> float:
+        """Ratio of THREAT votes to total votes (excluding abstentions)."""
+        if self.total_votes == 0:
+            return 0.0
+        return self.threat_vote_count / self.total_votes
+
+    @property
+    def weighted_ratio(self) -> float:
+        """Ratio of weighted THREAT votes to weighted SAFE votes."""
+        if self.weighted_safe_score == 0:
+            return float("inf") if self.weighted_threat_score > 0 else 0.0
+        return self.weighted_threat_score / self.weighted_safe_score
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to JSON-serializable dictionary."""
+        # Handle infinity for JSON serialization (JSON doesn't support Infinity)
+        ratio = self.weighted_ratio
+        if ratio == float("inf"):
+            ratio = 999.0  # Sentinel value for "infinitely more threat than safe"
+
+        return {
+            "decision": self.decision.value,
+            "confidence": self.confidence,
+            "preset_used": self.preset_used,
+            "per_head_votes": {
+                name: vote.to_dict() for name, vote in self.per_head_votes.items()
+            },
+            "aggregated_scores": self.aggregated_scores,
+            "decision_rule_triggered": self.decision_rule_triggered,
+            "threat_vote_count": self.threat_vote_count,
+            "safe_vote_count": self.safe_vote_count,
+            "abstain_vote_count": self.abstain_vote_count,
+            "weighted_threat_score": self.weighted_threat_score,
+            "weighted_safe_score": self.weighted_safe_score,
+            "weighted_ratio": ratio,
+        }
+
+
+@dataclass(frozen=True)
+class HeadOutput:
+    """Raw output from a classifier head before voting.
+
+    This is an intermediate data structure used to pass head outputs
+    to the voting engine without coupling to GemmaClassificationResult.
+
+    Attributes:
+        head_name: Name of the head (binary, family, severity, technique, harm)
+        prediction: Predicted class label
+        confidence: Confidence in the prediction
+        probabilities: Full probability distribution (optional)
+        is_threat_indicator: Whether this prediction indicates a threat
+    """
+
+    head_name: str
+    prediction: str
+    confidence: float
+    probabilities: tuple[float, ...] | None = None
+    is_threat_indicator: bool = False
+
+    def __post_init__(self) -> None:
+        """Validate head output."""
+        if not 0.0 <= self.confidence <= 1.0:
+            raise ValueError(f"confidence must be 0-1, got {self.confidence}")

raxe/domain/models.py

@@ -0,0 +1,82 @@
+"""Core domain models for RAXE CE.
+
+All models are immutable value objects (frozen=True).
+Pure domain layer - no I/O operations.
+"""
+from dataclasses import dataclass
+from enum import Enum
+
+# Re-export Detection and ScanResult from executor for convenience
+from raxe.domain.engine.executor import Detection, ScanResult
+
+
+class ThreatType(Enum):
+    """Categories of threats.
+
+    Maps to rule families but provides more semantic categorization.
+    """
+    PROMPT_INJECTION = "PROMPT_INJECTION"
+    JAILBREAK = "JAILBREAK"
+    PII_LEAK = "PII_LEAK"
+    DATA_EXFILTRATION = "DATA_EXFIL"
+    SECURITY = "SECURITY"
+    QUALITY = "QUALITY"
+    CUSTOM = "CUSTOM"
+
+
+@dataclass(frozen=True)
+class ScanRequest:
+    """Request to scan text for threats.
+
+    Immutable value object for scan input.
+
+    Attributes:
+        text: The text to scan for threats
+        context: Optional context metadata (user_id, session_id, etc.)
+        rule_filters: Optional list of rule IDs to apply (None = all rules)
+        max_text_length: Maximum allowed text length (default 1MB)
+    """
+    text: str
+    context: dict[str, str] | None = None
+    rule_filters: list[str] | None = None
+    max_text_length: int = 1_000_000  # 1MB default limit
+
+    def __post_init__(self) -> None:
+        """Validate request after construction.
+
+        Raises:
+            ValueError: If validation fails
+        """
+        if not self.text:
+            raise ValueError("Text cannot be empty")
+
+        if len(self.text) > self.max_text_length:
+            raise ValueError(
+                f"Text exceeds maximum length of {self.max_text_length} "
+                f"(got {len(self.text)} chars)"
+            )
+
+        if self.rule_filters is not None and not self.rule_filters:
+            raise ValueError("rule_filters cannot be empty list (use None for all rules)")
+
+    @property
+    def text_length(self) -> int:
+        """Length of text to scan."""
+        return len(self.text)
+
+    @property
+    def has_filters(self) -> bool:
+        """True if rule filters are specified."""
+        return self.rule_filters is not None
+
+
+# Note: BlockAction and ScanPolicy have been removed.
+# The advanced policy system in domain/policies/ replaces this functionality.
+
+
+__all__ = [
+    "Detection",
+    "ScanRequest",
+    "ScanResult",
+    "ThreatType",
+]

raxe/domain/packs/__init__.py

@@ -0,0 +1,17 @@
+"""Domain models for rule packs.
+
+Pure domain layer - NO I/O operations.
+"""
+from raxe.domain.packs.models import (
+    PackManifest,
+    PackRule,
+    PackType,
+    RulePack,
+)
+
+__all__ = [
+    "PackManifest",
+    "PackRule",
+    "PackType",
+    "RulePack",
+]