raxe 0.4.6__py3-none-any.whl

raxe/domain/inline_suppression.py

@@ -0,0 +1,176 @@
+"""Inline suppression types for SDK scan() method.
+
+This module provides PURE domain types for inline suppressions passed
+directly to scan() calls. These are separate from file-based suppressions
+(.raxe/suppressions.yaml) and take precedence over them.
+
+Example:
+    # Simple list of patterns
+    result = raxe.scan(text, suppress=["pi-001", "jb-*"])
+
+    # With action override
+    result = raxe.scan(text, suppress=[
+        "pi-001",  # Default SUPPRESS action
+        {"pattern": "jb-*", "action": "FLAG", "reason": "Under review"}
+    ])
+
+Clean Architecture:
+    - This is PURE domain logic
+    - No I/O, no file access, no logging
+    - Converts inline specs to Suppression objects
+"""
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, TypeAlias
+
+from raxe.domain.suppression import (
+    Suppression,
+    SuppressionAction,
+    SuppressionValidationError,
+)
+
+# Type alias for inline suppression specification
+# Can be either a string pattern or a dict with pattern, action, reason
+InlineSuppressionSpec: TypeAlias = str | dict[str, Any]
+
+
+@dataclass(frozen=True)
+class InlineSuppressionConfig:
+    """Configuration for an inline suppression.
+
+    Represents a suppression passed directly to scan() rather than
+    loaded from a configuration file.
+
+    Attributes:
+        pattern: Rule ID pattern (supports wildcards: pi-*, jb-*)
+        action: Action to take when matched (default: SUPPRESS)
+        reason: Optional reason for suppression
+    """
+
+    pattern: str
+    action: SuppressionAction = SuppressionAction.SUPPRESS
+    reason: str = "Inline suppression"
+
+
+def parse_inline_suppression(spec: InlineSuppressionSpec) -> Suppression:
+    """Parse an inline suppression specification into a Suppression object.
+
+    Accepts either:
+    - String pattern: "pi-001" or "jb-*"
+    - Dict with full config: {"pattern": "jb-*", "action": "FLAG", "reason": "..."}
+
+    Args:
+        spec: Inline suppression specification
+
+    Returns:
+        Suppression object ready for use
+
+    Raises:
+        SuppressionValidationError: If specification is invalid
+    """
+    if isinstance(spec, str):
+        # Simple string pattern - default to SUPPRESS action
+        return Suppression(
+            pattern=spec,
+            reason="Inline suppression",
+            action=SuppressionAction.SUPPRESS,
+            created_at=datetime.now(timezone.utc).isoformat(),
+            created_by="inline",
+        )
+
+    if isinstance(spec, dict):
+        # Dict with full config
+        pattern = spec.get("pattern")
+        if not pattern:
+            raise SuppressionValidationError(
+                "Inline suppression dict must have 'pattern' key"
+            )
+
+        # Parse action
+        action_str = spec.get("action", "SUPPRESS")
+        if isinstance(action_str, SuppressionAction):
+            action = action_str
+        else:
+            try:
+                action = SuppressionAction(str(action_str).upper())
+            except ValueError:
+                valid_actions = [a.value for a in SuppressionAction]
+                raise SuppressionValidationError(
+                    f"Invalid action '{action_str}'. "
+                    f"Valid actions: {', '.join(valid_actions)}"
+                ) from None
+
+        # Get reason (optional, has default)
+        reason = spec.get("reason", "Inline suppression")
+        if not reason:
+            reason = "Inline suppression"
+
+        return Suppression(
+            pattern=str(pattern),
+            reason=str(reason),
+            action=action,
+            created_at=datetime.now(timezone.utc).isoformat(),
+            created_by="inline",
+        )
+
+    raise SuppressionValidationError(
+        f"Invalid inline suppression type: {type(spec).__name__}. "
+        "Expected str or dict."
+    )
+
+
+def parse_inline_suppressions(
+    specs: list[InlineSuppressionSpec] | None,
+) -> list[Suppression]:
+    """Parse a list of inline suppression specifications.
+
+    Args:
+        specs: List of inline suppression specs, or None
+
+    Returns:
+        List of Suppression objects
+
+    Raises:
+        SuppressionValidationError: If any specification is invalid
+    """
+    if not specs:
+        return []
+
+    return [parse_inline_suppression(spec) for spec in specs]
+
+
+def merge_suppressions(
+    config_suppressions: list[Suppression],
+    inline_suppressions: list[Suppression],
+) -> list[Suppression]:
+    """Merge config file suppressions with inline suppressions.
+
+    Inline suppressions take precedence over config file suppressions
+    for the same pattern.
+
+    Args:
+        config_suppressions: Suppressions from config files
+        inline_suppressions: Suppressions passed to scan()
+
+    Returns:
+        Merged list with inline taking precedence
+    """
+    # Build dict with config suppressions first
+    merged: dict[str, Suppression] = {}
+    for supp in config_suppressions:
+        merged[supp.pattern] = supp
+
+    # Override with inline suppressions (they take precedence)
+    for supp in inline_suppressions:
+        merged[supp.pattern] = supp
+
+    return list(merged.values())
+
+
+__all__ = [
+    "InlineSuppressionConfig",
+    "InlineSuppressionSpec",
+    "merge_suppressions",
+    "parse_inline_suppression",
+    "parse_inline_suppressions",
+]

