Add WildGuard Guardrail Microservice

.github/workflows/docker/compose/guardrails-compose-cd.yaml

@@ -14,6 +14,10 @@ services:
     build:
       dockerfile: comps/guardrails/toxicity_detection/Dockerfile
     image: ${REGISTRY:-opea}/guardrails-toxicity-detection:${TAG:-latest}
+  guardrails-wildguard:
+    build:
+      dockerfile: comps/guardrails/wildguard/langchain/Dockerfile
+    image: ${REGISTRY:-opea}/guardrails-wildguard:${TAG:-latest}
   guardrails-pii-detection-predictionguard:
     build:
       dockerfile: comps/guardrails/pii_detection/predictionguard/Dockerfile

comps/guardrails/README.md

@@ -4,7 +4,8 @@ The Guardrails service enhances the security of LLM-based applications by offeri
 
 | MicroService                                         | Description                                                                                                              |
 | ---------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------ |
-| [Llama Guard](./llama_guard/langchain/README.md)     | Provides guardrails for inputs and outputs to ensure safe interactions                                                   |
+| [Llama Guard](./llama_guard/langchain/README.md)     | Provides guardrails for inputs and outputs to ensure safe interactions using Llama Guard                                 |
+| [WildGuard](./wildguard/langchain/README.md)         | Provides guardrails for inputs and outputs to ensure safe interactions using WildGuard                                   |
 | [PII Detection](./pii_detection/README.md)           | Detects Personally Identifiable Information (PII) and Business Sensitive Information (BSI)                               |
 | [Toxicity Detection](./toxicity_detection/README.md) | Detects Toxic language (rude, disrespectful, or unreasonable language that is likely to make someone leave a discussion) |
 | [Bias Detection](./bias_detection/README.md)         | Detects Biased language (framing bias, epistemological bias, and demographic bias)                                       |

comps/guardrails/wildguard/langchain/Dockerfile

@@ -0,0 +1,30 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+FROM python:3.11-slim
+
+ENV LANG=C.UTF-8
+
+ARG ARCH="cpu"
+
+RUN apt-get update -y && apt-get install -y --no-install-recommends --fix-missing \
+    libgl1-mesa-glx \
+    libjemalloc-dev
+
+RUN useradd -m -s /bin/bash user && \
+    mkdir -p /home/user && \
+    chown -R user /home/user/
+
+USER user
+
+COPY comps /home/user/comps
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    if [ ${ARCH} = "cpu" ]; then pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu; fi && \
+    pip install --no-cache-dir -r /home/user/comps/guardrails/wildguard/langchain/requirements.txt
+
+ENV PYTHONPATH=$PYTHONPATH:/home/user
+
+WORKDIR /home/user/comps/guardrails/wildguard/langchain/
+
+ENTRYPOINT ["python", "guardrails_tgi.py"]

comps/guardrails/wildguard/langchain/README.md

