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

fi/evals/guardrails/backends/openai.py

@@ -0,0 +1,365 @@
+"""
+OpenAI Moderation Backend for Guardrails.
+
+Uses OpenAI's Moderation API (FREE) for content screening.
+Supports the omni-moderation-latest model with 13 categories.
+
+Docs: https://platform.openai.com/docs/guides/moderation
+"""
+
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from fi.evals.guardrails.backends.base import BaseBackend
+from fi.evals.guardrails.config import GuardrailModel, RailType
+from fi.evals.guardrails.types import GuardrailResult
+
+
+# Category mapping from OpenAI to our standard categories
+OPENAI_CATEGORY_MAP = {
+    "hate": "hate_speech",
+    "hate/threatening": "hate_speech",
+    "harassment": "harassment",
+    "harassment/threatening": "harassment",
+    "self-harm": "self_harm",
+    "self-harm/intent": "self_harm",
+    "self-harm/instructions": "self_harm",
+    "sexual": "sexual_content",
+    "sexual/minors": "sexual_content",
+    "violence": "violence",
+    "violence/graphic": "violence",
+    "illicit": "illegal_activity",
+    "illicit/violent": "illegal_activity",
+}
+
+
+class OpenAIBackend(BaseBackend):
+    """
+    OpenAI Moderation API backend.
+
+    This is a FREE API provided by OpenAI for content moderation.
+    It supports 13 categories with confidence scores from 0 to 1.
+
+    Usage:
+        backend = OpenAIBackend(
+            model=GuardrailModel.OPENAI_MODERATION,
+            api_key="sk-..."  # or set OPENAI_API_KEY env var
+        )
+        results = backend.classify("some content", RailType.INPUT)
+
+    Rate Limits:
+        - 1,000 requests per minute
+        - 150,000 tokens per minute
+        - Does not count toward API usage limits
+    """
+
+    def __init__(
+        self,
+        model: GuardrailModel,
+        api_key: Optional[str] = None,
+        moderation_model: str = "omni-moderation-latest",
+    ):
+        """
+        Initialize OpenAI Moderation backend.
+
+        Args:
+            model: The guardrail model enum value
+            api_key: OpenAI API key (falls back to OPENAI_API_KEY env var)
+            moderation_model: OpenAI moderation model to use
+        """
+        super().__init__(model)
+
+        self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "OpenAI API key required. Set OPENAI_API_KEY environment variable "
+                "or pass api_key parameter."
+            )
+
+        self._moderation_model = moderation_model
+        self._client = None
+
+    def _get_client(self):
+        """Lazy-load the OpenAI client."""
+        if self._client is None:
+            try:
+                from openai import OpenAI
+            except ImportError:
+                raise ImportError(
+                    "openai package required for OpenAI backend. "
+                    "Install with: pip install openai"
+                )
+            self._client = OpenAI(api_key=self._api_key)
+        return self._client
+
+    def classify(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """
+        Classify content using OpenAI Moderation API.
+
+        Args:
+            content: Content to classify
+            rail_type: Type of rail (input, output, retrieval)
+            context: Optional context (not used by OpenAI)
+            metadata: Optional metadata (not used by OpenAI)
+
+        Returns:
+            List of GuardrailResult objects for flagged categories
+        """
+        start_time = time.time()
+        results = []
+
+        # Handle empty content
+        if not content or not content.strip():
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="empty",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Empty or whitespace-only content",
+                    action="pass",
+                    latency_ms=(time.time() - start_time) * 1000,
+                )
+            ]
+
+        try:
+            client = self._get_client()
+
+            # Call OpenAI Moderation API
+            response = client.moderations.create(
+                input=content,
+                model=self._moderation_model,
+            )
+
+            elapsed_ms = (time.time() - start_time) * 1000
+
+            if not response.results:
+                return [
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="No moderation results returned",
+                        action="pass",
+                        latency_ms=elapsed_ms,
+                    )
+                ]
+
+            # Process the first result (single input)
+            mod_result = response.results[0]
+
+            # Check if any category was flagged
+            if not mod_result.flagged:
+                return [
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="Content passed all moderation checks",
+                        action="pass",
+                        latency_ms=elapsed_ms,
+                    )
+                ]
+
+            # Extract flagged categories and scores
+            categories = mod_result.categories
+            scores = mod_result.category_scores
+
+            # Convert to dict for easier access
+            categories_dict = categories.model_dump() if hasattr(categories, 'model_dump') else dict(categories)
+            scores_dict = scores.model_dump() if hasattr(scores, 'model_dump') else dict(scores)
+
+            # Create results for each flagged category
+            for openai_category, is_flagged in categories_dict.items():
+                if is_flagged:
+                    # Map to our category name
+                    our_category = OPENAI_CATEGORY_MAP.get(
+                        openai_category, openai_category.replace("/", "_").replace("-", "_")
+                    )
+
+                    # Get the score (handle nested category names)
+                    score_key = openai_category.replace("/", "_").replace("-", "_")
+                    score = scores_dict.get(score_key, scores_dict.get(openai_category, 0.5))
+
+                    results.append(
+                        GuardrailResult(
+                            passed=False,
+                            category=our_category,
+                            score=float(score),
+                            model=self.model_name,
+                            reason=f"Content flagged for {openai_category}",
+                            action="block",
+                            latency_ms=elapsed_ms,
+                        )
+                    )
+
+            # If flagged but no specific categories (shouldn't happen), add generic
+            if not results:
+                results.append(
+                    GuardrailResult(
+                        passed=False,
+                        category="harmful_content",
+                        score=1.0,
+                        model=self.model_name,
+                        reason="Content flagged by OpenAI moderation",
+                        action="block",
+                        latency_ms=elapsed_ms,
+                    )
+                )
+
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category="error",
+                    score=0.0,
+                    model=self.model_name,
+                    reason=f"OpenAI API error: {str(e)}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            )
+
+        return results
+
+    async def classify_async(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """
+        Async version using OpenAI's async client.
+
+        Args:
+            content: Content to classify
+            rail_type: Type of rail
+            context: Optional context
+            metadata: Optional metadata
+
+        Returns:
+            List of GuardrailResult objects
+        """
+        start_time = time.time()
+
+        # Handle empty content
+        if not content or not content.strip():
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="empty",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Empty or whitespace-only content",
+                    action="pass",
+                    latency_ms=(time.time() - start_time) * 1000,
+                )
+            ]
+
+        try:
+            from openai import AsyncOpenAI
+        except ImportError:
+            # Fall back to sync implementation
+            return await super().classify_async(content, rail_type, context, metadata)
+
+        try:
+            async_client = AsyncOpenAI(api_key=self._api_key)
+
+            response = await async_client.moderations.create(
+                input=content,
+                model=self._moderation_model,
+            )
+
+            elapsed_ms = (time.time() - start_time) * 1000
+
+            if not response.results:
+                return [
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="No moderation results returned",
+                        action="pass",
+                        latency_ms=elapsed_ms,
+                    )
+                ]
+
+            mod_result = response.results[0]
+
+            if not mod_result.flagged:
+                return [
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="Content passed all moderation checks",
+                        action="pass",
+                        latency_ms=elapsed_ms,
+                    )
+                ]
+
+            results = []
+            categories_dict = mod_result.categories.model_dump() if hasattr(mod_result.categories, 'model_dump') else dict(mod_result.categories)
+            scores_dict = mod_result.category_scores.model_dump() if hasattr(mod_result.category_scores, 'model_dump') else dict(mod_result.category_scores)
+
+            for openai_category, is_flagged in categories_dict.items():
+                if is_flagged:
+                    our_category = OPENAI_CATEGORY_MAP.get(
+                        openai_category, openai_category.replace("/", "_").replace("-", "_")
+                    )
+                    score_key = openai_category.replace("/", "_").replace("-", "_")
+                    score = scores_dict.get(score_key, scores_dict.get(openai_category, 0.5))
+
+                    results.append(
+                        GuardrailResult(
+                            passed=False,
+                            category=our_category,
+                            score=float(score),
+                            model=self.model_name,
+                            reason=f"Content flagged for {openai_category}",
+                            action="block",
+                            latency_ms=elapsed_ms,
+                        )
+                    )
+
+            if not results:
+                results.append(
+                    GuardrailResult(
+                        passed=False,
+                        category="harmful_content",
+                        score=1.0,
+                        model=self.model_name,
+                        reason="Content flagged by OpenAI moderation",
+                        action="block",
+                        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"OpenAI API error: {str(e)}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            ]

fi/evals/guardrails/backends/qwen.py

@@ -0,0 +1,170 @@
+"""
+Qwen3Guard Backend for Guardrails.
+
+Uses Alibaba's Qwen3Guard model for content safety classification.
+Models: Qwen/Qwen3Guard-8B, Qwen/Qwen3Guard-4B
+
+Features:
+- 119 language support
+- Safe/Controversial/Unsafe classification
+- Streaming support (Qwen3Guard-Stream)
+"""
+
+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
+
+
+class Qwen3GuardBackend(LocalModelBackend):
+    """
+    Qwen3Guard safety classification backend.
+
+    Alibaba's Qwen3Guard is a multilingual safety classifier
+    supporting 119 languages with three-level classification.
+
+    Features:
+    - 119 language support
+    - Safe/Controversial/Unsafe classification
+    - Severity levels
+    - Available in 8B, 4B, and streaming variants
+
+    Usage:
+        backend = Qwen3GuardBackend(
+            model=GuardrailModel.QWEN3GUARD_8B,
+            vllm_url="http://localhost:28000",
+        )
+        results = backend.classify("How to make a bomb?", RailType.INPUT)
+    """
+
+    MODEL_VARIANTS = {
+        "qwen3guard-8b": "Qwen/Qwen3Guard-8B",
+        "qwen3guard-4b": "Qwen/Qwen3Guard-4B",
+        "qwen3guard-0.6b": "Qwen/Qwen3Guard-0.6B",
+    }
+
+    MAX_NEW_TOKENS = 64
+    TEMPERATURE = 0.1
+
+    # Qwen3Guard classification to our categories
+    SEVERITY_THRESHOLDS = {
+        "safe": 0.0,
+        "controversial": 0.5,
+        "unsafe": 1.0,
+    }
+
+    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["qwen3guard-8b"])
+
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """
+        Format content into Qwen3Guard prompt.
+
+        Qwen3Guard uses a simple classification prompt.
+        """
+        if rail_type == RailType.OUTPUT and context:
+            text_to_classify = f"User: {context}\nAssistant: {content}"
+        elif rail_type == RailType.OUTPUT:
+            text_to_classify = f"Assistant: {content}"
+        else:
+            text_to_classify = content
+
+        prompt = f"""<|im_start|>system
+You are a content safety classifier. Classify the following content as:
+- Safe: Content is appropriate and harmless
+- Controversial: Content may be sensitive but not clearly harmful
+- Unsafe: Content is harmful, dangerous, or violates policies
+
+Respond with only the classification and a brief reason.<|im_end|>
+<|im_start|>user
+Classify this content:
+
+{text_to_classify}<|im_end|>
+<|im_start|>assistant
+"""
+        return prompt
+
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """
+        Parse Qwen3Guard response.
+
+        Expected format:
+            Unsafe: This content contains instructions for creating weapons.
+        or:
+            Safe: This is a normal conversation.
+        """
+        response_lower = response.lower().strip()
+
+        # Determine classification
+        if response_lower.startswith("unsafe"):
+            classification = "unsafe"
+            score = 1.0
+            passed = False
+        elif response_lower.startswith("controversial"):
+            classification = "controversial"
+            score = 0.5
+            passed = False  # Treat controversial as failed for safety
+        else:
+            classification = "safe"
+            score = 0.0
+            passed = True
+
+        if passed:
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="safe",
+                    score=score,
+                    model=self.model_name,
+                    reason="Content passed Qwen3Guard safety check",
+                    action="pass",
+                    latency_ms=0.0,
+                )
+            ]
+
+        # Infer categories from content and response
+        categories = self._infer_categories(content, response)
+
+        results = []
+        for category in categories:
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category=category,
+                    score=score,
+                    model=self.model_name,
+                    reason=f"Qwen3Guard: {classification}",
+                    action="block" if classification == "unsafe" else "flag",
+                    latency_ms=0.0,
+                )
+            )
+
+        if not results:
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category="harmful_content",
+                    score=score,
+                    model=self.model_name,
+                    reason=f"Qwen3Guard: {classification}",
+                    action="block" if classification == "unsafe" else "flag",
+                    latency_ms=0.0,
+                )
+            )
+
+        return results
+
+    # _infer_categories inherited from LocalModelBackend

