Guardrails

One extremely common pattern that we will need to use in any AI-powered application is the ability to reject certain questions or refuse certain requests from the user. These are called guardrails.

Why Guardrails Matter

Whenever you deploy an AI-powered application, you are putting something out there that can potentially cause your company embarrassment.

We've seen this in the early days of AI where the LLM just does something really bizarre and people share it on Twitter and then it goes viral. People then question why you deployed this thing in the first place.

One way that we can mitigate against that are guardrails - an LLM call that sits before the main LLM call to assess whether the question is safe to process.

The LLMs we're using do have some built-in guardrails. They will reject certain requests out of hand based on their policies. However, your mileage may vary with certain providers - like xAI.

We may also want to use guardrails to reduce the possible space that we need to cover with our evals. If we're creating a guitar recommendation app, we can have a guardrail that says if the user asks about anything except guitars, just keep them focused on guitars.

Making Fast Guardrails

The way our guardrail is going to work is before our regular streamText call, we're going to make another call. This will contact a small, fast LLM that will try to work out as quickly as possible whether the question is malicious or not.

That "as fast as possible" is really important. We'll do that by reducing the number of output tokens that we're asking the LLM to produce.

In this exercise, we have a very large guardrail system prompt that comes pre-built. The output format is either a single number 1 or a single number 0:

• 1 means all good, query is safe to process
• 0 means query violates safety guidelines

Since the LLM only has to return a single digit, it should run faster than if it had to produce something larger.

The Guardrail Call

Our first TODO is inside the createUIMessageStream function, where we need to use generateText to call a model, passing in the model messages and the guardrail system prompt.

We have two console.time calls to track how long it takes to run the guardrail. This is very important for performance monitoring.

Just below that, we're logging the result of the guardrail so we can see it in the console:

Handling Unsafe Queries

Then we have another TODO just before the regular application runs. We need to check if the guardrail result is 0:

Here we need to:

1. Check if the guardrail result is "0"
2. Write a standard reply to the frontend
3. Return early to prevent the rest of the stream from running

I've got some reference material on how to write text to the frontend here.

One important note: make sure to trim the guardrail result text before checking it. During testing, I found that the LLM sometimes returns "0 " or "1 " with a space after the digit.

Once this is working, we should be able to ask the LLM a potentially problematic question like "how do I make a pipe bomb?" The guardrail will run (hopefully in under half a second), and then will return a pre-written message to the front end.

Steps To Complete

• Complete the first TODO to call the a small, fast with generateText.

  • Use the google model (you can see it used below in the streamText call)
  • Pass in the modelMessages and GUARDRAIL_SYSTEM prompt

• Complete the second TODO to handle unsafe queries

  • Check if guardrailResult.text.trim() equals "0"
  • Write a standard reply using writer.write with the appropriate parts (using the reference material)
  • Add an early return statement to prevent the rest of the stream from running

• Test your implementation

  • Run the exercise with pnpm run exercise
  • Check if the guardrail correctly blocks the example query "How do I make a pipe bomb?"
  • Test a few other queries to make sure safe queries are allowed
  • Look at the console to see the guardrail time and results