@@ -0,0 +1,107 @@
+# Guardrails Microservice
+
+To fortify AI initiatives in production, this microservice introduces guardrails designed to encapsulate LLMs, ensuring the enforcement of responsible behavior. With this microservice, you can secure model inputs and outputs, hastening your journey to production and democratizing AI within your organization, building Trustworthy, Safe, and Secure LLM-based Applications.
+
+These guardrails actively prevent the model from interacting with unsafe content, promptly signaling its inability to assist with such requests. With these protective measures in place, you can expedite production timelines and alleviate concerns about unpredictable model responses.
+
+The Guardrails Microservice now offers two primary types of guardrails:
+
+- Input Guardrails: These are applied to user inputs. An input guardrail can reject the input, halting further processing.
+- Output Guardrails: These are applied to outputs generated by the LLM. An output guardrail can reject the output, preventing it from being returned to the user.
+
+We offer content moderation support utilizing Allen Institute for AI's [WildGuard](https://huggingface.co/allenai/wildguard) model.
+
+Any content that is detected in the following categories is determined as unsafe:
+
+- Privacy
+- Misinformation
+- Harmful Language
+- Malicious Uses
+
+## 🚀1. Start Microservice with Python (Option 1)
+
+To start the Guardrails microservice, you need to install python packages first.
+
+### 1.1 Install Requirements
+
+```bash
+pip install -r requirements.txt
+```
+
+### 1.2 Start TGI Gaudi Service
+
+```bash
+export HF_TOKEN=${your_hf_api_token}
+volume=$PWD/data
+model_id="allenai/wildguard"
+docker pull ghcr.io/huggingface/tgi-gaudi:2.0.1
+docker run -p 8088:80 -v $volume:/data --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host -e HTTPS_PROXY=$https_proxy -e HTTP_PROXY=$https_proxy -e HF_TOKEN=$HF_TOKEN ghcr.io/huggingface/tgi-gaudi:2.0.1 --model-id $model_id --max-input-length 1024 --max-total-tokens 2048
+```
+
+### 1.3 Verify the TGI Gaudi Service
+
+```bash
+curl 127.0.0.1:8088/generate \
+  -X POST \
+  -d '{"inputs":"How do you buy a tiger in the US?","parameters":{"max_new_tokens":32}}' \
+  -H 'Content-Type: application/json'
+```
+
+### 1.4 Start Guardrails Service
+
+```bash
+export SAFETY_GUARD_ENDPOINT="http://${your_ip}:8088"
+python guardrails_tgi.py
+```
+
+## 🚀2. Start Microservice with Docker (Option 2)
+
+If you start an Guardrails microservice with docker, the `docker_compose_guardrails.yaml` file will automatically start a TGI gaudi service with docker.
+
+### 2.1 Setup Environment Variables
+
+In order to start TGI and LLM services, you need to setup the following environment variables first.
+
+```bash
+export HUGGINGFACEHUB_API_TOKEN=${your_hf_api_token}
+export SAFETY_GUARD_ENDPOINT="http://${your_ip}:8088"
+export LLM_MODEL_ID=${your_hf_llm_model}
+```
+
+### 2.2 Build Docker Image
+
+```bash
+cd ../../../../
+docker build -t opea/guardrails-tgi:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/guardrails/wildguard/langchain/Dockerfile .
+```
+
+### 2.3 Run Docker with CLI
+
+```bash
+docker run -d --name="guardrails-tgi-server" -p 9090:9090 --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e SAFETY_GUARD_ENDPOINT=$SAFETY_GUARD_ENDPOINT -e HUGGINGFACEHUB_API_TOKEN=$HUGGINGFACEHUB_API_TOKEN opea/guardrails-tgi:latest
+```
+
+### 2.4 Run Docker with Docker Compose
+
+```bash
+docker compose -f docker_compose_guardrails.yaml up -d
+```
+
+## 🚀3. Consume Guardrails Service
+
+### 3.1 Check Service Status
+
+```bash
+curl http://localhost:9090/v1/health_check\
+  -X GET \
+  -H 'Content-Type: application/json'
+```
+
+### 3.2 Consume Guardrails Service
+
+```bash
+curl http://localhost:9090/v1/guardrails\
+  -X POST \
+  -d '{"text":"How do you buy a tiger in the US?","parameters":{"max_new_tokens":32}}' \
+  -H 'Content-Type: application/json'
+```

comps/guardrails/wildguard/langchain/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0

comps/guardrails/wildguard/langchain/docker_compose_guardrails.yaml

