Python guardrails

GuardrailAppEnvironment turns a Python function into a guardrail app. You write a function that takes the request and returns a verdict. The class serves the HTTP endpoints the gateway calls, and gives the app the Guardrail type the console looks for.

Install

pip install 'flyteplugins-union[guardrails]>=0.19.1'

Requires flyteplugins-union 0.19.1 or later. The guardrails extra adds FastAPI and uvicorn.

Write a pre-call guardrail

This guardrail refuses prompts that contain a Social Security number or an email address, and answers one known question from a cache:

pii_guardrail.py
                
SSN = re.compile(r"\b\d{3}-\d{2}-\d{4}\b")
EMAIL = re.compile(r"\b[\w.%+-]+@[\w.-]+\.[A-Za-z]{2,}\b")

# Questions answered without calling a model.
CACHE = {
    "what is our refund window?": "Refunds are accepted within 30 days of purchase.",
}

def check_prompt(request: GuardrailRequest) -> Verdict:
    prompt = request.last_user_message()
    if SSN.search(prompt):
        return Verdict.block("Remove the Social Security number and try again.", reason="pii_ssn")
    if EMAIL.search(prompt):
        return Verdict.block("Remove the email address and try again.", reason="pii_email")
    if answer := CACHE.get(prompt.strip().lower()):
        return Verdict.cached(answer)
    return Verdict.allow()

A check receives a GuardrailRequest and returns a Verdict:

Return	Effect
`Verdict.allow()`	The call proceeds.
`Verdict.block(message, reason)`	The caller gets `403` `guardrail_blocked` with `message`. `reason` is a short label recorded with the request’s usage.
`Verdict.cached(content)`	The caller gets `content` as the model’s answer, and the model is never called. Pre-call stage only.

reason must be lowercase letters, digits, _, or -, up to 64 characters, for example pii_ssn. It is recorded, so it must never contain user content. Verdict.block() raises ValueError for any other value. message is shown to the caller and is not recorded.

GuardrailRequest carries the fields described in the guardrail protocol, plus two helpers:

last_user_message() returns the most recent user message.
text() returns the text the stage is about: all message content on the pre-call stage, the model’s answer on the post-call stage.

Deploy it

Wrap the check in a GuardrailAppEnvironment:

pii_guardrail.py
                
env = GuardrailAppEnvironment(
    name="pii-guardrail",
    pre=check_prompt,
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    # Keep a replica warm: every guarded call waits on this app.
    scaling=flyte.app.Scaling(replicas=(1, 4)),
)

GuardrailAppEnvironment accepts every flyte.app.AppEnvironment option, such as resources, scaling, secrets, and image. The default image installs flyteplugins-union[guardrails]. If your checks need more packages, pass your own image and include flyteplugins-union[guardrails]>=0.19.1 in it.

Deploy it on the gateway’s cluster:

flyte deploy pii_guardrail.py env

Then register it in Settings → Inference Gateway → Guardrails as the pre-call guardrail.

Leave requires_auth at its default. The gateway reaches the app over the cluster-local network and never goes through the app’s public endpoint.

Check answers after the model

Pass a post function to check the model’s answer. One app can serve both stages. Checks can be async:

post_guardrail.py
                
# AWS access key IDs and PEM private keys.
CREDENTIAL = re.compile(r"\bAKIA[0-9A-Z]{16}\b|-----BEGIN [A-Z ]*PRIVATE KEY-----")

async def check_prompt(request: GuardrailRequest) -> Verdict:
    if CREDENTIAL.search(request.text()):
        return Verdict.block("Prompts may not contain credentials.", reason="credential_in_prompt")
    return Verdict.allow()

async def check_answer(request: GuardrailRequest) -> Verdict:
    # On the post stage, request.text() is the model's answer.
    if CREDENTIAL.search(request.text()):
        return Verdict.block("The answer contained a credential and was withheld.", reason="credential_in_answer")
    return Verdict.allow()

env = GuardrailAppEnvironment(
    name="credential-guardrail",
    pre=check_prompt,
    post=check_answer,
    scaling=flyte.app.Scaling(replicas=(1, 4)),
)

Add guardrail routes to an existing FastAPI app

If you already run a FastAPI app, mount the guardrail routes on it with guardrail_router(), and give the environment the Guardrail type:

fastapi_guardrail.py
                
app = FastAPI()

@app.get("/topics")
def topics() -> list[str]:
    return list(BANNED_TOPICS)

app.include_router(guardrail_router(pre=check_prompt))

env = FastAPIAppEnvironment(
    name="topic-guardrail",
    app=app,
    # The gateway only accepts apps of type "Guardrail" as guardrails.
    type="Guardrail",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-union[guardrails]>=0.19.1"),
)

GuardrailAppEnvironment also accepts an existing app, as app=, and adds the routes to it.

Test locally

The guardrail is an ordinary FastAPI app, so you can test it without deploying, using FastAPI’s test client:

        
    
from fastapi.testclient import TestClient

from pii_guardrail import env

client = TestClient(env.app)
verdict = client.post(
    "/guardrail/pre",
    json={"stage": "pre", "input": {"messages": [{"role": "user", "content": "My SSN is 123-45-6789"}]}},
).json()
assert verdict["decision"] == "fail"