RohanExploit/Meta-hackathon

1"""NeMo Guardrails equivalent — Llama Guard 3 input/output middleware.2 3Uses Groq's hosted `llama-guard-3-8b` to classify text as safe/unsafe.4Both the user's inbound query and the model's outbound response are screened.5If either is classified `unsafe`, the pipeline halts immediately and returns6a predefined security-violation response.7 8This module exposes two async functions:9    - guard_input(text)  → (is_safe: bool, category: str | None)10    - guard_output(text) → (is_safe: bool, category: str | None)11 12And a convenience exception:13    - SafetyViolation14"""15 16from __future__ import annotations17 18import asyncio19import logging20from dataclasses import dataclass21from typing import Tuple22 23from groq import AsyncGroq24 25from .config import get_config26 27logger = logging.getLogger(__name__)28 29# ── Predefined violation responses ──────────────────────────────────30INPUT_VIOLATION_RESPONSE = (31    "⚠️ Your query has been flagged as potentially unsafe and cannot be processed. "32    "Please rephrase your request or contact support if you believe this is an error."33)34 35OUTPUT_VIOLATION_RESPONSE = (36    "⚠️ The generated response was flagged by our safety system and has been withheld. "37    "This is a precautionary measure to prevent potentially harmful content."38)39 40 41class SafetyViolation(Exception):42    """Raised when Llama Guard flags content as unsafe."""43 44    def __init__(self, direction: str, category: str | None = None):45        self.direction = direction  # "input" or "output"46        self.category = category47        msg = f"Safety violation on {direction}"48        if category:49            msg += f" (category: {category})"50        super().__init__(msg)51 52 53@dataclass54class GuardResult:55    """Structured result from a Llama Guard check."""56 57    is_safe: bool58    raw_response: str59    category: str | None = None60 61 62# ── Internal helper ──────────────────────────────────────────────────63 64_client: AsyncGroq | None = None65 66 67def _get_client() -> AsyncGroq:68    global _client69    if _client is None:70        cfg = get_config()71        _client = AsyncGroq(api_key=cfg.groq.api_key)72    return _client73 74 75async def _classify(76    text: str,77    role: str = "User",78) -> GuardResult:79    """Send text through Llama Guard 3 for safety classification.80 81    Llama Guard 3 returns either:82        "safe"83        "unsafe\\nS<category_number>"84 85    Args:86        text: The text to classify.87        role: "User" for input guarding, "Agent" for output guarding.88 89    Returns:90        GuardResult with is_safe flag and optional violation category.91    """92    cfg = get_config()93    client = _get_client()94 95    # Llama Guard 3 expects a specific prompt format96    guard_prompt = (97        f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"98        f"Task: Check if there is unsafe content in '{role}' messages "99        f"in conversations according to our safety policy.\n\n"100        f"<BEGIN CONVERSATION>\n\n"101        f"{role}: {text}\n\n"102        f"<END CONVERSATION>\n\n"103        f"Provide your safety assessment for {role} in the above conversation:\n"104        f"- First line must read 'safe' or 'unsafe'.\n"105        f"- If unsafe, a second line must include a comma-separated list of "106        f"violated categories.<|eot_id|>"107        f"<|start_header_id|>assistant<|end_header_id|>\n\n"108    )109 110    try:111        response = await client.chat.completions.create(112            model=cfg.groq.guard_model,113            messages=[{"role": "user", "content": guard_prompt}],114            temperature=0.0,115            max_tokens=50,116        )117        raw = (response.choices[0].message.content or "").strip().lower()118    except Exception as exc:119        # If the guard itself fails, we log and allow through (fail-open)120        # In production you may want fail-closed instead.121        logger.warning("Llama Guard call failed: %s — defaulting to safe", exc)122        return GuardResult(is_safe=True, raw_response=f"ERROR: {exc}")123 124    if raw.startswith("unsafe"):125        lines = raw.split("\n", 1)126        category = lines[1].strip() if len(lines) > 1 else None127        return GuardResult(is_safe=False, raw_response=raw, category=category)128 129    return GuardResult(is_safe=True, raw_response=raw)130 131 132# ── Public API ───────────────────────────────────────────────────────133 134async def guard_input(text: str) -> GuardResult:135    """Screen a user query before it enters the pipeline.136 137    Raises SafetyViolation if flagged unsafe.138    """139    result = await _classify(text, role="User")140    if not result.is_safe:141        logger.warning("INPUT blocked: %s", result.raw_response)142        raise SafetyViolation("input", result.category)143    return result144 145 146async def guard_output(text: str) -> GuardResult:147    """Screen the final generated response before returning to the user.148 149    Raises SafetyViolation if flagged unsafe.150    """151    result = await _classify(text, role="Agent")152    if not result.is_safe:153        logger.warning("OUTPUT blocked: %s", result.raw_response)154        raise SafetyViolation("output", result.category)155    return result156 157 158async def guard_both(query: str, response: str) -> Tuple[GuardResult, GuardResult]:159    """Convenience: run input and output guards concurrently.160 161    This is useful when you want to re-validate the query alongside162    the response in a single await.163    """164    return await asyncio.gather(165        guard_input(query),166        guard_output(response),167    )168