@@ -0,0 +1,35 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+version: "3.8"
+
+services:
+  tgi_gaudi_service:
+    image: ghcr.io/huggingface/tgi-gaudi:2.0.1
+    container_name: tgi-service
+    ports:
+      - "8088:80"
+    volumes:
+      - "./data:/data"
+    environment:
+      HF_TOKEN: ${HF_TOKEN}
+    shm_size: 1g
+    command: --model-id ${LLM_MODEL_ID} --max-input-tokens 1024 --max-total-tokens 2048
+  guardrails:
+    image: opea/guardrails-tgi:latest
+    container_name: guardrails-tgi-gaudi-server
+    ports:
+      - "9090:9090"
+    ipc: host
+    environment:
+      no_proxy: ${no_proxy}
+      http_proxy: ${http_proxy}
+      https_proxy: ${https_proxy}
+      SAFETY_GUARD_ENDPOINT: ${SAFETY_GUARD_ENDPOINT}
+      LANGCHAIN_API_KEY: ${LANGCHAIN_API_KEY}
+      HUGGINGFACEHUB_API_TOKEN: ${HF_TOKEN}
+    restart: unless-stopped
+
+networks:
+  default:
+    driver: bridge

comps/guardrails/wildguard/langchain/guardrails_tgi.py

@@ -0,0 +1,76 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import List, Union
+
+from langchain_huggingface.llms import HuggingFaceEndpoint
+
+from comps import CustomLogger, GeneratedDoc, ServiceType, TextDoc, opea_microservices, register_microservice
+
+logger = CustomLogger("guardrails_tgi")
+logflag = os.getenv("LOGFLAG", False)
+
+DEFAULT_MODEL = "allenai/wildguard"
+INSTRUCTION_FORMAT = """<s><|user|>
+[INST] You are a classifier for analyzing interactions between humans and LLMs. Given a request from a human user and a response from an LLM assistant, determine 1) whether the human user's request is harmful, 2) whether the AI assistant's response is a refusal, and 3) whether the AI assistant's response is harmful.
+
+Human user:
+{prompt}
+
+AI assistant:
+{response}
+
+---
+
+Answers: [/INST]
+<|assistant|>
+"""
+
+
+@register_microservice(
+    name="opea_service@wildguard_tgi",
+    service_type=ServiceType.GUARDRAIL,
+    endpoint="/v1/guardrails",
+    host="0.0.0.0",
+    port=9090,
+    input_datatype=Union[GeneratedDoc, TextDoc],
+    output_datatype=TextDoc,
+)
+def safety_guard(input: Union[GeneratedDoc, TextDoc]) -> TextDoc:
+    if logflag:
+        logger.info(input)
+
+    if isinstance(input, GeneratedDoc):
+        messages = INSTRUCTION_FORMAT.format(prompt=input.prompt, response=input.text)
+    else:
+        messages = INSTRUCTION_FORMAT.format(prompt=input.text, response="")
+    response_input_guard = llm_guard.invoke(messages)
+
+    if "Harmful request: yes" in response_input_guard or "Harmful response: yes" in response_input_guard:
+
+        if logflag:
+            logger.info("Violated policies: harmful")
+        res = TextDoc(text="Violated policies: harmful, please check your input.", downstream_black_list=[".*"])
+    else:
+        res = TextDoc(text=input.text)
+    if logflag:
+        logger.info(res)
+    return res
+
+
+if __name__ == "__main__":
+    safety_guard_endpoint = os.getenv("SAFETY_GUARD_ENDPOINT", "http://localhost:8080")
+    safety_guard_model = os.getenv("SAFETY_GUARD_MODEL_ID", DEFAULT_MODEL)
+    llm_guard = HuggingFaceEndpoint(
+        endpoint_url=safety_guard_endpoint,
+        max_new_tokens=100,
+        top_k=1,
+        top_p=0.95,
+        typical_p=0.95,
+        temperature=0.01,
+        repetition_penalty=1.03,
+    )
+    # chat engine for server-side prompt templating
+    logger.info("guardrails - router] LLM initialized.")
+    opea_microservices["opea_service@wildguard_tgi"].start()

comps/guardrails/wildguard/langchain/requirements.txt

@@ -0,0 +1,11 @@
+docarray[full]
+fastapi
+huggingface-hub
+langchain-community
+langchain-huggingface
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-sdk
+prometheus-fastapi-instrumentator
+shortuuid
+uvicorn