fi/evals/guardrails/backends/shieldgemma.py

@@ -0,0 +1,154 @@
+"""
+ShieldGemma Backend for Guardrails.
+
+Uses Google's ShieldGemma model for content safety classification.
+Model: google/shieldgemma-2b
+
+Features:
+- Lightweight (2B parameters)
+- Instruction-tuned for safety
+- Yes/No classification
+"""
+
+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
+
+
+class ShieldGemmaBackend(LocalModelBackend):
+    """
+    ShieldGemma safety classification backend.
+
+    Google's ShieldGemma is a lightweight instruction-tuned safety
+    classifier based on the Gemma architecture.
+
+    Features:
+    - Lightweight (2B parameters)
+    - Fast inference
+    - Yes/No classification with optional explanation
+    - Good for edge deployment
+
+    Usage:
+        backend = ShieldGemmaBackend(
+            model=GuardrailModel.SHIELDGEMMA_2B,
+            vllm_url="http://localhost:28000",
+        )
+        results = backend.classify("How to make a bomb?", RailType.INPUT)
+    """
+
+    HF_MODEL_NAME = "google/shieldgemma-2b"
+    MAX_NEW_TOKENS = 64
+    TEMPERATURE = 0.1
+
+    # Safety categories ShieldGemma can detect
+    SAFETY_CATEGORIES = [
+        "dangerous",
+        "harassment",
+        "hate",
+        "sexual",
+        "self_harm",
+    ]
+
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """
+        Format content into ShieldGemma prompt.
+
+        ShieldGemma uses a structured safety check format.
+        """
+        if rail_type == RailType.OUTPUT and context:
+            text_to_check = f"User message: {context}\nAI response: {content}"
+            task = "Check if the AI response is safe and appropriate."
+        elif rail_type == RailType.OUTPUT:
+            text_to_check = f"AI response: {content}"
+            task = "Check if this AI response is safe and appropriate."
+        else:
+            text_to_check = f"User message: {content}"
+            task = "Check if this user message requests harmful content."
+
+        prompt = f"""<start_of_turn>user
+You are a safety classifier. {task}
+
+Content to evaluate:
+{text_to_check}
+
+Does this content violate safety policies? Answer "Yes" if it contains dangerous, harmful, harassing, hateful, sexual, or self-harm content. Answer "No" if it is safe.
+
+Answer with only "Yes" or "No" followed by a brief explanation.<end_of_turn>
+<start_of_turn>model
+"""
+        return prompt
+
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """
+        Parse ShieldGemma response.
+
+        Expected format:
+            Yes - This content contains instructions for violence.
+        or:
+            No - This is a safe request.
+        """
+        response_clean = response.strip()
+        response_lower = response_clean.lower()
+
+        # Check for Yes/No at the start
+        is_violation = response_lower.startswith("yes")
+        is_safe = response_lower.startswith("no")
+
+        if is_safe and not is_violation:
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="safe",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Content passed ShieldGemma safety check",
+                    action="pass",
+                    latency_ms=0.0,
+                )
+            ]
+
+        # Content is unsafe
+        categories = self._infer_categories(content, response_clean)
+
+        results = []
+        for category in categories:
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category=category,
+                    score=1.0,
+                    model=self.model_name,
+                    reason=f"ShieldGemma: {response_clean[:100]}",
+                    action="block",
+                    latency_ms=0.0,
+                )
+            )
+
+        if not results:
+            results.append(
+                GuardrailResult(
+                    passed=False,
+                    category="harmful_content",
+                    score=1.0,
+                    model=self.model_name,
+                    reason="ShieldGemma classified as unsafe",
+                    action="block",
+                    latency_ms=0.0,
+                )
+            )
+
+        return results
+
+    # _infer_categories inherited from LocalModelBackend