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

fi/evals/guardrails/gateway.py

@@ -0,0 +1,437 @@
+"""
+Guardrails Gateway - High-level API for content screening.
+
+Provides a simple, ergonomic interface for screening content
+with automatic backend management and context managers.
+"""
+
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncIterator, Iterator, List, Optional
+
+from fi.evals.guardrails.base import Guardrails
+from fi.evals.guardrails.config import GuardrailModel, GuardrailsConfig, AggregationStrategy
+from fi.evals.guardrails.types import GuardrailsResponse
+from fi.evals.guardrails.discovery import discover_backends, get_backend_details
+
+
+class ScreeningSession:
+    """A screening session for synchronous operations."""
+
+    def __init__(self, guardrails: Guardrails):
+        self._guardrails = guardrails
+        self._history: List[GuardrailsResponse] = []
+
+    def input(self, content: str, metadata: Optional[dict] = None) -> GuardrailsResponse:
+        """Screen user input.
+
+        Args:
+            content: The user input to screen.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            GuardrailsResponse with screening results.
+        """
+        result = self._guardrails.screen_input(content, metadata=metadata)
+        self._history.append(result)
+        return result
+
+    def output(
+        self,
+        content: str,
+        context: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> GuardrailsResponse:
+        """Screen LLM output.
+
+        Args:
+            content: The LLM response to screen.
+            context: Optional original user query for context.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            GuardrailsResponse with screening results.
+        """
+        result = self._guardrails.screen_output(content, context=context, metadata=metadata)
+        self._history.append(result)
+        return result
+
+    def retrieval(
+        self,
+        chunks: List[str],
+        query: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> List[GuardrailsResponse]:
+        """Screen retrieval chunks.
+
+        Args:
+            chunks: List of document chunks to screen.
+            query: Optional user query for context.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            List of GuardrailsResponse, one per chunk.
+        """
+        results = self._guardrails.screen_retrieval(chunks, query=query, metadata=metadata)
+        self._history.extend(results)
+        return results
+
+    @property
+    def history(self) -> List[GuardrailsResponse]:
+        """Get the history of all screening results in this session."""
+        return self._history.copy()
+
+    @property
+    def all_passed(self) -> bool:
+        """Check if all screenings in this session passed."""
+        return all(r.passed for r in self._history)
+
+
+class AsyncScreeningSession:
+    """A screening session for async operations."""
+
+    def __init__(self, guardrails: Guardrails):
+        self._guardrails = guardrails
+        self._history: List[GuardrailsResponse] = []
+
+    async def input(self, content: str, metadata: Optional[dict] = None) -> GuardrailsResponse:
+        """Screen user input asynchronously.
+
+        Args:
+            content: The user input to screen.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            GuardrailsResponse with screening results.
+        """
+        result = await self._guardrails.screen_input_async(content, metadata=metadata)
+        self._history.append(result)
+        return result
+
+    async def output(
+        self,
+        content: str,
+        context: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> GuardrailsResponse:
+        """Screen LLM output asynchronously.
+
+        Args:
+            content: The LLM response to screen.
+            context: Optional original user query for context.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            GuardrailsResponse with screening results.
+        """
+        result = await self._guardrails.screen_output_async(content, context=context, metadata=metadata)
+        self._history.append(result)
+        return result
+
+    async def retrieval(
+        self,
+        chunks: List[str],
+        query: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> List[GuardrailsResponse]:
+        """Screen retrieval chunks asynchronously.
+
+        Args:
+            chunks: List of document chunks to screen.
+            query: Optional user query for context.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            List of GuardrailsResponse, one per chunk.
+        """
+        results = await self._guardrails.screen_retrieval_async(chunks, query=query, metadata=metadata)
+        self._history.extend(results)
+        return results
+
+    async def batch(
+        self,
+        contents: List[str],
+        metadata: Optional[dict] = None,
+    ) -> List[GuardrailsResponse]:
+        """Screen multiple contents in batch asynchronously.
+
+        Args:
+            contents: List of contents to screen.
+            metadata: Optional metadata for the request.
+
+        Returns:
+            List of GuardrailsResponse, one per content.
+        """
+        results = await self._guardrails.screen_batch_async(contents, metadata=metadata)
+        self._history.extend(results)
+        return results
+
+    @property
+    def history(self) -> List[GuardrailsResponse]:
+        """Get the history of all screening results in this session."""
+        return self._history.copy()
+
+    @property
+    def all_passed(self) -> bool:
+        """Check if all screenings in this session passed."""
+        return all(r.passed for r in self._history)
+
+
+class GuardrailsGateway:
+    """High-level gateway for content screening.
+
+    Provides convenient methods for creating screening sessions
+    and managing guardrails configuration.
+
+    Example:
+        # Simple usage
+        gateway = GuardrailsGateway()
+        result = gateway.screen("Hello world")
+
+        # With context manager
+        with gateway.screening() as session:
+            input_result = session.input("user message")
+            if input_result.passed:
+                response = call_llm("user message")
+                output_result = session.output(response)
+
+        # Async context manager
+        async with gateway.screening_async() as session:
+            input_result = await session.input("user message")
+            output_result = await session.output(response)
+    """
+
+    def __init__(
+        self,
+        models: Optional[List[GuardrailModel]] = None,
+        config: Optional[GuardrailsConfig] = None,
+        auto_discover: bool = False,
+    ):
+        """Initialize the gateway.
+
+        Args:
+            models: List of models to use. If None, uses default.
+            config: Full configuration object. Takes precedence over models.
+            auto_discover: If True and no models specified, auto-discover available backends.
+        """
+        if config:
+            self._config = config
+        elif models:
+            self._config = GuardrailsConfig(models=models)
+        elif auto_discover:
+            available = discover_backends()
+            if not available:
+                raise ValueError("No backends available. Set API keys or start VLLM server.")
+            self._config = GuardrailsConfig(models=available[:1])  # Use first available
+        else:
+            self._config = GuardrailsConfig()
+
+        self._guardrails = Guardrails(config=self._config)
+
+    @classmethod
+    def with_openai(cls, api_key: Optional[str] = None) -> "GuardrailsGateway":
+        """Create a gateway using OpenAI Moderation (FREE).
+
+        Args:
+            api_key: OpenAI API key. If None, uses OPENAI_API_KEY env var.
+
+        Returns:
+            GuardrailsGateway configured for OpenAI.
+        """
+        import os
+        if api_key:
+            os.environ["OPENAI_API_KEY"] = api_key
+
+        config = GuardrailsConfig(
+            models=[GuardrailModel.OPENAI_MODERATION],
+            timeout_ms=30000,
+        )
+        return cls(config=config)
+
+    @classmethod
+    def with_azure(
+        cls,
+        endpoint: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ) -> "GuardrailsGateway":
+        """Create a gateway using Azure Content Safety.
+
+        Args:
+            endpoint: Azure endpoint URL. If None, uses env var.
+            api_key: Azure API key. If None, uses env var.
+
+        Returns:
+            GuardrailsGateway configured for Azure.
+        """
+        import os
+        if endpoint:
+            os.environ["AZURE_CONTENT_SAFETY_ENDPOINT"] = endpoint
+        if api_key:
+            os.environ["AZURE_CONTENT_SAFETY_KEY"] = api_key
+
+        config = GuardrailsConfig(
+            models=[GuardrailModel.AZURE_CONTENT_SAFETY],
+            timeout_ms=30000,
+        )
+        return cls(config=config)
+
+    @classmethod
+    def with_local_model(
+        cls,
+        model: GuardrailModel,
+        vllm_url: Optional[str] = None,
+    ) -> "GuardrailsGateway":
+        """Create a gateway using a local model via VLLM.
+
+        Args:
+            model: The local model to use (e.g., GuardrailModel.WILDGUARD_7B).
+            vllm_url: VLLM server URL. If None, uses env var.
+
+        Returns:
+            GuardrailsGateway configured for the local model.
+        """
+        import os
+        if vllm_url:
+            os.environ["VLLM_SERVER_URL"] = vllm_url
+
+        config = GuardrailsConfig(
+            models=[model],
+            timeout_ms=60000,  # Local models may be slower
+        )
+        return cls(config=config)
+
+    @classmethod
+    def with_ensemble(
+        cls,
+        models: List[GuardrailModel],
+        aggregation: AggregationStrategy = AggregationStrategy.ANY,
+        parallel: bool = True,
+    ) -> "GuardrailsGateway":
+        """Create a gateway with ensemble of multiple backends.
+
+        Args:
+            models: List of models to use.
+            aggregation: How to combine results (ANY, ALL, MAJORITY).
+            parallel: Whether to run backends in parallel.
+
+        Returns:
+            GuardrailsGateway configured for ensemble mode.
+        """
+        config = GuardrailsConfig(
+            models=models,
+            aggregation=aggregation,
+            parallel=parallel,
+            timeout_ms=60000,
+        )
+        return cls(config=config)
+
+    @classmethod
+    def auto(cls) -> "GuardrailsGateway":
+        """Create a gateway that auto-discovers available backends.
+
+        Returns:
+            GuardrailsGateway with the best available backend.
+
+        Raises:
+            ValueError: If no backends are available.
+        """
+        return cls(auto_discover=True)
+
+    def screen(self, content: str, metadata: Optional[dict] = None) -> GuardrailsResponse:
+        """Quick screen content (alias for screen_input).
+
+        Args:
+            content: Content to screen.
+            metadata: Optional metadata.
+
+        Returns:
+            GuardrailsResponse with results.
+        """
+        return self._guardrails.screen_input(content, metadata=metadata)
+
+    async def screen_async(self, content: str, metadata: Optional[dict] = None) -> GuardrailsResponse:
+        """Quick screen content asynchronously.
+
+        Args:
+            content: Content to screen.
+            metadata: Optional metadata.
+
+        Returns:
+            GuardrailsResponse with results.
+        """
+        return await self._guardrails.screen_input_async(content, metadata=metadata)
+
+    @contextmanager
+    def screening(self) -> Iterator[ScreeningSession]:
+        """Create a synchronous screening session.
+
+        Yields:
+            ScreeningSession for screening operations.
+
+        Example:
+            with gateway.screening() as session:
+                input_result = session.input("user message")
+                if input_result.passed:
+                    response = generate_response()
+                    output_result = session.output(response)
+        """
+        session = ScreeningSession(self._guardrails)
+        yield session
+
+    @asynccontextmanager
+    async def screening_async(self) -> AsyncIterator[AsyncScreeningSession]:
+        """Create an async screening session.
+
+        Yields:
+            AsyncScreeningSession for async screening operations.
+
+        Example:
+            async with gateway.screening_async() as session:
+                input_result = await session.input("user message")
+                if input_result.passed:
+                    response = await generate_response()
+                    output_result = await session.output(response)
+        """
+        session = AsyncScreeningSession(self._guardrails)
+        yield session
+
+    @property
+    def available_backends(self) -> List[GuardrailModel]:
+        """Get list of available backends."""
+        return discover_backends()
+
+    @property
+    def configured_models(self) -> List[GuardrailModel]:
+        """Get list of configured models."""
+        return self._config.models
+
+    @staticmethod
+    def discover() -> List[GuardrailModel]:
+        """Discover available backends.
+
+        Returns:
+            List of available GuardrailModel values.
+        """
+        return discover_backends()
+
+    @staticmethod
+    def get_details() -> dict:
+        """Get detailed backend availability info.
+
+        Returns:
+            Dict mapping model names to availability details.
+        """
+        return get_backend_details()
+
+
+# Convenience alias
+Gateway = GuardrailsGateway
+
+
+__all__ = [
+    "GuardrailsGateway",
+    "Gateway",
+    "ScreeningSession",
+    "AsyncScreeningSession",
+    "discover_backends",
+    "get_backend_details",
+]

fi/evals/guardrails/registry.py

@@ -0,0 +1,231 @@
+"""
+Model Registry for Guardrails.
+
+Central registry for all supported guardrail models with metadata
+about backends, model types, and requirements.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Type
+
+from fi.evals.guardrails.config import GuardrailModel
+
+
+@dataclass
+class ModelInfo:
+    """Information about a guardrail model."""
+    model: GuardrailModel
+    backend_class: str  # String reference to avoid circular imports
+    backend_module: str  # Module containing the backend class
+    model_type: str  # "api", "local", or "vllm"
+    hf_model_name: Optional[str] = None  # HuggingFace model name
+    vram_required_gb: Optional[float] = None  # Minimum VRAM in GB
+    description: str = ""
+    is_gated: bool = False  # Requires HF token
+
+
+# Central registry of all supported models
+MODEL_REGISTRY: Dict[str, ModelInfo] = {
+    # Turing (FutureAGI API)
+    "turing_flash": ModelInfo(
+        model=GuardrailModel.TURING_FLASH,
+        backend_class="TuringBackend",
+        backend_module="fi.evals.guardrails.backends.turing",
+        model_type="api",
+        description="Fast binary classification via FutureAGI API",
+    ),
+    "turing_safety": ModelInfo(
+        model=GuardrailModel.TURING_SAFETY,
+        backend_class="TuringBackend",
+        backend_module="fi.evals.guardrails.backends.turing",
+        model_type="api",
+        description="Detailed safety analysis via FutureAGI API",
+    ),
+
+    # OpenAI (Free API)
+    "openai-moderation": ModelInfo(
+        model=GuardrailModel.OPENAI_MODERATION,
+        backend_class="OpenAIBackend",
+        backend_module="fi.evals.guardrails.backends.openai",
+        model_type="api",
+        description="OpenAI Moderation API (FREE, 13 categories)",
+    ),
+
+    # Azure (Paid API)
+    "azure-content-safety": ModelInfo(
+        model=GuardrailModel.AZURE_CONTENT_SAFETY,
+        backend_class="AzureBackend",
+        backend_module="fi.evals.guardrails.backends.azure",
+        model_type="api",
+        description="Azure Content Safety API (4 categories)",
+    ),
+
+    # WildGuard (Local)
+    "wildguard-7b": ModelInfo(
+        model=GuardrailModel.WILDGUARD_7B,
+        backend_class="WildGuardBackend",
+        backend_module="fi.evals.guardrails.backends.wildguard",
+        model_type="local",
+        hf_model_name="allenai/wildguard",
+        vram_required_gb=8.0,
+        description="AllenAI WildGuard safety classifier",
+        is_gated=True,
+    ),
+
+    # LlamaGuard (Local)
+    "llamaguard-3-8b": ModelInfo(
+        model=GuardrailModel.LLAMAGUARD_3_8B,
+        backend_class="LlamaGuardBackend",
+        backend_module="fi.evals.guardrails.backends.llamaguard",
+        model_type="local",
+        hf_model_name="meta-llama/Llama-Guard-3-8B",
+        vram_required_gb=16.0,
+        description="Meta LlamaGuard 3 8B safety classifier",
+        is_gated=True,
+    ),
+    "llamaguard-3-1b": ModelInfo(
+        model=GuardrailModel.LLAMAGUARD_3_1B,
+        backend_class="LlamaGuardBackend",
+        backend_module="fi.evals.guardrails.backends.llamaguard",
+        model_type="local",
+        hf_model_name="meta-llama/Llama-Guard-3-1B",
+        vram_required_gb=4.0,
+        description="Meta LlamaGuard 3 1B safety classifier (lightweight)",
+        is_gated=True,
+    ),
+
+    # Granite Guardian (Local)
+    "granite-guardian-3.3-8b": ModelInfo(
+        model=GuardrailModel.GRANITE_GUARDIAN_8B,
+        backend_class="GraniteGuardianBackend",
+        backend_module="fi.evals.guardrails.backends.granite",
+        model_type="local",
+        hf_model_name="ibm-granite/granite-guardian-3.3-8b",
+        vram_required_gb=16.0,
+        description="IBM Granite Guardian 3.3 8B",
+    ),
+    "granite-guardian-3.2-5b": ModelInfo(
+        model=GuardrailModel.GRANITE_GUARDIAN_5B,
+        backend_class="GraniteGuardianBackend",
+        backend_module="fi.evals.guardrails.backends.granite",
+        model_type="local",
+        hf_model_name="ibm-granite/granite-guardian-3.2-5b",
+        vram_required_gb=10.0,
+        description="IBM Granite Guardian 3.2 5B (lightweight)",
+    ),
+
+    # Qwen3Guard (Local)
+    "qwen3guard-8b": ModelInfo(
+        model=GuardrailModel.QWEN3GUARD_8B,
+        backend_class="Qwen3GuardBackend",
+        backend_module="fi.evals.guardrails.backends.qwen",
+        model_type="local",
+        hf_model_name="Qwen/Qwen3Guard-8B",
+        vram_required_gb=16.0,
+        description="Alibaba Qwen3Guard 8B (119 languages)",
+    ),
+    "qwen3guard-4b": ModelInfo(
+        model=GuardrailModel.QWEN3GUARD_4B,
+        backend_class="Qwen3GuardBackend",
+        backend_module="fi.evals.guardrails.backends.qwen",
+        model_type="local",
+        hf_model_name="Qwen/Qwen3Guard-4B",
+        vram_required_gb=8.0,
+        description="Alibaba Qwen3Guard 4B (lightweight, 119 languages)",
+    ),
+
+    "qwen3guard-0.6b": ModelInfo(
+        model=GuardrailModel.QWEN3GUARD_0_6B,
+        backend_class="Qwen3GuardBackend",
+        backend_module="fi.evals.guardrails.backends.qwen",
+        model_type="local",
+        hf_model_name="Qwen/Qwen3Guard-0.6B",
+        vram_required_gb=1.0,
+        description="Alibaba Qwen3Guard 0.6B (ultra-lightweight, 119 languages)",
+    ),
+
+    # Generic LLM as guard (prompted for safety classification)
+    "llama3.2-3b": ModelInfo(
+        model=GuardrailModel.LLAMA_3_2_3B,
+        backend_class="GenericLLMGuardBackend",
+        backend_module="fi.evals.guardrails.backends.generic_llm",
+        model_type="local",
+        hf_model_name="meta-llama/Llama-3.2-3B-Instruct",
+        vram_required_gb=4.0,
+        description="Llama 3.2 3B as prompted safety classifier",
+    ),
+
+    # ShieldGemma (Local)
+    "shieldgemma-2b": ModelInfo(
+        model=GuardrailModel.SHIELDGEMMA_2B,
+        backend_class="ShieldGemmaBackend",
+        backend_module="fi.evals.guardrails.backends.shieldgemma",
+        model_type="local",
+        hf_model_name="google/shieldgemma-2b",
+        vram_required_gb=4.0,
+        description="Google ShieldGemma 2B (lightweight, fast)",
+    ),
+}
+
+
+def get_model_info(model: GuardrailModel) -> Optional[ModelInfo]:
+    """
+    Get model information from registry.
+
+    Args:
+        model: GuardrailModel enum value
+
+    Returns:
+        ModelInfo or None if not found
+    """
+    return MODEL_REGISTRY.get(model.value)
+
+
+def get_backend_class(model: GuardrailModel) -> Type:
+    """
+    Get the backend class for a model.
+
+    Args:
+        model: GuardrailModel enum value
+
+    Returns:
+        Backend class
+
+    Raises:
+        ValueError: If model not found in registry
+        ImportError: If backend module cannot be imported
+    """
+    info = get_model_info(model)
+    if not info:
+        raise ValueError(f"Model {model.value} not found in registry")
+
+    # Dynamic import to avoid circular dependencies
+    import importlib
+    module = importlib.import_module(info.backend_module)
+    return getattr(module, info.backend_class)
+
+
+def list_models(model_type: Optional[str] = None) -> List[ModelInfo]:
+    """
+    List all models in registry.
+
+    Args:
+        model_type: Filter by type ("api", "local", "vllm")
+
+    Returns:
+        List of ModelInfo objects
+    """
+    models = list(MODEL_REGISTRY.values())
+    if model_type:
+        models = [m for m in models if m.model_type == model_type]
+    return models
+
+
+def list_api_models() -> List[ModelInfo]:
+    """List all API-based models."""
+    return list_models(model_type="api")
+
+
+def list_local_models() -> List[ModelInfo]:
+    """List all local models."""
+    return list_models(model_type="local")