guardix 0.1.0__py3-none-any.whl

guardix/__init__.py

@@ -0,0 +1,57 @@
+"""guardix — Universal LLM prompt guard against injection attacks."""
+
+from typing import Any, Optional
+
+from .core import Guardial, Policy, Decision
+from .exceptions import GuardBlocked, GuardError
+from .config import Config
+from .responses import is_blocked_response
+
+__version__ = "0.1.0"
+__all__ = [
+    "Guardial",
+    "Policy",
+    "Decision",
+    "GuardBlocked",
+    "GuardError",
+    "Config",
+    "guard_client",
+    "is_blocked_response",
+]
+
+
+def guard_client(client: Any, guardial: Optional[Guardial] = None, provider: Optional[str] = None) -> Any:
+    """Wrap any supported LLM client with prompt guarding in one line.
+
+    Auto-detects the client type:
+    - ``messages.create``            -> Anthropic
+    - ``models.generate_content``    -> Gemini (google-genai)
+    - ``chat.completions.create``    -> OpenAI and all OpenAI-compatible
+      providers (Azure OpenAI, Groq, OpenRouter, Together, ...)
+
+    ``provider`` overrides the name used in logs (e.g. "groq", "openrouter").
+
+    Usage:
+        from guardix import guard_client
+        client = guard_client(OpenAI())
+        client.chat.completions.create(...)  # guarded, never raises on block
+    """
+    from .providers import AnthropicAdapter, GeminiAdapter, OpenAIAdapter
+
+    messages = getattr(client, "messages", None)
+    if messages is not None and callable(getattr(messages, "create", None)):
+        return AnthropicAdapter(client, guardial=guardial)
+
+    models = getattr(client, "models", None)
+    if models is not None and callable(getattr(models, "generate_content", None)):
+        return GeminiAdapter(client, guardial=guardial)
+
+    chat = getattr(client, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    if completions is not None and callable(getattr(completions, "create", None)):
+        return OpenAIAdapter(client, guardial=guardial, provider_name=provider or "openai")
+
+    raise TypeError(
+        "Unsupported client: expected an object with messages.create (Anthropic), "
+        "models.generate_content (Gemini), or chat.completions.create (OpenAI-compatible)."
+    )

guardix/config.py

@@ -0,0 +1,54 @@
+"""Configuration for Guardial guard engine."""
+
+from typing import Any, Callable, Dict, List, Optional
+
+from .detectors.base import BaseDetector
+
+
+class Config:
+    """Guard configuration."""
+
+    DEFAULT_POLICIES: Dict[str, Dict[str, Any]] = {
+        "permissive": {"threshold": 0.9, "fail_mode": "open", "log_level": "INFO"},
+        "standard": {"threshold": 0.7, "fail_mode": "open", "log_level": "INFO"},
+        "strict": {"threshold": 0.5, "fail_mode": "closed", "log_level": "DEBUG"},
+    }
+
+    def __init__(
+        self,
+        policy: str = "standard",
+        threshold: Optional[float] = None,
+        fail_mode: Optional[str] = None,
+        log_level: Optional[str] = None,
+        log_sink: Optional[Callable[[Dict[str, Any]], None]] = None,
+        log_file: Optional[str] = "logs/guardix.jsonl",
+        custom_detectors: Optional[List[BaseDetector]] = None,
+        mask_raw_prompt: bool = True,
+        block_mode: str = "mock",
+        block_message: Optional[str] = None,
+    ) -> None:
+        defaults = self.DEFAULT_POLICIES.get(policy, self.DEFAULT_POLICIES["standard"])
+        self.policy = policy
+        self.threshold = threshold if threshold is not None else defaults["threshold"]
+        self.fail_mode = fail_mode if fail_mode is not None else defaults["fail_mode"]
+        self.log_level = log_level if log_level is not None else defaults["log_level"]
+        self.log_sink = log_sink
+        # When set, structured logs go to this file (folder auto-created)
+        # and console output is suppressed.
+        self.log_file = log_file
+        self.custom_detectors = custom_detectors or []
+        self.mask_raw_prompt = mask_raw_prompt
+        # "mock": blocked calls return a provider-shaped mock response so the
+        # pipeline never breaks. "raise": blocked calls raise GuardBlocked.
+        self.block_mode = block_mode
+        self.block_message = block_message
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "policy": self.policy,
+            "threshold": self.threshold,
+            "fail_mode": self.fail_mode,
+            "log_level": self.log_level,
+            "mask_raw_prompt": self.mask_raw_prompt,
+            "block_mode": self.block_mode,
+        }