raxe/domain/ml/__init__.py

@@ -0,0 +1,133 @@
+"""L2 ML-based threat detection.
+
+This module defines the protocol and interfaces for L2 machine learning
+threat detection that augments L1 rule-based detection.
+
+L2 provides:
+- Semantic threat detection using Gemma embeddings
+- 5-head classification (is_threat, family, severity, technique, harm_types)
+- Multilabel harm type classification
+- Context-aware analysis with explainability
+- Probabilistic predictions with confidence scores
+
+L2 uses Gemma-based ONNX models containing:
+- EmbeddingGemma-300M ONNX model (256d embeddings)
+- 5 classifier heads (is_threat, threat_family, severity, primary_technique, harm_types)
+- Gemma tokenizer
+- Label configuration
+
+Note: ML features require optional dependencies: pip install raxe[ml]
+"""
+
+# Protocol and base types (no ML dependencies required)
+# Domain models (no ML dependencies required)
+from raxe.domain.ml.gemma_models import (
+    DEFAULT_HARM_THRESHOLDS,
+    GemmaClassificationResult,
+    HarmType,
+    MultilabelResult,
+    PrimaryTechnique,
+    Severity,
+    ThreatFamily,
+)
+from raxe.domain.ml.l2_config import (
+    L2Config,
+    L2EnsembleConfig,
+    L2ThresholdConfig,
+    create_example_config,
+    get_l2_config,
+    load_l2_config,
+    reset_l2_config,
+    set_l2_config,
+)
+from raxe.domain.ml.protocol import (
+    L2Detector,
+    L2Prediction,
+    L2Result,
+    L2ThreatType,
+)
+from raxe.domain.ml.scoring_models import (
+    ActionType,
+    ScoringMode,
+    ScoringResult,
+    ScoringThresholds,
+    ThreatLevel,
+    ThreatScore,
+)
+from raxe.domain.ml.stub_detector import StubL2Detector
+
+# ML-dependent imports (require numpy, onnxruntime, transformers)
+# These are lazily imported to allow basic SDK usage without ML deps
+_ML_AVAILABLE = False
+_ML_IMPORT_ERROR = None
+
+try:
+    from raxe.domain.ml.gemma_detector import (
+        GemmaL2Detector,
+        create_gemma_detector,
+    )
+    from raxe.domain.ml.threat_scorer import HierarchicalThreatScorer
+    _ML_AVAILABLE = True
+except ImportError as e:
+    _ML_IMPORT_ERROR = str(e)
+    # Provide stub implementations for when ML deps are missing
+    GemmaL2Detector = None  # type: ignore[assignment, misc]
+    HierarchicalThreatScorer = None  # type: ignore[assignment, misc]
+
+    def create_gemma_detector(*args, **kwargs):  # type: ignore[misc]
+        """Stub: ML dependencies not installed."""
+        raise ImportError(
+            f"ML features require optional dependencies. "
+            f"Install with: pip install raxe[ml]\n"
+            f"Original error: {_ML_IMPORT_ERROR}"
+        )
+
+
+def is_ml_available() -> bool:
+    """Check if ML dependencies are available."""
+    return _ML_AVAILABLE
+
+
+def get_ml_import_error() -> str | None:
+    """Get the ML import error message if ML is not available."""
+    return _ML_IMPORT_ERROR
+
+
+__all__ = [
+    # Availability check
+    "get_ml_import_error",
+    "is_ml_available",
+    # Domain Models (Gemma)
+    "DEFAULT_HARM_THRESHOLDS",
+    "GemmaClassificationResult",
+    "HarmType",
+    "MultilabelResult",
+    "PrimaryTechnique",
+    "Severity",
+    "ThreatFamily",
+    # Configuration
+    "L2Config",
+    "L2EnsembleConfig",
+    "L2ThresholdConfig",
+    "create_example_config",
+    "get_l2_config",
+    "load_l2_config",
+    "reset_l2_config",
+    "set_l2_config",
+    # Threat Scoring
+    "ActionType",
+    "ScoringMode",
+    "ScoringResult",
+    "ScoringThresholds",
+    "ThreatLevel",
+    "ThreatScore",
+    # L2 Detection
+    "GemmaL2Detector",
+    "HierarchicalThreatScorer",
+    "L2Detector",
+    "L2Prediction",
+    "L2Result",
+    "L2ThreatType",
+    "StubL2Detector",
+    "create_gemma_detector",
+]

