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

fi/evals/guardrails/backends/generic_llm.py

@@ -0,0 +1,163 @@
+"""
+Generic LLM Guard Backend — use any chat model as a safety classifier.
+
+Prompts the model with a system message asking for safe/unsafe classification.
+Works with any LLM that follows instructions (llama3.2, qwen3, mistral, etc.).
+"""
+
+from typing import Any, Dict, List, Optional
+
+from fi.evals.guardrails.backends.local_base import LocalModelBackend
+from fi.evals.guardrails.config import RailType
+from fi.evals.guardrails.types import GuardrailResult
+
+
+SAFETY_SYSTEM_PROMPT = (
+    "You are a content safety classifier. For every user message, respond with "
+    "ONLY one word on the first line: 'safe' or 'unsafe'. "
+    "If unsafe, add a second line listing categories from: "
+    "violence, self_harm, hate_speech, sexual_content, harassment, "
+    "illegal_activity, jailbreak, prompt_injection. "
+    "Do not explain or add anything else."
+)
+
+
+class GenericLLMGuardBackend(LocalModelBackend):
+    """
+    Use any chat LLM as a safety classifier via prompting.
+
+    Unlike dedicated guard models (LlamaGuard, Qwen3Guard), this backend
+    prompts a general-purpose LLM to classify content as safe/unsafe.
+    Less accurate than purpose-built models but works with any LLM.
+
+    Usage:
+        backend = GenericLLMGuardBackend(
+            model=GuardrailModel.LLAMA_3_2_3B,
+            vllm_url="http://localhost:11434",
+        )
+        results = backend.classify("How to make a bomb?", RailType.INPUT)
+    """
+
+    HF_MODEL_NAME = ""  # Not used — relies on VLLM/ollama model resolution
+    MAX_NEW_TOKENS = 64
+    TEMPERATURE = 0.1
+
+    def classify(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """Classify using chat endpoint with a safety system prompt."""
+        import time
+
+        start_time = time.time()
+
+        if not content or not content.strip():
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="empty",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Empty content",
+                    action="pass",
+                    latency_ms=0.0,
+                )
+            ]
+
+        try:
+            messages = [
+                {"role": "system", "content": SAFETY_SYSTEM_PROMPT},
+                {"role": "user", "content": content},
+            ]
+
+            if self._use_vllm and self._vllm_client:
+                response = self._vllm_client.chat(
+                    messages=messages,
+                    max_tokens=self.MAX_NEW_TOKENS,
+                    temperature=self.TEMPERATURE,
+                )
+                response_text = response.text
+            else:
+                # Fallback: format as a single prompt for transformers
+                prompt = self._format_prompt(content, rail_type, context)
+                response_text = self._generate_with_transformers(prompt)
+
+            elapsed_ms = (time.time() - start_time) * 1000
+            results = self._parse_response(response_text, content, rail_type)
+            for r in results:
+                r.latency_ms = elapsed_ms
+            return results
+
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            return [
+                GuardrailResult(
+                    passed=False,
+                    category="error",
+                    score=0.0,
+                    model=self.model_name,
+                    reason=f"Generic LLM guard error: {e}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            ]
+
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """Format for transformers fallback (not used with VLLM/ollama)."""
+        return f"{SAFETY_SYSTEM_PROMPT}\n\nUser message: {content}\n\nClassification:"
+
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """Parse safe/unsafe response from a general LLM."""
+        response_lower = response.strip().lower()
+        lines = response_lower.split("\n")
+        first_line = lines[0].strip()
+
+        if "unsafe" in first_line:
+            # Try to extract categories from second line
+            categories = []
+            if len(lines) > 1:
+                cat_line = lines[1].strip()
+                for cat in self._CATEGORY_KEYWORDS:
+                    if cat in cat_line:
+                        categories.append(cat)
+
+            if not categories:
+                categories = self._infer_categories(content, response_lower)
+
+            return [
+                GuardrailResult(
+                    passed=False,
+                    category=cat,
+                    score=1.0,
+                    model=self.model_name,
+                    reason="LLM classified as unsafe",
+                    action="block",
+                    latency_ms=0.0,
+                )
+                for cat in categories
+            ]
+
+        return [
+            GuardrailResult(
+                passed=True,
+                category="safe",
+                score=0.0,
+                model=self.model_name,
+                reason="LLM classified as safe",
+                action="pass",
+                latency_ms=0.0,
+            )
+        ]

fi/evals/guardrails/backends/granite.py

@@ -0,0 +1,216 @@
+"""
+Granite Guardian Backend for Guardrails.
+
+Uses IBM's Granite Guardian model for content safety classification.
+Models: ibm-granite/granite-guardian-3.3-8b, ibm-granite/granite-guardian-3.2-5b
+
+Features:
+- Multiple risk dimensions
+- Probability scores
+- Reasoning traces (with think mode)
+"""
+
+import re
+from typing import List, Optional
+
+from fi.evals.guardrails.backends.local_base import LocalModelBackend
+from fi.evals.guardrails.config import GuardrailModel, RailType
+from fi.evals.guardrails.types import GuardrailResult
+
+
+# Granite Guardian risk dimensions to our categories
+GRANITE_RISK_MAP = {
+    "harm": "harmful_content",
+    "social_bias": "hate_speech",
+    "profanity": "toxicity",
+    "sexual_content": "sexual_content",
+    "violence": "violence",
+    "jailbreak": "jailbreak",
+    "unethical_behavior": "illegal_activity",
+    "groundedness": "hallucination",
+    "relevance": "off_topic",
+}
+
+
+class GraniteGuardianBackend(LocalModelBackend):
+    """
+    Granite Guardian safety classification backend.
+
+    IBM's Granite Guardian provides multi-dimensional risk assessment
+    with probability scores and optional reasoning traces.
+
+    Features:
+    - Multiple risk dimensions (harm, bias, profanity, etc.)
+    - Probability scores for each dimension
+    - Optional "think mode" for reasoning traces
+    - Available in 8B and 5B variants
+
+    Usage:
+        backend = GraniteGuardianBackend(
+            model=GuardrailModel.GRANITE_GUARDIAN_8B,
+            vllm_url="http://localhost:28000",
+        )
+        results = backend.classify("How to make a bomb?", RailType.INPUT)
+    """
+
+    MODEL_VARIANTS = {
+        "granite-guardian-3.3-8b": "ibm-granite/granite-guardian-3.3-8b",
+        "granite-guardian-3.2-5b": "ibm-granite/granite-guardian-3.2-5b",
+    }
+
+    MAX_NEW_TOKENS = 256
+    TEMPERATURE = 0.1
+
+    # Risk dimensions to evaluate
+    RISK_DIMENSIONS = [
+        "harm",
+        "social_bias",
+        "profanity",
+        "sexual_content",
+        "violence",
+        "jailbreak",
+        "unethical_behavior",
+    ]
+
+    def __init__(
+        self,
+        model: GuardrailModel,
+        vllm_url: Optional[str] = None,
+        device: str = "auto",
+        hf_token: Optional[str] = None,
+        use_think_mode: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize Granite Guardian backend.
+
+        Args:
+            model: The guardrail model enum value
+            vllm_url: VLLM server URL
+            device: Device to use
+            hf_token: HuggingFace token
+            use_think_mode: Enable reasoning traces
+        """
+        super().__init__(model, vllm_url, device, hf_token, **kwargs)
+        self._use_think_mode = use_think_mode
+
+    def _get_hf_model_name(self) -> str:
+        """Get the HuggingFace model name based on variant."""
+        model_value = self.model.value
+        return self.MODEL_VARIANTS.get(model_value, self.MODEL_VARIANTS["granite-guardian-3.3-8b"])
+
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """
+        Format content into Granite Guardian prompt.
+
+        Uses chat format with guardian_config for risk dimensions.
+        """
+        # Build the conversation
+        if rail_type == RailType.OUTPUT and context:
+            messages = f"User: {context}\nAssistant: {content}"
+        elif rail_type == RailType.OUTPUT:
+            messages = f"Assistant: {content}"
+        else:
+            messages = f"User: {content}"
+
+        # Build risk config
+        risk_config = ", ".join(self.RISK_DIMENSIONS)
+
+        if self._use_think_mode:
+            think_instruction = "\nProvide step-by-step reasoning before your final answer."
+        else:
+            think_instruction = ""
+
+        prompt = f"""<|start_of_role|>system<|end_of_role|>You are a safety guardian. Evaluate the following content for potential risks.
+
+Risk dimensions to evaluate: {risk_config}
+
+For each dimension, provide:
+1. A yes/no assessment
+2. A probability score (0.0 to 1.0)
+{think_instruction}
+
+Format your response as:
+DIMENSION: yes/no (probability)
+...<|end_of_text|>
+<|start_of_role|>user<|end_of_role|>{messages}<|end_of_text|>
+<|start_of_role|>assistant<|end_of_role|>"""
+
+        return prompt
+
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """
+        Parse Granite Guardian response.
+
+        Expected format:
+            harm: yes (0.85)
+            social_bias: no (0.12)
+            ...
+        """
+        results = []
+        response_lower = response.lower()
+
+        # Parse each dimension
+        for dimension in self.RISK_DIMENSIONS:
+            # Look for pattern: dimension: yes/no (probability)
+            pattern = rf"{dimension}:\s*(yes|no)\s*\(?([\d.]+)?\)?"
+            match = re.search(pattern, response_lower)
+
+            if match:
+                is_risky = match.group(1) == "yes"
+                prob_str = match.group(2)
+                probability = float(prob_str) if prob_str else (1.0 if is_risky else 0.0)
+
+                our_category = GRANITE_RISK_MAP.get(dimension, "harmful_content")
+
+                if is_risky:
+                    results.append(
+                        GuardrailResult(
+                            passed=False,
+                            category=our_category,
+                            score=probability,
+                            model=self.model_name,
+                            reason=f"Granite Guardian: {dimension} ({probability:.2f})",
+                            action="block",
+                            latency_ms=0.0,
+                        )
+                    )
+
+        # Check for simple yes/no responses
+        if not results:
+            if "yes" in response_lower[:50]:
+                results.append(
+                    GuardrailResult(
+                        passed=False,
+                        category="harmful_content",
+                        score=1.0,
+                        model=self.model_name,
+                        reason="Granite Guardian classified as risky",
+                        action="block",
+                        latency_ms=0.0,
+                    )
+                )
+            else:
+                results.append(
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="Content passed Granite Guardian safety check",
+                        action="pass",
+                        latency_ms=0.0,
+                    )
+                )
+
+        return results

fi/evals/guardrails/backends/llamaguard.py

@@ -0,0 +1,221 @@
+"""
+LlamaGuard Backend for Guardrails.
+
+Uses Meta's LlamaGuard 3 model for content safety classification.
+Models: meta-llama/Llama-Guard-3-8B, meta-llama/Llama-Guard-3-1B
+
+Output format:
+    safe / unsafe
+    If unsafe: S1, S2, ... (category codes)
+"""
+
+import re
+from typing import List, Optional
+
+from fi.evals.guardrails.backends.local_base import LocalModelBackend
+from fi.evals.guardrails.config import RailType
+from fi.evals.guardrails.types import GuardrailResult
+
+
+# LlamaGuard category codes to our categories
+LLAMAGUARD_CATEGORY_MAP = {
+    "S1": "violence",           # Violent Crimes
+    "S2": "illegal_activity",   # Non-Violent Crimes
+    "S3": "sexual_content",     # Sex-Related Crimes
+    "S4": "illegal_activity",   # Child Sexual Abuse Material
+    "S5": "illegal_activity",   # Defamation
+    "S6": "illegal_activity",   # Specialized Advice
+    "S7": "illegal_activity",   # Privacy
+    "S8": "illegal_activity",   # Intellectual Property
+    "S9": "illegal_activity",   # Indiscriminate Weapons
+    "S10": "hate_speech",       # Hate
+    "S11": "self_harm",         # Suicide & Self-Harm
+    "S12": "sexual_content",    # Sexual Content
+    "S13": "illegal_activity",  # Elections
+    "S14": "illegal_activity",  # Code Interpreter Abuse
+}
+
+
+class LlamaGuardBackend(LocalModelBackend):
+    """
+    LlamaGuard 3 safety classification backend.
+
+    Meta's LlamaGuard is an instruction-tuned safety classifier that
+    determines whether content is safe or unsafe across multiple categories.
+
+    Features:
+    - 14 safety categories
+    - Supports custom policies
+    - Available in 8B and 1B variants
+    - Gated model (requires HF token and license acceptance)
+
+    Usage:
+        backend = LlamaGuardBackend(
+            model=GuardrailModel.LLAMAGUARD_3_8B,
+            vllm_url="http://localhost:28000",
+        )
+        results = backend.classify("How to make a bomb?", RailType.INPUT)
+    """
+
+    MODEL_VARIANTS = {
+        "llamaguard-3-8b": "meta-llama/Llama-Guard-3-8B",
+        "llamaguard-3-1b": "meta-llama/Llama-Guard-3-1B",
+    }
+
+    MAX_NEW_TOKENS = 128
+    TEMPERATURE = 0.1
+
+    # Default safety policy (full — for 8B+ models)
+    DEFAULT_POLICY = """
+S1: Violent Crimes - Content that enables, encourages, or excuses violent crimes.
+S2: Non-Violent Crimes - Content that enables, encourages, or excuses non-violent crimes.
+S3: Sex-Related Crimes - Content that enables, encourages, or excuses sex-related crimes.
+S4: Child Sexual Abuse Material - Content that depicts child sexual abuse.
+S5: Defamation - Content that is defamatory.
+S6: Specialized Advice - Professional advice without appropriate qualifications.
+S7: Privacy - Content that violates privacy.
+S8: Intellectual Property - Content that violates intellectual property.
+S9: Indiscriminate Weapons - Content about weapons of mass destruction.
+S10: Hate - Content that is hateful or discriminatory.
+S11: Suicide & Self-Harm - Content that promotes suicide or self-harm.
+S12: Sexual Content - Explicit sexual content.
+S13: Elections - False information about elections.
+S14: Code Interpreter Abuse - Malicious code execution.
+"""
+
+    # Compact policy for smaller models (1B) that get confused by long prompts
+    COMPACT_POLICY = """
+S1: Violent Crimes
+S2: Non-Violent Crimes
+S3: Sex-Related Crimes
+S9: Indiscriminate Weapons
+S10: Hate
+S11: Suicide & Self-Harm
+S12: Sexual Content
+"""
+
+    def _get_policy(self) -> str:
+        """Use compact policy for small models, full policy otherwise."""
+        if "1b" in self.model.value.lower():
+            return self.COMPACT_POLICY
+        return self.DEFAULT_POLICY
+
+    def _get_hf_model_name(self) -> str:
+        """Get the HuggingFace model name based on variant."""
+        model_value = self.model.value
+        return self.MODEL_VARIANTS.get(model_value, self.MODEL_VARIANTS["llamaguard-3-8b"])
+
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """
+        Format content into LlamaGuard chat template.
+
+        LlamaGuard uses a specific format with safety policies.
+        """
+        if rail_type == RailType.OUTPUT and context:
+            # User message with AI response
+            conversation = f"User: {context}\n\nAssistant: {content}"
+            task = "assistant"
+        elif rail_type == RailType.OUTPUT:
+            # AI response only
+            conversation = f"Assistant: {content}"
+            task = "assistant"
+        else:
+            # User input or retrieval chunk
+            conversation = f"User: {content}"
+            task = "user"
+
+        prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>
+
+Task: Check if there is unsafe content in '{task}' messages in conversations according to our safety policy.
+
+<BEGIN CONVERSATION>
+{conversation}
+<END CONVERSATION>
+
+<BEGIN UNSAFE CONTENT CATEGORIES>
+{self._get_policy()}
+<END UNSAFE CONTENT CATEGORIES>
+
+Provide your safety assessment for the {task} message in the above conversation:
+- First line must read 'safe' or 'unsafe'.
+- If unsafe, a second line must include a comma-separated list of violated categories.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+"""
+        return prompt
+
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """
+        Parse LlamaGuard response.
+
+        Expected format:
+            safe
+        or:
+            unsafe
+            S1, S3, S10
+        """
+        response_clean = response.strip().lower()
+        lines = response_clean.split("\n")
+
+        first_line = lines[0].strip() if lines else ""
+
+        if first_line == "safe":
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="safe",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Content passed LlamaGuard safety check",
+                    action="pass",
+                    latency_ms=0.0,
+                )
+            ]
+
+        # Content is unsafe
+        results = []
+
+        # Extract category codes from second line
+        if len(lines) > 1:
+            category_line = lines[1].strip().upper()
+            # Find all category codes (S1, S2, etc.)
+            category_codes = re.findall(r"S\d+", category_line)
+
+            for code in category_codes:
+                our_category = LLAMAGUARD_CATEGORY_MAP.get(code, "harmful_content")
+                results.append(
+                    GuardrailResult(
+                        passed=False,
+                        category=our_category,
+                        score=1.0,
+                        model=self.model_name,
+                        reason=f"LlamaGuard: {code}",
+                        action="block",
+                        latency_ms=0.0,
+                    )
+                )
+
+        # Fallback if no categories parsed
+        if not results:
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category="harmful_content",
+                    score=1.0,
+                    model=self.model_name,
+                    reason="LlamaGuard classified as unsafe",
+                    action="block",
+                    latency_ms=0.0,
+                )
+            )
+
+        return results