netra-sdk 0.1.0__py3-none-any.whl

netra/decorators.py

@@ -0,0 +1,167 @@
+"""Netra decorator utilities.
+
+This module provides decorators for common patterns in Netra SDK.
+Decorators can be applied to both functions and classes.
+"""
+
+import functools
+import inspect
+import json
+from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, TypeVar, Union, cast
+
+from opentelemetry import trace
+
+from .config import Config
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+F_Callable = TypeVar("F_Callable", bound=Callable[..., Any])
+C = TypeVar("C", bound=type)
+
+
+def _serialize_value(value: Any) -> str:
+    """Safely serialize a value to string for span attributes."""
+    try:
+        if isinstance(value, (str, int, float, bool, type(None))):
+            return str(value)
+        elif isinstance(value, (list, dict, tuple)):
+            return json.dumps(value, default=str)[:1000]  # Limit size
+        else:
+            return str(value)[:1000]  # Limit size
+    except Exception:
+        return str(type(value).__name__)
+
+
+def _add_span_attributes(
+    span: trace.Span, func: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any], entity_type: str
+) -> None:
+    """Helper function to add span attributes from function parameters."""
+    span.set_attribute(f"{Config.LIBRARY_NAME}.entity.type", entity_type)
+
+    try:
+        sig = inspect.signature(func)
+        param_names = list(sig.parameters.keys())
+        input_data = {}
+
+        for i, arg in enumerate(args):
+            if i < len(param_names):
+                param_name = param_names[i]
+                if param_name not in ("self", "cls"):
+                    input_data[param_name] = _serialize_value(arg)
+
+        for key, value in kwargs.items():
+            input_data[key] = _serialize_value(value)
+
+        if input_data:
+            span.set_attribute(f"{Config.LIBRARY_NAME}.entity.input", json.dumps(input_data))
+
+    except Exception as e:
+        span.set_attribute(f"{Config.LIBRARY_NAME}.input_error", str(e))
+
+
+def _add_output_attributes(span: trace.Span, result: Any) -> None:
+    """Helper function to add output attributes to span."""
+    try:
+        serialized_output = _serialize_value(result)
+        span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output", serialized_output)
+    except Exception as e:
+        span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output_error", str(e))
+
+
+def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optional[str] = None) -> Callable[P, R]:
+    module_name = func.__name__
+    is_async = inspect.iscoroutinefunction(func)
+    span_name = name if name is not None else func.__name__
+
+    if is_async:
+
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            tracer = trace.get_tracer(module_name)
+            with tracer.start_as_current_span(span_name) as span:
+                _add_span_attributes(span, func, args, kwargs, entity_type)
+                try:
+                    result = await cast(Awaitable[Any], func(*args, **kwargs))
+                    _add_output_attributes(span, result)
+                    return result
+                except Exception as e:
+                    span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                    span.record_exception(e)
+                    raise
+
+        return cast(Callable[P, R], async_wrapper)
+
+    else:
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            tracer = trace.get_tracer(module_name)
+            with tracer.start_as_current_span(span_name) as span:
+                _add_span_attributes(span, func, args, kwargs, entity_type)
+                try:
+                    result = func(*args, **kwargs)
+                    _add_output_attributes(span, result)
+                    return result
+                except Exception as e:
+                    span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                    span.record_exception(e)
+                    raise
+
+        return cast(Callable[P, R], sync_wrapper)
+
+
+def _wrap_class_methods(cls: C, entity_type: str, name: Optional[str] = None) -> C:
+    class_name = name if name is not None else cls.__name__
+    for attr_name in cls.__dict__:
+        attr = getattr(cls, attr_name)
+        if attr_name.startswith("_"):
+            continue
+        if callable(attr) and inspect.isfunction(attr):
+            method_span_name = f"{class_name}.{attr_name}"
+            wrapped_method = _create_function_wrapper(attr, entity_type, method_span_name)
+            setattr(cls, attr_name, wrapped_method)
+    return cls
+
+
+def workflow(
+    target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
+) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
+    def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
+        if inspect.isclass(obj):
+            return _wrap_class_methods(cast(C, obj), "workflow", name)
+        else:
+            return _create_function_wrapper(cast(Callable[P, R], obj), "workflow", name)
+
+    if target is not None:
+        return decorator(target)
+    return decorator
+
+
+def agent(
+    target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
+) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
+    def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
+        if inspect.isclass(obj):
+            return _wrap_class_methods(cast(C, obj), "agent", name)
+        else:
+            return _create_function_wrapper(cast(Callable[P, R], obj), "agent", name)
+
+    if target is not None:
+        return decorator(target)
+    return decorator
+
+
+def task(
+    target: Union[Callable[P, R], C, None] = None, *, name: Optional[str] = None
+) -> Union[Callable[P, R], C, Callable[[Callable[P, R]], Callable[P, R]]]:
+    def decorator(obj: Union[Callable[P, R], C]) -> Union[Callable[P, R], C]:
+        if inspect.isclass(obj):
+            return _wrap_class_methods(cast(C, obj), "task", name)
+        else:
+            # When obj is a function, it should be type Callable[P, R]
+            return _create_function_wrapper(cast(Callable[P, R], obj), "task", name)
+
+    if target is not None:
+        return decorator(target)
+    return decorator

