Merge branch 'trustyai-explainability:main' into incubation

Author: ruivieira

README.md

@@ -13,7 +13,7 @@ At the moment, the following detectors are supported:
 - `builtIn` -- Small, lightweight detection functions that are deployed out-of-the-box alongside the [Guardrails Orchestrator](https://github.com/foundation-model-stack/fms-guardrails-orchestrator). The built-in detectors provide a number of heuristic or algorithmic detection functions, such as:
   - Regex-based detections, with pre-written regexes for flagging various Personally Identifiable Information items like emails or phone numbers, as well as the ability to provide custom regexes
   - File-type validations, for verifying if model input/output is valid JSON, XML, or YAML
-
+  - "Custom" detectors which can be defined with raw Python code, which allows for easy declaration of custom detection functions that can perform arbitrarily complex logic.
 
 ## Building
 

detectors/Dockerfile.builtIn

@@ -21,12 +21,14 @@ ARG CACHEBUST=1
 RUN echo "$CACHEBUST"
 COPY ./common /app/detectors/common
 COPY ./built_in/ /app
+ENV PROMETHEUS_MULTIPROC_DIR="/tmp/prometheus_multiproc_dir"
+RUN mkdir -p $PROMETHEUS_MULTIPROC_DIR && chmod 777 $PROMETHEUS_MULTIPROC_DIR
 
 EXPOSE 8080
 
 # for backwards compatibility with existing k8s deployment configs
 RUN mkdir /app/bin &&\
-     echo '#!/bin/bash' > /app/bin/regex-detector &&\
+    echo '#!/bin/bash' > /app/bin/regex-detector &&\
     echo "uvicorn app:app --workers 4 --host 0.0.0.0 --port 8080 --log-config /app/detectors/common/log_conf.yaml" >> /app/bin/regex-detector &&\
     chmod +x /app/bin/regex-detector
 CMD ["/app/bin/regex-detector"]

detectors/Dockerfile.hf

@@ -96,11 +96,12 @@ FROM builder
 WORKDIR /app
 ARG CACHEBUST=1
 RUN echo "$CACHEBUST"
-COPY ./common /common
+COPY ./common /app/detectors/common
+COPY ./huggingface/detector.py /app/detectors/huggingface/
+RUN mkdir /common; cp /app/detectors/common/log_conf.yaml /common/
 COPY ./huggingface/app.py /app
-COPY ./huggingface/detector.py /app
 
 EXPOSE 8000
-CMD ["uvicorn", "app:app", "--workers", "4", "--host", "0.0.0.0", "--port", "8000", "--log-config", "/common/log_conf.yaml"]
+CMD ["uvicorn", "app:app", "--workers", "1", "--host", "0.0.0.0", "--port", "8000", "--log-config", "/common/log_conf.yaml"]
 
 # gunicorn main:app --workers 4 --worker-class uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000

detectors/built_in/app.py

@@ -7,8 +7,8 @@
 from custom_detectors_wrapper import CustomDetectorRegistry
 from file_type_detectors import FileTypeDetectorRegistry
 
-from prometheus_fastapi_instrumentator import Instrumentator
-from prometheus_client import Gauge
+from prometheus_client import generate_latest, CONTENT_TYPE_LATEST, CollectorRegistry, multiprocess
+from starlette.responses import Response
 from detectors.common.scheme import ContentAnalysisHttpRequest,  ContentsAnalysisResponse
 from detectors.common.app import DetectorBaseAPI as FastAPI
 
@@ -21,17 +21,23 @@ async def lifespan(app: FastAPI):
         CustomDetectorRegistry()
     ]:
         app.set_detector(detector_registry, detector_registry.registry_name)
-        detector_registry.add_instruments(app.state.instruments)
+        detector_registry.set_instruments(app.state.instruments)
     yield
     app.cleanup_detector()
 
 
 app = FastAPI(lifespan=lifespan)
-Instrumentator().instrument(app).expose(app)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
+@app.get("/metrics")
+def metrics():
+    registry = CollectorRegistry()
+    multiprocess.MultiProcessCollector(registry)
+    data = generate_latest(registry)
+    return Response(data, media_type=CONTENT_TYPE_LATEST)
+
 @app.post("/api/v1/text/contents", response_model=ContentsAnalysisResponse)
 def detect_content(request: ContentAnalysisHttpRequest, raw_request: Request):
     logger.info(f"Request for {request.detector_params}")

detectors/built_in/base_detector_registry.py

@@ -1,19 +1,16 @@
-import contextlib
 import logging
 from abc import ABC, abstractmethod
-import time
-from http.client import HTTPException
+from fastapi import HTTPException
 from typing import List
 
+from detectors.common.instrumented_detector import InstrumentedDetector
 from detectors.common.scheme import ContentAnalysisResponse
 
-class BaseDetectorRegistry(ABC):
+class BaseDetectorRegistry(InstrumentedDetector, ABC):
     def __init__(self, registry_name):
+        super().__init__(registry_name)
         self.registry = None
-        self.registry_name = registry_name
 
