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

fi/evals/guardrails/backends/turing.py

@@ -0,0 +1,235 @@
+"""
+Turing Backend for Guardrails.
+
+Uses FutureAGI's Turing API for content screening.
+- turing_flash: Fast binary classification
+- turing_safety: Higher accuracy, detailed categories
+"""
+
+import time
+from typing import Dict, List, Optional, Any
+
+from fi.evals.guardrails.backends.base import BaseBackend
+from fi.evals.guardrails.config import RailType, GuardrailModel
+from fi.evals.guardrails.types import GuardrailResult
+from fi.evals.protect import Protect
+from fi.utils.utils import get_keys_from_env
+
+
+class TuringBackend(BaseBackend):
+    """
+    Backend for Turing models (turing_flash, turing_safety).
+
+    Uses the existing Protect class for API calls to FutureAGI.
+    """
+
+    def __init__(
+        self,
+        model: GuardrailModel,
+        fi_api_key: Optional[str] = None,
+        fi_secret_key: Optional[str] = None,
+        fi_base_url: Optional[str] = None,
+    ):
+        """
+        Initialize Turing backend.
+
+        Args:
+            model: The Turing model to use
+            fi_api_key: FutureAGI API key
+            fi_secret_key: FutureAGI secret key
+            fi_base_url: Base URL for API
+        """
+        super().__init__(model)
+
+        # Get keys from environment if not provided
+        if not fi_api_key or not fi_secret_key:
+            env_api_key, env_secret_key = get_keys_from_env()
+            fi_api_key = fi_api_key or env_api_key
+            fi_secret_key = fi_secret_key or env_secret_key
+
+        self._protect = Protect(
+            fi_api_key=fi_api_key,
+            fi_secret_key=fi_secret_key,
+            fi_base_url=fi_base_url,
+        )
+
+        # Determine if this is flash or safety mode
+        self._use_flash = model == GuardrailModel.TURING_FLASH
+
+    def classify(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """
+        Classify content using Turing API.
+
+        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()
+        results = []
+
+        # Handle empty or whitespace-only 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:
+            if self._use_flash:
+                # Use ProtectFlash for fast binary classification
+                response = self._protect.protect(
+                    inputs=content,
+                    use_flash=True,
+                    timeout=30000,
+                )
+            else:
+                # Use full Protect with content moderation rules
+                response = self._protect.protect(
+                    inputs=content,
+                    protect_rules=[
+                        {"metric": "content_moderation"},
+                        {"metric": "security"},
+                        {"metric": "bias_detection"},
+                    ],
+                    timeout=30000,
+                )
+
+            elapsed_ms = (time.time() - start_time) * 1000
+
+            # Parse response
+            is_harmful = response.get("status") == "failed"
+            failed_rule = response.get("failed_rule", [])
+            reasons = response.get("reasons", [])
+
+            if is_harmful:
+                # Determine categories from failed rules
+                categories = self._extract_categories(failed_rule, reasons)
+                for category in categories:
+                    results.append(
+                        GuardrailResult(
+                            passed=False,
+                            category=category,
+                            score=1.0,
+                            model=self.model_name,
+                            reason=reasons[0] if reasons else None,
+                            action="block",
+                            latency_ms=elapsed_ms,
+                        )
+                    )
+
+                # If no specific categories, use generic harmful_content
+                if not results:
+                    results.append(
+                        GuardrailResult(
+                            passed=False,
+                            category="harmful_content",
+                            score=1.0,
+                            model=self.model_name,
+                            reason=reasons[0] if reasons else "Content flagged as harmful",
+                            action="block",
+                            latency_ms=elapsed_ms,
+                        )
+                    )
+            else:
+                # Content passed
+                results.append(
+                    GuardrailResult(
+                        passed=True,
+                        category="safe",
+                        score=0.0,
+                        model=self.model_name,
+                        reason="Content passed all checks",
+                        action="pass",
+                        latency_ms=elapsed_ms,
+                    )
+                )
+
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            results.append(
+                GuardrailResult(
+                    passed=False,  # Fail closed by default
+                    category="error",
+                    score=0.0,
+                    model=self.model_name,
+                    reason=f"Error during classification: {str(e)}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            )
+
+        return results
+
+    def _extract_categories(
+        self,
+        failed_rules: List[str],
+        reasons: List[str],
+    ) -> List[str]:
+        """
+        Extract category names from failed rules and reasons.
+
+        Args:
+            failed_rules: List of failed rule names
+            reasons: List of reason strings
+
+        Returns:
+            List of category names
+        """
+        categories = set()
+
+        # Map rule names to categories
+        rule_to_category = {
+            "content_moderation": "toxicity",
+            "security": "prompt_injection",
+            "bias_detection": "hate_speech",
+            "data_privacy_compliance": "pii",
+            "ProtectFlash": "harmful_content",
+        }
+
+        if isinstance(failed_rules, list):
+            for rule in failed_rules:
+                if rule in rule_to_category:
+                    categories.add(rule_to_category[rule])
+                else:
+                    categories.add(rule.lower().replace(" ", "_"))
+        elif isinstance(failed_rules, str):
+            if failed_rules in rule_to_category:
+                categories.add(rule_to_category[failed_rules])
+            else:
+                categories.add(failed_rules.lower().replace(" ", "_"))
+
+        # Also parse reasons for additional context
+        for reason in reasons:
+            reason_lower = reason.lower() if reason else ""
+            if "jailbreak" in reason_lower:
+                categories.add("jailbreak")
+            if "injection" in reason_lower:
+                categories.add("prompt_injection")
+            if "harm" in reason_lower or "violence" in reason_lower:
+                categories.add("violence")
+            if "hate" in reason_lower:
+                categories.add("hate_speech")
+            if "self-harm" in reason_lower or "suicide" in reason_lower:
+                categories.add("self_harm")
+            if "sexual" in reason_lower:
+                categories.add("sexual_content")
+
+        return list(categories) if categories else ["harmful_content"]

fi/evals/guardrails/backends/vllm_client.py

@@ -0,0 +1,321 @@
+"""
+VLLM Client for Guardrails.
+
+Provides a client for OpenAI-compatible VLLM servers.
+Used by local model backends to communicate with VLLM instances.
+"""
+
+import os
+import time
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+import httpx
+
+
+@dataclass
+class VLLMResponse:
+    """Response from VLLM server."""
+    text: str
+    model: str
+    usage: Dict[str, int]
+    latency_ms: float
+
+
+class VLLMClient:
+    """
+    Client for OpenAI-compatible VLLM servers.
+
+    Supports both sync and async operations.
+
+    Usage:
+        client = VLLMClient(base_url="http://localhost:28000")
+
+        # Check health
+        if client.health_check():
+            response = client.generate("prompt", max_tokens=128)
+            print(response.text)
+
+        # Async usage
+        response = await client.generate_async("prompt", max_tokens=128)
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        model: Optional[str] = None,
+        timeout: float = 60.0,
+    ):
+        """
+        Initialize VLLM client.
+
+        Args:
+            base_url: VLLM server base URL (e.g., http://localhost:28000)
+            model: Model name (if not specified, uses server default)
+            timeout: Request timeout in seconds
+        """
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+        self.timeout = timeout
+
+    def health_check(self) -> bool:
+        """
+        Check if VLLM server is healthy.
+
+        Returns:
+            True if server is healthy, False otherwise
+        """
+        try:
+            with httpx.Client(timeout=5.0) as client:
+                # Try /health (VLLM), fall back to / (ollama)
+                response = client.get(f"{self.base_url}/health")
+                if response.status_code == 200:
+                    return True
+                response = client.get(self.base_url)
+                return response.status_code == 200
+        except Exception:
+            return False
+
+    async def health_check_async(self) -> bool:
+        """Async version of health_check."""
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                response = await client.get(f"{self.base_url}/health")
+                return response.status_code == 200
+        except Exception:
+            return False
+
+    def get_models(self) -> List[str]:
+        """
+        Get list of available models from VLLM server.
+
+        Returns:
+            List of model names
+        """
+        try:
+            with httpx.Client(timeout=self.timeout) as client:
+                response = client.get(f"{self.base_url}/v1/models")
+                if response.status_code == 200:
+                    data = response.json()
+                    return [m["id"] for m in data.get("data", [])]
+        except Exception:
+            pass
+        return []
+
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 128,
+        temperature: float = 0.1,
+        stop: Optional[List[str]] = None,
+    ) -> VLLMResponse:
+        """
+        Generate completion from VLLM server.
+
+        Args:
+            prompt: Input prompt
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            stop: Stop sequences
+
+        Returns:
+            VLLMResponse with generated text
+        """
+        start_time = time.time()
+
+        # Determine model to use
+        model = self.model or self._get_default_model()
+
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+        if stop:
+            payload["stop"] = stop
+
+        with httpx.Client(timeout=self.timeout) as client:
+            response = client.post(
+                f"{self.base_url}/v1/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        elapsed_ms = (time.time() - start_time) * 1000
+
+        return VLLMResponse(
+            text=data["choices"][0]["text"],
+            model=data.get("model", model),
+            usage=data.get("usage", {}),
+            latency_ms=elapsed_ms,
+        )
+
+    async def generate_async(
+        self,
+        prompt: str,
+        max_tokens: int = 128,
+        temperature: float = 0.1,
+        stop: Optional[List[str]] = None,
+    ) -> VLLMResponse:
+        """
+        Async version of generate.
+
+        Args:
+            prompt: Input prompt
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            stop: Stop sequences
+
+        Returns:
+            VLLMResponse with generated text
+        """
+        start_time = time.time()
+
+        model = self.model or await self._get_default_model_async()
+
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+        if stop:
+            payload["stop"] = stop
+
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.post(
+                f"{self.base_url}/v1/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        elapsed_ms = (time.time() - start_time) * 1000
+
+        return VLLMResponse(
+            text=data["choices"][0]["text"],
+            model=data.get("model", model),
+            usage=data.get("usage", {}),
+            latency_ms=elapsed_ms,
+        )
+
+    def chat(
+        self,
+        messages: List[Dict[str, str]],
+        max_tokens: int = 128,
+        temperature: float = 0.1,
+    ) -> VLLMResponse:
+        """
+        Generate chat completion from VLLM server.
+
+        Args:
+            messages: List of chat messages [{"role": "user", "content": "..."}]
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+
+        Returns:
+            VLLMResponse with generated text
+        """
+        start_time = time.time()
+
+        model = self.model or self._get_default_model()
+
+        payload = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+
+        with httpx.Client(timeout=self.timeout) as client:
+            response = client.post(
+                f"{self.base_url}/v1/chat/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        elapsed_ms = (time.time() - start_time) * 1000
+
+        return VLLMResponse(
+            text=data["choices"][0]["message"]["content"],
+            model=data.get("model", model),
+            usage=data.get("usage", {}),
+            latency_ms=elapsed_ms,
+        )
+
+    async def chat_async(
+        self,
+        messages: List[Dict[str, str]],
+        max_tokens: int = 128,
+        temperature: float = 0.1,
+    ) -> VLLMResponse:
+        """Async version of chat."""
+        start_time = time.time()
+
+        model = self.model or await self._get_default_model_async()
+
+        payload = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.post(
+                f"{self.base_url}/v1/chat/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        elapsed_ms = (time.time() - start_time) * 1000
+
+        return VLLMResponse(
+            text=data["choices"][0]["message"]["content"],
+            model=data.get("model", model),
+            usage=data.get("usage", {}),
+            latency_ms=elapsed_ms,
+        )
+
+    def _get_default_model(self) -> str:
+        """Get default model from server."""
+        models = self.get_models()
+        return models[0] if models else "default"
+
+    async def _get_default_model_async(self) -> str:
+        """Async version of _get_default_model."""
+        try:
+            async with httpx.AsyncClient(timeout=self.timeout) as client:
+                response = await client.get(f"{self.base_url}/v1/models")
+                if response.status_code == 200:
+                    data = response.json()
+                    models = [m["id"] for m in data.get("data", [])]
+                    return models[0] if models else "default"
+        except Exception:
+            pass
+        return "default"
+
+
+def get_vllm_url(model_name: str) -> Optional[str]:
+    """
+    Get VLLM server URL for a model from environment variables.
+
+    Checks for model-specific env var first, then falls back to generic.
+
+    Args:
+        model_name: Model name (e.g., "wildguard-7b")
+
+    Returns:
+        VLLM server URL or None
+    """
+    # Try model-specific env var (e.g., VLLM_WILDGUARD_7B_URL)
+    env_var = f"VLLM_{model_name.upper().replace('-', '_')}_URL"
+    url = os.environ.get(env_var)
+    if url:
+        return url
+
+    # Fall back to generic VLLM_SERVER_URL
+    return os.environ.get("VLLM_SERVER_URL")