Skip to content

Commit

Permalink
Add WildGuard Guardrail Microservice (#710)
Browse files Browse the repository at this point in the history
* add wildguard microservice

Signed-off-by: Daniel Deleon <[email protected]>

* fix paths for wildguard dir

Signed-off-by: Daniel Deleon <[email protected]>

* add README

Signed-off-by: Daniel Deleon <[email protected]>

* add wildguard to guardrail README table

Signed-off-by: Daniel Deleon <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* updates per code review

Signed-off-by: Daniel Deleon <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update tokens

Signed-off-by: Daniel Deleon <[email protected]>

* fix endpoint for testing

Signed-off-by: Daniel Deleon <[email protected]>

* fixed formating

Signed-off-by: Daniel Deleon <[email protected]>

* add more description of wildguard to distinguish from llamaguard

Signed-off-by: Daniel Deleon <[email protected]>

---------

Signed-off-by: Daniel Deleon <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Letong Han <[email protected]>
Co-authored-by: Abolfazl Shahbazi <[email protected]>
Co-authored-by: ZePan110 <[email protected]>
  • Loading branch information
5 people authored Oct 11, 2024
1 parent 367b3aa commit 5bb4046
Show file tree
Hide file tree
Showing 9 changed files with 348 additions and 1 deletion.
4 changes: 4 additions & 0 deletions .github/workflows/docker/compose/guardrails-compose-cd.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion comps/guardrails/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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) |
Expand Down
30 changes: 30 additions & 0 deletions comps/guardrails/wildguard/langchain/Dockerfile
Original file line number Diff line number Diff line change
@@ -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"]
107 changes: 107 additions & 0 deletions comps/guardrails/wildguard/langchain/README.md
Original file line number Diff line number Diff line change
@@ -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.

`allenai/wildguard` was fine-tuned from `mistralai/Mistral-7B-v0.3` on their own [`allenai/wildguardmix`](https://huggingface.co/datasets/allenai/wildguardmix) dataset. 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'
```
2 changes: 2 additions & 0 deletions comps/guardrails/wildguard/langchain/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
Original file line number Diff line number Diff line change
@@ -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
76 changes: 76 additions & 0 deletions comps/guardrails/wildguard/langchain/guardrails_tgi.py
Original file line number Diff line number Diff line change
@@ -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()
11 changes: 11 additions & 0 deletions comps/guardrails/wildguard/langchain/requirements.txt
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
#!/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

0 comments on commit 5bb4046

Please sign in to comment.