-        # prometheus
-        self.instruments = {}
 
     @abstractmethod
     def handle_request(self, content: str, detector_params: dict, headers: dict, **kwargs) -> List[ContentAnalysisResponse]:
@@ -22,35 +19,6 @@ def handle_request(self, content: str, detector_params: dict, headers: dict, **k
     def get_registry(self):
         return self.registry
 
-    def add_instruments(self, gauges):
-        self.instruments = gauges
-
-    def increment_detector_instruments(self, function_name: str, is_detection: bool):
-        """Increment the detection and request counters, automatically update rates"""
-        if self.instruments.get("requests"):
-            self.instruments["requests"].labels(self.registry_name, function_name).inc()
-
-        # The labels() function will initialize the counters if not already created.
-        # This prevents the counters not existing until they are first incremented
-        # If the counters have already been created, this is just a cheap dict.get() call
-        if self.instruments.get("errors"):
-            _ = self.instruments["errors"].labels(self.registry_name, function_name)
-        if self.instruments.get("runtime"):
-            _ = self.instruments["runtime"].labels(self.registry_name, function_name)
-
-        # create and/or increment the detection counter
-        if self.instruments.get("detections"):
-            detection_counter = self.instruments["detections"].labels(self.registry_name, function_name)
-            if is_detection:
-                detection_counter.inc()
-
-
-    def increment_error_instruments(self, function_name: str):
-        """Increment the error counter, update the rate gauges"""
-        if self.instruments.get("errors"):
-            self.instruments["errors"].labels(self.registry_name, function_name).inc()
-
-
     def throw_internal_detector_error(self, function_name: str, logger: logging.Logger, exception: Exception, increment_requests: bool):
         """consistent handling of internal errors within a detection function"""
         if increment_requests and self.instruments.get("requests"):
@@ -60,16 +28,6 @@ def throw_internal_detector_error(self, function_name: str, logger: logging.Logg
         raise HTTPException(status_code=500, detail="Detection error, check detector logs")
 
 
-    @contextlib.contextmanager
-    def instrument_runtime(self, function_name: str):
-        try:
-            start_time = time.time()
-            yield
-            if self.instruments.get("runtime"):
-                self.instruments["runtime"].labels(self.registry_name, function_name).inc(time.time() - start_time)
-        finally:
-            pass
-
     def get_detection_functions_from_params(self, params: dict):
         """Parse the request parameters to extract and normalize detection functions as iterable list"""
         if self.registry_name in params and isinstance(params[self.registry_name], (list, str)):

detectors/built_in/custom_detectors/custom_detectors.py

@@ -1,41 +1,36 @@
 """
-This is an example custom_detectors.py file. Here, you can define any arbitrary Python code as a
-Guardrail detector.
+This is an example custom_detectors.py file. Overwrite this file to define custom guardrailing
+logic!
 
-The following rules apply:
-1) Each function defined in this file (except for those starting with "_") will be registered as a detector
-2) Functions that accept a parameter "headers" will receive the inbound request headers as a parameter
-3) Functions may either return a boolean or a dict:
-    3a) Return values that evaluate to false (e.g., {}, "", None, etc) are treated as non-detections
-    3b) Boolean responses of "true" are considered a detection
-    3c) Dict response must be parseable as a ContentAnalysisResponse object (see example below)
-4) This code may not import "os", "subprocess", "sys", or "shutil" for security reasons
-5) This code may not call "eval", "exec", "open", "compile", or "input" for security reasons
+See [docs/custom_detectors.md](../../docs/custom_detectors.md) for more details.
 """
 
-# example boolean-returning function
-def over_100_characters(text: str) -> bool:
-    return len(text)>100
-
-# example dict-returning function
-def contains_word(text: str) -> dict:
-    detection = "apple" in text.lower()
-    if detection:
-        detection_position = text.find("apple")
-        return {
-            "start":detection_position,  # start position of detection in text
-            "end": detection_position+5, # end position of detection in text
-            "text": text, # "the flagged text, or some arbitrary message to return to the user"
-            "detection_type": "content_check", #detection_type -> use these fields to define your detector taxonomy as you see fit
-            "detection": "forbidden_word: apple", ##detection -> use these fields to define your detector taxonomy as you see fit
-            "score": 1.0 # the score/severity/probability of the detection
-        }
-    else:
-        return {}
-
-def _this_function_will_not_be_exposed():
-    pass
-
-def function_that_needs_headers(text: str, headers: dict) -> bool:
-    return headers['magic-key'] != "123"
+import time
+def slow_func(text: str) -> bool:
+    time.sleep(.25)
+    return False
+    
+from prometheus_client import Counter
 
+prompt_rejection_counter = Counter(
+    "trustyai_guardrails_system_prompt_rejections",
+    "Number of rejections by the system prompt",
+)
+  
+@use_instruments(instruments=[prompt_rejection_counter])  
+def has_metrics(text: str) -> bool:
+    if "sorry" in text:
+        prompt_rejection_counter.inc()
+    return False
+    
+background_metric = Counter(
+    "trustyai_guardrails_background_metric",
+    "Runs some logic in the background without blocking the /detections call"
+)
+@use_instruments(instruments=[background_metric])
+@non_blocking(return_value=False)
+def background_function(text: str) -> bool:
+    time.sleep(.25)
+    if "sorry" in text:
+        background_metric.inc()     
+    return False

detectors/built_in/custom_detectors_wrapper.py

@@ -1,18 +1,71 @@
 import ast
+import logging
+import importlib.util
+import inspect
+import functools
 import os
-import traceback
+import sys
 
+from concurrent.futures import ThreadPoolExecutor
 from fastapi import HTTPException
-import inspect
-import logging
 from typing import List, Optional, Callable
 
-
 from base_detector_registry import BaseDetectorRegistry
+from detectors.common.app import METRIC_PREFIX
 from detectors.common.scheme import ContentAnalysisResponse
 
 logger = logging.getLogger(__name__)
 
+def use_instruments(instruments: List):
+    """Use this decorator to register the provided Prometheus instruments with the main /metrics endpoint"""
+    def inner_layer_1(func):
+        @functools.wraps(func)
+        def inner_layer_2(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        # check to see if "func" is already decorated, and only add the prometheus instruments field into the original function
+        target = get_underlying_function(func)
+        setattr(target, "prometheus_instruments", instruments)
+        return inner_layer_2
+    return inner_layer_1
+
+def non_blocking(return_value):
+    """
+    Use this decorator to run the guardrail as a non-blocking background thread.
+
+    The `return_value` is returned instantly to the caller of the /api/v1/text/contents, while
+    the logic inside the function will run asynchronously in the background.
+    """
+    def inner_layer_1(func):
+        @functools.wraps(func)
+        def inner_layer_2(*args, **kwargs):
+            executor = getattr(non_blocking, "_executor", None)
+            if executor is None:
+                executor = ThreadPoolExecutor()
+                non_blocking._executor = executor
+            def runner():
+                try:
+                    func(*args, **kwargs)
+                except Exception as e:
+                    logging.error(f"Exception in non-blocking guardrail {func.__name__}: {e}")
+            executor.submit(runner)
+
+            # check to see if "func" is already decorated by `use_instruments`, and grab the instruments if so
+            target = get_underlying_function(func)
+            if hasattr(target, "prometheus_instruments"):
+                setattr(target, "prometheus_instruments", target.prometheus_instruments)
+            return return_value
+        return inner_layer_2
+    return inner_layer_1
+
+forbidden_names = [use_instruments.__name__, non_blocking.__name__]
+
+def get_underlying_function(func):
+    if hasattr(func, "__wrapped__"):
+        return get_underlying_function(func.__wrapped__)
+    return func
+
+
 def custom_func_wrapper(func: Callable, func_name: str, s: str, headers: dict) -> Optional[ContentAnalysisResponse]:
     """Convert a some f(text)->bool into a Detector response"""
     sig = inspect.signature(func)
@@ -92,17 +145,42 @@ class CustomDetectorRegistry(BaseDetectorRegistry):
     def __init__(self):
         super().__init__("custom")
 
+        # check the imported code for potential security issues
         issues = static_code_analysis(module_path = os.path.join(os.path.dirname(__file__), "custom_detectors", "custom_detectors.py"))
         if issues:
             logging.error(f"Detected {len(issues)} potential security issues inside the custom_detectors file: {issues}")
             raise ImportError(f"Unsafe code detected in custom_detectors:\n" + "\n".join(issues))
 
-        import custom_detectors.custom_detectors as custom_detectors
+        # grab custom detectors module
+        module_path = os.path.join(os.path.dirname(__file__), "custom_detectors", "custom_detectors.py")
+        spec = importlib.util.spec_from_file_location("custom_detectors.custom_detectors", module_path)
+        custom_detectors = importlib.util.module_from_spec(spec)
+
+        # inject any user utility functions into the code automatically
+        inject_imports = {
+            "use_instruments": use_instruments,
+            "non_blocking": non_blocking,
+        }
+        for name, mod in inject_imports.items():
+            setattr(custom_detectors, name, mod)
+
+        # load the module
+        sys.modules["custom_detectors.custom_detectors"] = custom_detectors
+        spec.loader.exec_module(custom_detectors)
 
         self.registry = {name: obj for name, obj
                          in inspect.getmembers(custom_detectors, inspect.isfunction)
-                         if not name.startswith("_")}
+                         if not name.startswith("_") and name not in forbidden_names}
         self.function_needs_headers = {name: "headers" in inspect.signature(obj).parameters for name, obj in self.registry.items() }
+
+        # check if functions have requested user prometheus metrics
+        for name, func in self.registry.items():
+            target = get_underlying_function(func)
+            if getattr(target, "prometheus_instruments", False):
+                instruments = target.prometheus_instruments
+                for instrument in instruments:
+                    super().add_instrument(instrument)
+
         logger.info(f"Registered the following custom detectors: {self.registry.keys()}")