guardix/core.py

@@ -0,0 +1,203 @@
+"""Core Guardial engine, policies, and decisions."""
+
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Callable
+
+from .config import Config
+from .logging_config import StructuredLogger
+from .exceptions import GuardBlocked, GuardError
+from .detectors.base import BaseDetector
+from .detectors.bert_detector import BertDetector
+
+
+@dataclass
+class Decision:
+    """Result of a guard scan."""
+
+    prompt_id: str
+    decision: str  # "ALLOW", "WARN", "BLOCK"
+    scores: Dict[str, float] = field(default_factory=dict)
+    threshold: float = 0.7
+    reason: str = ""
+    latency_ms: float = 0.0
+    provider: str = "unknown"
+    raw_prompt: Optional[str] = None
+    class_name: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "prompt_id": self.prompt_id,
+            "decision": self.decision,
+            "scores": self.scores,
+            "threshold": self.threshold,
+            "reason": self.reason,
+            "latency_ms": self.latency_ms,
+            "provider": self.provider,
+            "class_name": self.class_name,
+        }
+
+
+class Policy:
+    """Policy rules for guard decisions."""
+
+    def __init__(self, threshold: float = 0.7, warn_threshold: Optional[float] = None) -> None:
+        self.threshold = threshold
+        self.warn_threshold = warn_threshold if warn_threshold is not None else threshold * 0.85
+
+    def evaluate(self, max_score: float) -> str:
+        if max_score >= self.threshold:
+            return "BLOCK"
+        if max_score >= self.warn_threshold:
+            return "WARN"
+        return "ALLOW"
+
+
+class Guardial:
+    """Guard engine powered by fine-tuned BERT-mini."""
+
+    def __init__(
+        self,
+        policy: Optional[str] = None,
+        threshold: Optional[float] = None,
+        fail_mode: Optional[str] = None,
+        log_level: Optional[str] = None,
+        log_sink: Optional[Callable[[Dict[str, Any]], None]] = None,
+        log_file: Optional[str] = "logs/guardix.jsonl",
+        custom_detectors: Optional[List[BaseDetector]] = None,
+        mask_raw_prompt: bool = True,
+        block_mode: str = "mock",
+        block_message: Optional[str] = None,
+        config: Optional[Config] = None,
+    ) -> None:
+        if config is not None:
+            self.config = config
+        else:
+            self.config = Config(
+                policy=policy or "standard",
+                threshold=threshold,
+                fail_mode=fail_mode,
+                log_level=log_level,
+                log_sink=log_sink,
+                log_file=log_file,
+                custom_detectors=custom_detectors,
+                mask_raw_prompt=mask_raw_prompt,
+                block_mode=block_mode,
+                block_message=block_message,
+            )
+
+        self.policy = Policy(threshold=self.config.threshold)
+        self.logger = StructuredLogger(
+            level=self.config.log_level,
+            sink=self.config.log_sink,
+            log_file=getattr(self.config, "log_file", None),
+        )
+
+        self.detectors: List[BaseDetector] = [BertDetector()]
+        if self.config.custom_detectors:
+            self.detectors.extend(self.config.custom_detectors)
+
+    def analyze(self, prompt: str, provider: str = "unknown") -> Decision:
+        """Run BERT-mini detector and return a Decision. Never raises."""
+        prompt_id = str(uuid.uuid4())
+        start = time.perf_counter()
+        scores: Dict[str, float] = {}
+        class_name = ""
+
+        try:
+            for detector in self.detectors:
+                try:
+                    if hasattr(detector, "detect_and_classify"):
+                        score, class_name, _ = detector.detect_and_classify(prompt)
+                        scores[detector.name] = score
+                    else:
+                        score = detector.detect(prompt)
+                        scores[detector.name] = score
+                        if hasattr(detector, "classify"):
+                            try:
+                                cn, _ = detector.classify(prompt)
+                                class_name = cn
+                            except Exception:
+                                pass
+                except Exception as e:
+                    scores[detector.name] = 0.0
+                    self.logger.log_error(prompt_id, provider, f"{detector.name} failed: {e}", 0.0)
+                    if self.config.fail_mode == "closed":
+                        raise GuardError(f"{detector.name} failed: {e}", original_error=e)
+
+            max_score = max(scores.values()) if scores else 0.0
+            decision_label = self.policy.evaluate(max_score)
+
+            reasons = []
+            for name, score in scores.items():
+                if score >= self.policy.threshold:
+                    reasons.append(f"{name}={score:.2f}")
+            reason = f"Threshold exceeded by {', '.join(reasons)}" if reasons else "No detectors flagged"
+
+            latency_ms = (time.perf_counter() - start) * 1000
+
+            decision = Decision(
+                prompt_id=prompt_id,
+                decision=decision_label,
+                scores=scores,
+                threshold=self.config.threshold,
+                reason=reason,
+                latency_ms=latency_ms,
+                provider=provider,
+                raw_prompt=prompt if not self.config.mask_raw_prompt else None,
+                class_name=class_name,
+            )
+
+            self.logger.log_decision(
+                prompt_id=prompt_id,
+                provider=provider,
+                detector_results=scores,
+                decision=decision_label,
+                reason=reason,
+                latency_ms=latency_ms,
+                raw_prompt=prompt if not self.config.mask_raw_prompt else None,
+                level="DEBUG" if decision_label == "ALLOW" else "WARNING",
+            )
+
+            return decision
+
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start) * 1000
+            self.logger.log_error(prompt_id, provider, f"Guard engine failed: {e}", latency_ms)
+            if self.config.fail_mode == "closed":
+                raise GuardError(f"Guard engine failed: {e}", original_error=e)
+            return Decision(
+                prompt_id=prompt_id,
+                decision="ALLOW",
+                scores={},
+                threshold=self.config.threshold,
+                reason=f"Guard engine error (fail_open): {e}",
+                latency_ms=latency_ms,
+                provider=provider,
+            )
+
+    def guard(
+        self,
+        prompt: str,
+        provider: str = "unknown",
+        on_block: Optional[Callable[[Decision], Any]] = None,
+    ) -> Decision:
+        """Analyze the prompt and act on a BLOCK according to block_mode.
+
+        - ``on_block`` callback given: call it with the Decision.
+        - ``block_mode="raise"``: raise GuardBlocked (legacy behavior).
+        - ``block_mode="mock"`` (default): return the Decision unraised so
+          the caller (adapter/middleware/decorator) can substitute a
+          provider-shaped mock response and keep the pipeline alive.
+        """
+        decision = self.analyze(prompt, provider=provider)
+        if decision.decision == "BLOCK":
+            if on_block:
+                on_block(decision)
+            elif self.config.block_mode == "raise":
+                self.logger.log_block_action(
+                    decision.prompt_id, provider, "exception_raised", decision.reason
+                )
+                raise GuardBlocked(f"Prompt blocked: {decision.reason}", decision=decision)
+        return decision