netra/exceptions/__init__.py

@@ -0,0 +1,6 @@
+# File: netra/exceptions/__init__.py
+
+from .injection import InjectionException
+from .pii import PIIBlockedException
+
+__all__ = ["PIIBlockedException", "InjectionException"]

netra/exceptions/injection.py

@@ -0,0 +1,33 @@
+# File: netra/exceptions/injection.py
+
+from typing import Dict, List, Optional
+
+
+class InjectionException(Exception):
+    """
+    Raised when prompt injection is detected in input and blocking is enabled.
+
+    Attributes:
+        message (str): Human-readable explanation of why blocking occurred.
+        has_violation (bool): True if prompt injection was detected in the provided text.
+        violations (List[str]): List of violation types that were detected.
+        is_blocked (bool): True if blocking is enabled and prompt injection was detected.
+        violation_actions (Dict[str, List[str]]): Dictionary mapping action types to lists of violations.
+    """
+
+    def __init__(
+        self,
+        message: str = "Input blocked due to detected injection.",
+        has_violation: bool = True,
+        violations: Optional[List[str]] = None,
+        is_blocked: bool = True,
+        violation_actions: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        # Always pass the message to the base Exception constructor
+        super().__init__(message)
+
+        # Store structured attributes
+        self.has_violation: bool = has_violation
+        self.violations: List[str] = violations or []
+        self.is_blocked: bool = is_blocked
+        self.violation_actions: Dict[str, List[str]] = violation_actions or {}

netra/exceptions/pii.py

@@ -0,0 +1,46 @@
+# File: netra/exceptions/pii.py
+
+from typing import Any, Dict, List, Optional, Union
+
+
+class PIIBlockedException(Exception):
+    """
+    Raised when PII is detected in input and blocking is enabled.
+
+    Attributes:
+        message (str): Human-readable explanation of why blocking occurred.
+        has_pii (bool): True if PII was detected in the provided text.
+        pii_entities (Dict[str, int]): Mapping from PII label to number of occurrences.
+        masked_text (Union[str, List[Dict[str, str]], List[Any], None]): Input text after masking PII spans.
+            Can be a string for simple inputs, a list of dicts for chat messages,
+            or a list of BaseMessage objects for LangChain inputs.
+        is_blocked (bool): True if blocking is enabled and PII was detected.
+        pii_actions (Dict[str, List[str]]): Dictionary mapping action types to lists of PII entities.
+        original_text (Union[str, List[Dict[str, str]], List[str], List[Any], None]): The original text used to call the detect() method.
+            Can be a string, list of strings, list of dictionaries, or any other type.
+        hashed_entities (Dict[str, str]): Dictionary mapping hashed entity values to their original values.
+            Only populated when using Anonymizer for masking.
+    """
+
+    def __init__(
+        self,
+        message: str = "Input blocked due to detected PII.",
+        has_pii: bool = True,
+        pii_entities: Optional[Dict[str, int]] = None,
+        masked_text: Optional[Union[str, List[Dict[str, str]], List[Any]]] = None,
+        pii_actions: Optional[Dict[Any, List[str]]] = None,
+        is_blocked: bool = True,
+        original_text: Optional[Union[str, List[Dict[str, str]], List[str], List[Any]]] = None,
+        hashed_entities: Optional[Dict[str, str]] = None,
+    ) -> None:
+        # Always pass the message to the base Exception constructor
+        super().__init__(message)
+
+        # Store structured attributes
+        self.has_pii: bool = has_pii
+        self.pii_entities: Dict[str, int] = pii_entities or {}
+        self.masked_text: Optional[Union[str, List[Dict[str, str]], List[Any]]] = masked_text
+        self.pii_actions: Dict[Any, List[str]] = pii_actions or {}
+        self.is_blocked: bool = is_blocked
+        self.original_text: Optional[Union[str, List[Dict[str, str]], List[str], List[Any]]] = original_text
+        self.hashed_entities: Dict[str, str] = hashed_entities or {}

netra/input_scanner.py

@@ -0,0 +1,142 @@
+"""
+Input Scanner module for Netra SDK to implement LLM guard scanning options.
+
+This module provides a unified interface for scanning input prompts using
+various scanner implementations.
+"""
+
+import json
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Union
+
+from netra import Netra
+from netra.exceptions import InjectionException
+from netra.scanner import Scanner
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ScanResult:
+    """
+    Result of running input scanning on prompts.
+
+    Attributes:
+        has_violation: True if any violations were detected
+        violations: List of violation types that were detected
+        is_blocked: True if the input should be blocked
+        violation_actions: Dictionary mapping action types to lists of violations
+    """
+
+    has_violation: bool = False
+    violations: List[str] = field(default_factory=list)
+    is_blocked: bool = False
+    violation_actions: Dict[str, List[str]] = field(default_factory=dict)
+
+
+class ScannerType(Enum):
+    """
+    Enum representing the available scanner types.
+    """
+
+    PROMPT_INJECTION = "prompt_injection"
+
+
+class InputScanner:
+    """
+    A factory class for creating input scanners.
+    """
+
+    def __init__(self, scanner_types: List[Union[str, ScannerType]] = [ScannerType.PROMPT_INJECTION]):
+        self.scanner_types = scanner_types
+
+    @staticmethod
+    def _get_scanner(scanner_type: Union[str, ScannerType], **kwargs: Any) -> Scanner:
+        """
+        Factory function to get a scanner instance based on the specified type.
+
+        Args:
+            scanner_type: The type of scanner to create (e.g., "prompt_injection" or ScannerType.PROMPT_INJECTION)
+            **kwargs: Additional parameters to pass to the scanner constructor
+
+        Returns:
+            Scanner: An instance of the appropriate scanner
+
+        Raises:
+            ValueError: If the specified scanner type is not supported
+        """
+        if isinstance(scanner_type, ScannerType):
+            scanner_type = scanner_type.value
+
+        if scanner_type == ScannerType.PROMPT_INJECTION.value:
+            match_type = None
+            try:
+                # Try to import from llm_guard if available
+                from llm_guard.input_scanners.prompt_injection import MatchType
+
+                match_type = kwargs.get("match_type", MatchType.FULL)
+            except ImportError:
+                logger.warning(
+                    "llm-guard package is not installed. Using default match type. "
+                    "To enable full functionality, install with: pip install 'netra-sdk[llm_guard]'"
+                )
+
+            from netra.scanner import PromptInjection
+
+            threshold_value = kwargs.get("threshold", 0.5)
+            if not isinstance(threshold_value, (int, float)):
+                logger.info(f"Invalid threshold value: {threshold_value}")
+                threshold = 0.5
+            else:
+                threshold = float(threshold_value)
+
+            return PromptInjection(threshold=threshold, match_type=match_type)
+        else:
+            raise ValueError(f"Unsupported scanner type: {scanner_type}")
+
+    def scan(self, prompt: str, is_blocked: bool = False) -> ScanResult:
+        violations_detected = []
+        for scanner_type in self.scanner_types:
+            try:
+                scanner = self._get_scanner(scanner_type)
+                scanner.scan(prompt)
+            except ValueError as e:
+                raise ValueError(f"Invalid value type: {e}")
+            except InjectionException as error:
+                violations_detected.append(error.violations[0])
+
+        # Create dynamic violation actions mapping based on detected violations and blocking status
+        violations_actions = {}
+        if violations_detected:
+            if is_blocked:
+                violations_actions["BLOCK"] = violations_detected
+            else:
+                violations_actions["FLAG"] = violations_detected
+
+            Netra.set_custom_event(
+                event_name="violation_detected",
+                attributes={
+                    "has_violation": True,
+                    "violations": violations_detected,
+                    "is_blocked": is_blocked,
+                    "violation_actions": json.dumps(violations_actions),
+                },
+            )
+
+        if is_blocked and violations_detected:
+            raise InjectionException(
+                message=f"Input blocked: detected {', '.join(violations_detected)}.",
+                has_violation=True,
+                violations=violations_detected,
+                is_blocked=True,
+                violation_actions=violations_actions,
+            )
+
+        return ScanResult(
+            has_violation=bool(violations_detected),
+            violations=violations_detected,
+            violation_actions=violations_actions,
+            is_blocked=bool(is_blocked and violations_detected),
+        )

netra/instrumentation/__init__.py

@@ -0,0 +1,257 @@
+import logging
+from typing import Any, Callable, Optional, Set
+
+from traceloop.sdk import Instruments, Telemetry
+from traceloop.sdk.utils.package_check import is_package_installed
+
+from netra.instrumentation.instruments import CustomInstruments, NetraInstruments
+
+
+def init_instrumentations(
+    should_enrich_metrics: bool,
+    base64_image_uploader: Optional[Callable[[str, str, str], str]],
+    instruments: Optional[Set[NetraInstruments]] = None,
+    block_instruments: Optional[Set[NetraInstruments]] = None,
+) -> None:
+    from traceloop.sdk.tracing.tracing import init_instrumentations
+
+    traceloop_instruments = set()
+    traceloop_block_instruments = set()
+    netra_custom_instruments = set()
+    netra_custom_block_instruments = set()
+    if instruments is not None:
+        for instrument in instruments:
+            if instrument.origin == CustomInstruments:  # type: ignore[attr-defined]
+                netra_custom_instruments.add(getattr(CustomInstruments, instrument.name))
+            else:
+                traceloop_instruments.add(getattr(Instruments, instrument.name))
+    if block_instruments is not None:
+        for instrument in block_instruments:
+            if instrument.origin == CustomInstruments:  # type: ignore[attr-defined]
+                netra_custom_block_instruments.add(getattr(CustomInstruments, instrument.name))
+            else:
+                traceloop_block_instruments.add(getattr(Instruments, instrument.name))
+    traceloop_block_instruments.update(
+        {
+            Instruments.WEAVIATE,
+            Instruments.QDRANT,
+            Instruments.GOOGLE_GENERATIVEAI,
+            Instruments.MISTRAL,
+        }
+    )
+
+    init_instrumentations(
+        should_enrich_metrics=should_enrich_metrics,
+        base64_image_uploader=base64_image_uploader,
+        instruments=traceloop_instruments,
+        block_instruments=traceloop_block_instruments,
+    )
+
+    netra_custom_instruments = netra_custom_instruments or set(CustomInstruments)
+    netra_custom_instruments = netra_custom_instruments - netra_custom_block_instruments
+    # Initialize Google GenAI instrumentation.
+    if CustomInstruments.GOOGLE_GENERATIVEAI in netra_custom_instruments:
+        init_google_genai_instrumentation()
+
+    # Initialize FastAPI instrumentation.
+    if CustomInstruments.FASTAPI in netra_custom_instruments:
+        init_fastapi_instrumentation()
+
+    # Initialize Qdrant instrumentation.
+    if CustomInstruments.QDRANTDB in netra_custom_instruments:
+        init_qdrant_instrumentation()
+
+    # Initialize Weaviate instrumentation.
+    if CustomInstruments.WEAVIATEDB in netra_custom_instruments:
+        init_weviate_instrumentation()
+
+    # Initialize HTTPX instrumentation.
+    if CustomInstruments.HTTPX in netra_custom_instruments:
+        init_httpx_instrumentation()
+
+    # Initialize AIOHTTP instrumentation.
+    if CustomInstruments.AIOHTTP in netra_custom_instruments:
+        init_aiohttp_instrumentation()
+
+    # Initialize Cohere instrumentation.
+    if CustomInstruments.COHEREAI in netra_custom_instruments:
+        init_cohere_instrumentation()
+
+    if CustomInstruments.MISTRALAI in netra_custom_instruments:
+        init_mistral_instrumentor()
+
+
+def init_google_genai_instrumentation() -> bool:
+    """Initialize Google GenAI instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("google-genai"):
+            Telemetry().capture("instrumentation:genai:init")
+            from netra.instrumentation.google_genai import GoogleGenAiInstrumentor
+
+            instrumentor = GoogleGenAiInstrumentor(
+                exception_logger=lambda e: Telemetry().log_exception(e),
+            )
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing Google GenAI instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_fastapi_instrumentation() -> bool:
+    """Initialize FastAPI instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if not is_package_installed("fastapi"):
+            return True
+        from fastapi import FastAPI
+
+        original_init = FastAPI.__init__
+
+        def _patched_init(self: FastAPI, *args: Any, **kwargs: Any) -> None:
+            original_init(self, *args, **kwargs)
+
+            try:
+                from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+                FastAPIInstrumentor().instrument_app(self)
+            except Exception as e:
+                logging.warning(f"Failed to auto-instrument FastAPI: {e}")
+
+        FastAPI.__init__ = _patched_init
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing FastAPI instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_qdrant_instrumentation() -> bool:
+    """Initialize Qdrant instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("qdrant-client"):
+            from opentelemetry.instrumentation.qdrant import QdrantInstrumentor
+
+            instrumentor = QdrantInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing Qdrant instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_weviate_instrumentation() -> bool:
+    """Initialize Weaviate instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("weaviate-client"):
+            from netra.instrumentation.weaviate import WeaviateInstrumentor
+
+            instrumentor = WeaviateInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing Weaviate instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_httpx_instrumentation() -> bool:
+    """Initialize HTTPX instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("httpx"):
+            from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+
+            instrumentor = HTTPXClientInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing HTTPX instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_aiohttp_instrumentation() -> bool:
+    """Initialize AIOHTTP instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("aiohttp"):
+            from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
+
+            instrumentor = AioHttpClientInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing AIOHTTP instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_cohere_instrumentation() -> bool:
+    """Initialize Cohere instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("cohere"):
+            from netra.instrumentation.cohere import CohereInstrumentor
+
+            instrumentor = CohereInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing Cohere instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
+def init_mistral_instrumentor() -> bool:
+    """Initialize Mistral instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("mistralai"):
+            from netra.instrumentation.mistralai import MistralAiInstrumentor
+
+            instrumentor = MistralAiInstrumentor(
+                exception_logger=lambda e: Telemetry().log_exception(e),
+            )
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"-----Error initializing Mistral instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False