raxe/domain/ml/embedding_cache.py

@@ -0,0 +1,309 @@
+"""
+LRU Embedding Cache for L2 ML Inference
+
+Thread-safe LRU cache for storing text embeddings to speed up repeated
+inference on the same or similar prompts.
+
+Performance characteristics:
+- Cache hit: ~0.1ms (vs ~2ms for embedding generation)
+- Memory: ~3KB per cached embedding (768 floats * 4 bytes)
+- Thread-safe: Uses threading.Lock for concurrent access
+
+This module is part of the domain layer and contains PURE logic:
+- No I/O operations (database, network, file system)
+- No logging (statistics are returned, not logged)
+- Deterministic behavior for testability
+
+Example:
+    cache = EmbeddingCache(max_size=1000)
+
+    # Check cache before generating embeddings
+    cached = cache.get("Hello world")
+    if cached is None:
+        embedding = generate_embedding("Hello world")
+        cache.put("Hello world", embedding)
+    else:
+        embedding = cached
+
+    # Check cache statistics
+    print(f"Hit rate: {cache.hit_rate:.1%}")
+"""
+
+from __future__ import annotations
+
+import hashlib
+import threading
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass(frozen=True)
+class CacheStats:
+    """
+    Immutable cache statistics snapshot.
+
+    Attributes:
+        hits: Number of cache hits
+        misses: Number of cache misses
+        size: Current number of entries in cache
+        max_size: Maximum cache capacity
+        hit_rate: Ratio of hits to total requests (0.0-1.0)
+        evictions: Number of entries evicted due to capacity
+    """
+
+    hits: int
+    misses: int
+    size: int
+    max_size: int
+    hit_rate: float
+    evictions: int
+
+
+@dataclass
+class EmbeddingCache:
+    """
+    Thread-safe LRU cache for text embeddings.
+
+    Uses an OrderedDict for O(1) LRU operations and a threading.Lock
+    for thread safety. Cache keys are derived from SHA256 hashes of
+    the input text (first 16 characters for compactness).
+
+    Performance targets:
+    - Cache hit: ~0.1ms
+    - Cache miss overhead: ~0.01ms
+    - Memory per entry: ~3KB (768 floats * 4 bytes)
+
+    Attributes:
+        max_size: Maximum number of entries (default 1000)
+                  Set to 0 to disable caching
+
+    Example:
+        # Create cache with default size
+        cache = EmbeddingCache()
+
+        # Create cache with custom size
+        cache = EmbeddingCache(max_size=5000)
+
+        # Disable caching
+        cache = EmbeddingCache(max_size=0)
+
+        # Use cache
+        embedding = cache.get("test prompt")
+        if embedding is None:
+            embedding = model.encode("test prompt")
+            cache.put("test prompt", embedding)
+
+        # Check statistics
+        stats = cache.stats
+        print(f"Hit rate: {stats.hit_rate:.1%}")
+    """
+
+    max_size: int = 1000
+
+    # Private fields initialized in __post_init__
+    _cache: OrderedDict[str, "np.ndarray"] = field(
+        default_factory=OrderedDict, repr=False
+    )
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+    _hits: int = field(default=0, repr=False)
+    _misses: int = field(default=0, repr=False)
+    _evictions: int = field(default=0, repr=False)
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        if self.max_size < 0:
+            raise ValueError("max_size must be non-negative")
+
+    @staticmethod
+    def _compute_key(text: str) -> str:
+        """
+        Compute cache key from text using SHA256 hash.
+
+        Uses first 16 characters of hex digest for compactness while
+        maintaining extremely low collision probability.
+
+        Args:
+            text: Input text to hash
+
+        Returns:
+            16-character hex string cache key
+        """
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
+
+    @property
+    def enabled(self) -> bool:
+        """Return True if caching is enabled (max_size > 0)."""
+        return self.max_size > 0
+
+    @property
+    def size(self) -> int:
+        """Return current number of cached entries."""
+        with self._lock:
+            return len(self._cache)
+
+    @property
+    def hits(self) -> int:
+        """Return number of cache hits."""
+        with self._lock:
+            return self._hits
+
+    @property
+    def misses(self) -> int:
+        """Return number of cache misses."""
+        with self._lock:
+            return self._misses
+
+    @property
+    def evictions(self) -> int:
+        """Return number of cache evictions."""
+        with self._lock:
+            return self._evictions
+
+    @property
+    def hit_rate(self) -> float:
+        """
+        Return cache hit rate as ratio (0.0-1.0).
+
+        Returns:
+            Hit rate as float between 0.0 and 1.0.
+            Returns 0.0 if no requests have been made.
+        """
+        with self._lock:
+            total = self._hits + self._misses
+            if total == 0:
+                return 0.0
+            return self._hits / total
+
+    @property
+    def stats(self) -> CacheStats:
+        """
+        Return snapshot of current cache statistics.
+
+        Thread-safe: Takes a consistent snapshot under lock.
+
+        Returns:
+            CacheStats immutable dataclass with current statistics
+        """
+        with self._lock:
+            total = self._hits + self._misses
+            hit_rate = self._hits / total if total > 0 else 0.0
+            return CacheStats(
+                hits=self._hits,
+                misses=self._misses,
+                size=len(self._cache),
+                max_size=self.max_size,
+                hit_rate=hit_rate,
+                evictions=self._evictions,
+            )
+
+    def get(self, text: str) -> "np.ndarray | None":
+        """
+        Retrieve cached embedding for text.
+
+        If found, moves entry to end of OrderedDict (most recently used).
+        Thread-safe operation.
+
+        Args:
+            text: Input text to look up
+
+        Returns:
+            Cached numpy array embedding, or None if not found.
+            Returns None immediately if caching is disabled.
+        """
+        if not self.enabled:
+            return None
+
+        key = self._compute_key(text)
+
+        with self._lock:
+            if key in self._cache:
+                # Move to end (most recently used)
+                self._cache.move_to_end(key)
+                self._hits += 1
+                return self._cache[key]
+            else:
+                self._misses += 1
+                return None
+
+    def put(self, text: str, embedding: "np.ndarray") -> None:
+        """
+        Store embedding in cache.
+
+        If cache is at capacity, evicts least recently used entry.
+        Thread-safe operation.
+
+        Args:
+            text: Input text (used to compute cache key)
+            embedding: Numpy array embedding to cache
+
+        Note:
+            Does nothing if caching is disabled (max_size=0).
+            If the text is already cached, updates the embedding
+            and moves entry to most recently used position.
+        """
+        if not self.enabled:
+            return
+
+        key = self._compute_key(text)
+
+        with self._lock:
+            # If key exists, update and move to end
+            if key in self._cache:
+                self._cache.move_to_end(key)
+                self._cache[key] = embedding
+                return
+
+            # Check if we need to evict
+            if len(self._cache) >= self.max_size:
+                # Remove oldest (first) item
+                self._cache.popitem(last=False)
+                self._evictions += 1
+
+            # Add new entry at end (most recently used)
+            self._cache[key] = embedding
+
+    def clear(self) -> None:
+        """
+        Clear all cached entries.
+
+        Thread-safe operation. Resets hits/misses/evictions counters.
+        """
+        with self._lock:
+            self._cache.clear()
+            self._hits = 0
+            self._misses = 0
+            self._evictions = 0
+
+    def contains(self, text: str) -> bool:
+        """
+        Check if text is in cache without affecting statistics.
+
+        Thread-safe operation. Does not update hit/miss counters
+        or change LRU order.
+
+        Args:
+            text: Input text to check
+
+        Returns:
+            True if text is cached, False otherwise.
+            Returns False immediately if caching is disabled.
+        """
+        if not self.enabled:
+            return False
+
+        key = self._compute_key(text)
+
+        with self._lock:
+            return key in self._cache
+
+    def __len__(self) -> int:
+        """Return current number of cached entries."""
+        return self.size
+
+    def __contains__(self, text: str) -> bool:
+        """Support 'in' operator for cache membership check."""
+        return self.contains(text)