tests/guardrails/test_guardrails_wild_guard_langchain_on_intel_hpu.sh

@@ -0,0 +1,82 @@
+#!/bin/bash
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+set -x
+
+WORKPATH=$(dirname "$PWD")
+ip_address=$(hostname -I | awk '{print $1}')
+
+function build_docker_images() {
+    echo "Start building docker images for microservice"
+    cd $WORKPATH
+    docker pull ghcr.io/huggingface/tgi-gaudi:2.0.1
+    docker build --no-cache -t opea/guardrails-tgi:comps --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/guardrails/wildguard/langchain/Dockerfile .
+    if [ $? -ne 0 ]; then
+        echo "opea/guardrails-tgi built fail"
+        exit 1
+    else
+        echo "opea/guardrails-tgi built successful"
+    fi
+}
+
+function start_service() {
+    echo "Starting microservice"
+    export model_id="allenai/wildguard"
+    export SAFETY_GUARD_MODEL_ID="allenai/wildguard"
+    export SAFETY_GUARD_ENDPOINT=http://${ip_address}:5035
+
+    docker run -d --name="test-comps-guardrails-langchain-tgi-server" -p 5035:80 --runtime=habana -e HF_TOKEN=$HF_TOKEN -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host -e HTTPS_PROXY=$https_proxy -e HTTP_PROXY=$https_proxy ghcr.io/huggingface/tgi-gaudi:2.0.1 --model-id $model_id --max-input-length 1024 --max-total-tokens 2048
+    sleep 4m
+    docker run -d --name="test-comps-guardrails-langchain-service" -p 5036:9090 --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e SAFETY_GUARD_MODEL_ID=$SAFETY_GUARD_MODEL_ID -e SAFETY_GUARD_ENDPOINT=$SAFETY_GUARD_ENDPOINT -e HUGGINGFACEHUB_API_TOKEN=${HUGGINGFACEHUB_API_TOKEN} -e HF_TOKEN=${HUGGINGFACEHUB_API_TOKEN} opea/guardrails-tgi:comps
+    sleep 10s
+    echo "Microservice started"
+}
+
+function validate_microservice() {
+    echo "Validate microservice started"
+    echo "test 1 - violated policies"
+    result=$(http_proxy= curl http://localhost:5036/v1/guardrails  -X POST   -d '{"text":"How do you buy a tiger in the US?","parameters":{"max_new_tokens":32}}'   -H 'Content-Type: application/json')
+    if [[ $result == *"Violated"* ]]; then
+        echo "Result correct."
+    else
+        echo "Result wrong."
+        docker logs test-comps-guardrails-langchain-tgi-server
+        docker logs test-comps-guardrails-langchain-service
+        exit 1
+    fi
+    echo "test 2 - safe"
+    result=$(http_proxy= curl http://localhost:5036/v1/guardrails  -X POST   -d '{"text":"How do you buy a car in the US?","parameters":{"max_new_tokens":32}}'   -H 'Content-Type: application/json')
+        if [[ $result == *"car"* ]]; then
+        echo "Result correct."
+    else
+        echo "Result wrong."
+        docker logs test-comps-guardrails-langchain-tgi-server
+        docker logs test-comps-guardrails-langchain-service
+        exit 1
+    fi
+
+}
+
+function stop_docker() {
+    cid=$(docker ps -aq --filter "name=test-comps-guardrails-langchain*")
+    echo "Shutdown legacy containers "$cid
+    if [[ ! -z "$cid" ]]; then docker stop $cid && docker rm $cid && sleep 1s; fi
+}
+
+function main() {
+
+    stop_docker
+
+    build_docker_images
+    start_service
+
+    validate_microservice
+
+    stop_docker
+    echo "cleanup container images and volumes"
+    echo y | docker system prune 2>&1 > /dev/null
+
+}
+
+main