Merge pull request #57 from RobGeada/MetricsAndNonBlockingGuardrails

Feat: Custom user metrics and non-blocking guardrails

Author: RobGeada

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/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/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()}")
 
 

detectors/common/app.py

@@ -7,7 +7,7 @@
 import yaml
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
-from prometheus_client import Counter
+from prometheus_client import Counter, CollectorRegistry
 
 import logging
 
@@ -28,29 +28,30 @@
     dependencies=[],
 )
 
+METRIC_PREFIX = "trustyai_guardrails"
 
 class DetectorBaseAPI(FastAPI):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.state.detectors = {}
         self.state.instruments = {
             "detections": Counter(
-                "trustyai_guardrails_detections",
+                f"{METRIC_PREFIX}_detections",
                 "Number of detections per detector function",
                 ["detector_kind", "detector_name"]
             ),
             "requests": Counter(
-                "trustyai_guardrails_requests",
+                f"{METRIC_PREFIX}_requests",
                 "Number of requests per detector function",
                 ["detector_kind", "detector_name"]
             ),
             "errors": Counter(
-                "trustyai_guardrails_errors",
+                f"{METRIC_PREFIX}_errors",
                 "Number of errors per detector function",
                 ["detector_kind", "detector_name"]
             ),
             "runtime": Counter(
-                "trustyai_guardrails_runtime",
+                f"{METRIC_PREFIX}_runtime",
                 "Total runtime of a detector function- this is the induced latency of this guardrail",
                 ["detector_kind", "detector_name"]
             )
@@ -62,7 +63,6 @@ def __init__(self, *args, **kwargs):
         self.add_exception_handler(StarletteHTTPException, self.http_exception_handler)
         self.add_api_route("/health", health, description="Check if server is alive")
 
-
     async def validation_exception_handler(self, request, exc):
         errors = exc.errors()
         if len(errors) > 0 and errors[0]["type"] == "missing":

detectors/common/instrumented_detector.py

@@ -17,8 +17,11 @@ def instrument_runtime(self, function_name: str):
         finally:
             pass
 
-    def add_instruments(self, gauges):
-        self.instruments = gauges
+    def set_instruments(self, instruments):
+        self.instruments = instruments
+
+    def add_instrument(self, instrument):
+        self.instruments[instrument._name] = instrument
 
     def increment_detector_instruments(self, function_name: str, is_detection: bool):
         """Increment the detection and request counters, automatically update rates"""

detectors/huggingface/app.py

@@ -16,7 +16,7 @@
 async def lifespan(app: FastAPI):
     detector = Detector()
     app.set_detector(detector, detector.model_name)
-    detector.add_instruments(app.state.instruments)
+    detector.set_instruments(app.state.instruments)
     yield
     # Clean up the ML models and release the resources
     detector: Detector = app.get_detector()

docs/custom_detectors.md

@@ -0,0 +1,41 @@
+# Custom Detectors
+You can overwrite [custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) to create
+custom detectors in the built-in server based off of arbitrary Python code. This lets you quickly and flexibly
+create your own detection logic!
+
+The following rules apply:
+1) Each function defined in the `custom_detectors.py` 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
+   * see the `function_that_needs_headers` example in [custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) for usage
+3) Functions that are intended to be used as detectors must either return a `bool` or a `dict`: 
+   1) Return values that evaluate to false (e.g., `{}`, `""`, `None`, etc) are treated as non-detections 
+   2) Boolean responses of `true` are considered a detection 
+      * see the `over_100_characters` example in [custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) for usage
+   3) Dict response that are parseable as a `ContentAnalysisResponse` object are considered a detection
+      * see the `contains_word` example in [custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) for usage
+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
+
+
+## Utility Decorators
+The following decorators are also available, and are automatically imported into the custom_detectors.py file:
+
+### `@use_instruments(instruments=[$INSTRUMENT_1, ..., $INSTRUMENT_N])`
+ Use this decorator to register your own Prometheus instruments with the server's main
+ `/metrics` registry. See the `function_that_has_prometheus_metrics` example 
+ in [custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) for usage.
+
+### `@non_blocking(return_value=$RETURN_VALUE)`
+Use this decorator to indicate that the logic inside this function should run in a non-blocking
+background thread. The guardrail function will immediately return $RETURN_VALUE while launching 
+your function logic into a background thread.
+
+This enables a number of use-cases, such as:
+* Producing some background analysis metric over the input/output, without adding latency to the system
+* Performing "silent" guardrailing, e.g., adding information to a server or alerting admins
+
+See the `background_function` example in
+[custom_detectors.py](detectors/built_in/custom_detectors/custom_detectors.py) for usage.
+
+## More Examples
+For a "real-world" example, check out the [TrustyAI custom detectors demo](https://github.com/trustyai-explainability/trustyai-llm-demo/blob/main/custom-detectors/custom_detectors.py)!

tests/conftest.py

@@ -1,6 +1,8 @@
 import os
+import shutil
 import sys
 import pytest
+import tempfile
 
 
 @pytest.fixture(autouse=True)
@@ -24,3 +26,22 @@ def setup_imports():
             sys.path.insert(0, path)
             print(f"Added to sys.path: {path}")
 
+@pytest.fixture(scope="session", autouse=True)
+def prometheus_multiproc_dir():
+    """
+    Create a temporary directory for PROMETHEUS_MULTIPROC_DIR and set the environment variable.
+    """
+    tmpdir = tempfile.mkdtemp(prefix="prometheus_multiproc_")
+    os.environ["PROMETHEUS_MULTIPROC_DIR"] = tmpdir
+    yield tmpdir
+    # Cleanup will be handled by the next fixture
+
+@pytest.fixture(scope="session", autouse=True)
+def cleanup_prometheus_multiproc_dir(request, prometheus_multiproc_dir):
+    """
+    Cleanup the PROMETHEUS_MULTIPROC_DIR after the test session.
+    """
+    yield
+    shutil.rmtree(prometheus_multiproc_dir, ignore_errors=True)
+    if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+        del os.environ["PROMETHEUS_MULTIPROC_DIR"]