guardix/decorators.py

@@ -0,0 +1,123 @@
+"""Decorators for easy guarding of LLM call functions."""
+
+from functools import wraps
+from typing import Any, Callable, Optional
+
+from .core import Guardial, Decision
+from .exceptions import GuardBlocked
+from .responses import (
+    anthropic_blocked_response,
+    gemini_blocked_response,
+    openai_blocked_response,
+)
+
+_RESPONSE_BUILDERS = {
+    "openai": openai_blocked_response,
+    "anthropic": anthropic_blocked_response,
+    "gemini": gemini_blocked_response,
+}
+
+
+def guardial_guard(
+    policy: Optional[str] = None,
+    threshold: Optional[float] = None,
+    fail_mode: Optional[str] = None,
+    on_block: Optional[Callable[[Decision], Any]] = None,
+    provider: str = "unknown",
+    block_mode: str = "mock",
+    response_format: str = "openai",
+    block_message: Optional[str] = None,
+) -> Callable[..., Any]:
+    """Decorator that guards a function taking a prompt or messages argument.
+
+    By default a blocked prompt does NOT raise: the wrapped function is
+    skipped and a provider-shaped mock response (``response_format``:
+    "openai", "anthropic", or "gemini") is returned, so the pipeline keeps
+    flowing. Pass ``block_mode="raise"`` for the old GuardBlocked behavior.
+
+    Usage:
+        @guardial_guard(policy="strict")
+        def chat(messages):
+            return openai_client.chat.completions.create(model="gpt-4", messages=messages)
+    """
+    g = Guardial(
+        policy=policy,
+        threshold=threshold,
+        fail_mode=fail_mode,
+        block_mode=block_mode,
+        block_message=block_message,
+    )
+    builder = _RESPONSE_BUILDERS.get(response_format, openai_blocked_response)
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            messages = kwargs.get("messages")
+            if messages is None and args:
+                first = args[0]
+                if isinstance(first, str):
+                    messages = [{"role": "user", "content": first}]
+                elif isinstance(first, list):
+                    messages = first
+            prompt = _messages_to_prompt(messages)
+            decision = g.guard(prompt, provider=provider, on_block=on_block)
+            if decision.decision == "BLOCK" and on_block is None and block_mode == "mock":
+                g.logger.log_block_action(
+                    decision.prompt_id, provider, "mock_response", decision.reason
+                )
+                if response_format == "gemini":
+                    return builder(decision, message=block_message)
+                return builder(decision, model=kwargs.get("model", "unknown"), message=block_message)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def guardial_audit(
+    policy: Optional[str] = None,
+    provider: str = "unknown",
+    log_sink: Optional[Callable[[Any], None]] = None,
+) -> Callable[..., Any]:
+    """Decorator that only audits (never blocks) LLM calls.
+
+    Usage:
+        @guardial_audit(policy="standard")
+        def chat(messages):
+            return openai_client.chat.completions.create(model="gpt-4", messages=messages)
+    """
+    g = Guardial(policy=policy, fail_mode="open", log_sink=log_sink)
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            messages = kwargs.get("messages")
+            if messages is None and args:
+                first = args[0]
+                if isinstance(first, str):
+                    messages = [{"role": "user", "content": first}]
+                elif isinstance(first, list):
+                    messages = first
+            prompt = _messages_to_prompt(messages)
+            g.analyze(prompt, provider=provider)  # audit only, never blocks
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def _messages_to_prompt(messages: Any) -> str:
+    if messages is None:
+        return ""
+    if isinstance(messages, str):
+        return messages
+    parts = []
+    for msg in messages:
+        if isinstance(msg, dict):
+            content = msg.get("content", "")
+            parts.append(content)
+        else:
+            parts.append(str(msg))
+    return "\n".join(parts)

guardix/detectors/__init__.py

@@ -0,0 +1,6 @@
+"""Detector exports — single BERT-mini model replaces all rule-based detectors."""
+
+from .base import BaseDetector
+from .bert_detector import BertDetector
+
+__all__ = ["BaseDetector", "BertDetector"]

guardix/detectors/base.py

@@ -0,0 +1,14 @@
+"""Base detector interface."""
+
+from abc import ABC, abstractmethod
+
+
+class BaseDetector(ABC):
+    """Abstract base class for all prompt injection detectors."""
+
+    name: str = "base"
+
+    @abstractmethod
+    def detect(self, prompt: str) -> float:
+        """Return a confidence score between 0.0 and 1.0."""
+        ...

guardix/detectors/bert_detector.py

@@ -0,0 +1,108 @@
+"""Single BERT-mini PyTorch detector — replaces all 16 old rule detectors."""
+
+import re
+import threading
+from typing import Dict, List, Tuple
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+from .base import BaseDetector
+
+
+MODEL_ID = "PraneshJs/promptgaurd"
+
+# Process-wide cache so every Guardial/BertDetector instance shares one
+# loaded model instead of re-downloading and re-loading per instance.
+_MODEL_CACHE: Dict[str, Tuple] = {}
+
+_SENTENCE_SPLIT = re.compile(r"(?<=[.!?\n])\s+")
+
+
+def _load_model(model_id: str) -> Tuple:
+    if model_id not in _MODEL_CACHE:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSequenceClassification.from_pretrained(model_id).to(device)
+        model.eval()
+        # Fast tokenizers are not thread-safe ("Already borrowed"); the
+        # shared cache means concurrent callers must serialize inference.
+        _MODEL_CACHE[model_id] = (tokenizer, model, device, threading.Lock())
+    return _MODEL_CACHE[model_id]
+
+
+class BertDetector(BaseDetector):
+    """Binary detector using fine-tuned BERT-mini (safe/attack).
+
+    Returns 0.0–1.0 attack probability. Also provides class-level prediction.
+    Long prompts cannot bypass detection via truncation: the prompt is scored
+    as overlapping 128-token sliding windows AND as individual sentences (so a
+    short injection buried in benign text gets an undiluted look), all in one
+    batched forward pass. The worst (most attack-like) segment wins.
+    """
+
+    name = "bert_mini"
+
+    def __init__(self, model_id: str = MODEL_ID):
+        self.tokenizer, self.model, self.device, self._lock = _load_model(model_id)
+        self.max_len = 128
+        self.stride = 64
+
+    def _segments(self, prompt: str) -> List[str]:
+        """Full prompt plus per-sentence segments for undiluted scoring.
+
+        Sentence segments are only added when the prompt exceeds the model's
+        max length: a short prompt cannot hide an injection via truncation,
+        and scoring its sentences in isolation causes false positives
+        (e.g. a bare "You are a helpful assistant." reads as role hijacking).
+        """
+        segments = [prompt]
+        n_tokens = len(self.tokenizer(prompt, truncation=False)["input_ids"])
+        if n_tokens > self.max_len:
+            sentences = [s.strip() for s in _SENTENCE_SPLIT.split(prompt) if len(s.strip()) >= 8]
+            if len(sentences) > 1:
+                segments.extend(sentences)
+        return segments
+
+    def _predict(self, prompt: str) -> Tuple[float, int, float]:
+        """Run one batched inference over all sliding windows and sentences.
+
+        Returns (attack_prob, pred_id, confidence) taken from the segment
+        with the highest attack probability.
+        """
+        with self._lock:
+            inputs = self.tokenizer(
+                self._segments(prompt),
+                truncation=True,
+                padding=True,
+                max_length=self.max_len,
+                stride=self.stride,
+                return_overflowing_tokens=True,
+                return_tensors="pt",
+            )
+            inputs.pop("overflow_to_sample_mapping", None)
+            inputs = inputs.to(self.device)
+            with torch.no_grad():
+                logits = self.model(**inputs).logits
+        probs = torch.softmax(logits, dim=-1)  # (num_segments, 2)
+        worst = int(torch.argmax(probs[:, 1]))
+        segment_probs = probs[worst]
+        pred_id = int(torch.argmax(segment_probs))
+        return float(segment_probs[1]), pred_id, float(segment_probs[pred_id])
+
+    def detect(self, prompt: str) -> float:
+        """Return 0.0-1.0 attack probability."""
+        return self.detect_and_classify(prompt)[0]
+
+    def classify(self, prompt: str) -> tuple:
+        """Return (class_name, confidence)."""
+        _, class_name, confidence = self.detect_and_classify(prompt)
+        return (class_name, confidence)
+
+    def detect_and_classify(self, prompt: str) -> Tuple[float, str, float]:
+        """Return (attack_prob, class_name, confidence) from a single inference."""
+        if not prompt or not prompt.strip():
+            return (0.0, "safe", 1.0)
+        attack_prob, pred_id, confidence = self._predict(prompt)
+        labels = self.model.config.id2label
+        return (attack_prob, labels[pred_id], confidence)

guardix/exceptions.py

@@ -0,0 +1,19 @@
+"""Exceptions raised by guardix."""
+
+from typing import Optional
+
+
+class GuardError(Exception):
+    """Raised when the guard itself fails in fail-closed mode."""
+
+    def __init__(self, message: str, original_error: Optional[Exception] = None) -> None:
+        super().__init__(message)
+        self.original_error = original_error
+
+
+class GuardBlocked(Exception):
+    """Raised when a prompt is blocked by the guard."""
+
+    def __init__(self, message: str, decision: "Decision") -> None:
+        super().__init__(message)
+        self.decision = decision