Adding guardrails to Embabel agents

Embabel: Adding guardrails to the Agent

In this blog post, you'll follow me as I add guardrails to my Embabel sample application. I wrote an extensive blog about guardrails, and now I want to apply this to my Embabel sample application. If you're not familiar with the sample application, please refer to the post at the link below. I keep this post up to date. Recently, I upgraded the sample to version 0.3.5 and added observability.

Building Agents with Embabel: A Hands-On Introduction

A short recap on guardrails

When reading about guardrails, you'll find that different frameworks use different approaches. However, two types of guardrails are generally identified.

One that verifies user input or assistant output. If everything is fine, the result is passed; if there is a problem, the request is stopped. An example is a user request to have the agent do things we do not allow, which can be prompt hacking. Another example is an agent response that hallucinates, even though we want it to use only the provided context.

The second type is a transformer. These guardrails transform the request or the assistant response. A good example here is PII data: we do not want to send PII to the LLM; therefore, we replace it with placeholders before sending the request.

For more information, check my other blog:

https://jettro.dev/building-ai-agents-safely-pii-jailbreaks-and-real-guardrails-a52245a5c624

Guardrails in Embabel

Embabel supports guardrails out of the box. There are two types of guardrails: UserInputGuardRail and the AssistantMessageGuardRail. The names are clear enough to understand when the guardrails are applied during the agent loop. Below is an example of calling the LLM with a Guardrail applied.

return operationContext.ai()
        .withLlm(
                LlmOptions.fromCriteria(AutoModelSelectionCriteria.INSTANCE)
                        .withTemperature(0.2) // Higher temperature for more creative output
        ).withPromptContributor(Personas.EXTRACTOR)
        .withGuardRails(new PIIUserInputGuardRail(piiAnalyzerClient, presidioProperties.piiTypes()))
        .withToolLoopTransformers(transformer)
        .withToolGroup("mcp-firecrawl")
        .createObject(String.format("""
                Fetch the content of the blog post from the URL that is provided by the user.
                If the user does not provide a URL or if the URL is not valid, return an error 
                message with the problem.
                Provide the content without any boilerplate or additional information.
                Extract all the image urls from the page and return them in a list.

                # User input
                %s
                """, userInput.getContent().trim()), BlogPost.class);

When working with guardrails in Embabel, it is essential to understand the impact of the validation severity. A guardrail returns a ValidationResult. The result can contain validation errors. Each error has a severity. If one of them is of type CRITICAL, the guardrail will block and throw an exception of type GuardRailViolationException.

In Embebal, the guardrail does not support transformations. You need an additional component for that. More on that later in the blog. The first focus is the guardrail. For this example, I create a PII detector guardrail.

Implement a PII detection guardrail

The goal of the guardrail is to detect specific PII and generate validation errors for each detected entity. The guardrail should raise warnings only; it must not block the request. At a later stage, I want to use the found warnings to replace the PII data with placeholders.

PII data detection is done through Presidio. The next section explains how you can work with Presidio.

Next is the code for the complete Guardrail.

package dev.jettro.blogpromotor.agent;

import com.embabel.agent.api.validation.guardrails.UserInputGuardRail;
import com.embabel.agent.core.Blackboard;
import com.embabel.common.core.validation.ValidationError;
import com.embabel.common.core.validation.ValidationResult;
import com.embabel.common.core.validation.ValidationSeverity;
import dev.jettro.blogpromotor.presidio.AnalyzeRequest;
import dev.jettro.blogpromotor.presidio.PresidioAnalyzerClient;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;

/**
 * This class implements a user input guard rail that checks for PII (Personally Identifiable Information) in user input.
 * It uses the Presidio Analyzer API to analyze the input and identify PII entities. If PII is found, it returns a validation
 * error with a level of severity of WARNING, indicating that the input contains sensitive information.
 */
public class PIIUserInputGuardRail implements UserInputGuardRail {
    private static final Logger logger = LoggerFactory.getLogger(PIIUserInputGuardRail.class);
    public static final String PII_ANALYZE_RESULT_KEY = "pii_analyze_result";

    private final PresidioAnalyzerClient presidioAnalyzerClient;
    private final List<String> piiTypes;

    public PIIUserInputGuardRail(PresidioAnalyzerClient presidioAnalyzerClient, List<String> piiTypes) {
        this.presidioAnalyzerClient = presidioAnalyzerClient;
        this.piiTypes = piiTypes;
    }

    @NotNull
    @Override
    public String getName() {
        return "PIIUserInputGuardRail";
    }

    @NotNull
    @Override
    public String getDescription() {
        return "Finds PII data in user input and rejects it.";
    }

    @NotNull
    @Override
    public ValidationResult validate(String input, @NotNull Blackboard blackboard) {
        logger.info("Validating input: {}", input);

        var request = AnalyzeRequest.builder()
                .text(input)
                .language("en")
                .entities(piiTypes)
                .build();

        var analyzeResult = presidioAnalyzerClient.analyze(request);

        blackboard.set(PII_ANALYZE_RESULT_KEY, analyzeResult);

        if (analyzeResult.isEmpty()) {
            return new ValidationResult(true, List.of());
        }

        var errors = analyzeResult.stream()
                .map(result -> {
                    var foundValue = input.substring(result.start(), result.end());
                    return new ValidationError("pii",
                            String.format("Found entity of type %s: with value '%s'", result.entityType(), foundValue),
                            ValidationSeverity.WARNING);
                })
                .toList();

        return new ValidationResult(false, errors);
    }
}

Notice how I store the analysis result on the Blackboard. The blackboard is provided by the framework and, therefore, is easily accessible. In the current scenario, I add the presidio result to the blackboard.

Interact with Microsoft Presidio

You can find more information about Presidio in the blog post about guardrails I mentioned earlier. You can also go to the original documentation. I run Presidio as a Docker container.

docker pull mcr.microsoft.com/presidio-analyzer
docker run -d -p 5002:3000 mcr.microsoft.com/presidio-analyzer:latest

Wrapping a rest endpoint with Spring Boot is fun; you just need some config, and you are good to go. First, the endpoint and the PII types we want to extract.

@ConfigurationProperties(prefix = "presidio.analyzer")
public record PresidioProperties(
        String baseUrl,
        List<String> piiTypes
) {
}

Next, you initialise the client to connect to the endpoint.

package dev.jettro.blogpromotor.presidio;

import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.client.RestClient;
import org.springframework.web.client.support.RestClientAdapter;
import org.springframework.web.service.invoker.HttpServiceProxyFactory;

@Configuration
@EnableConfigurationProperties(PresidioProperties.class)
public class PresidioConfig {

    @Bean
    public PresidioAnalyzerClient presidioAnalyzerClient(PresidioProperties properties,
                                                         RestClient.Builder restClientBuilder) {
        RestClient restClient = restClientBuilder
                .baseUrl(properties.baseUrl())
                .build();
        RestClientAdapter adapter = RestClientAdapter.create(restClient);
        HttpServiceProxyFactory factory = HttpServiceProxyFactory.builderFor(adapter).build();
        return factory.createClient(PresidioAnalyzerClient.class);
    }
}

Finally, the service that you can use in other components.

package dev.jettro.blogpromotor.presidio;

import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.service.annotation.HttpExchange;
import org.springframework.web.service.annotation.PostExchange;

import java.util.List;

@HttpExchange
public interface PresidioAnalyzerClient {

    @PostExchange("/analyze")
    List<AnalyzeResult> analyze(@RequestBody AnalyzeRequest request);
}

Check the code for the guardrail to see how you can use the PresidioAnalyzerClient.

Use the blackboard for transformations

To transform content before sending it off to an LLM, you can implement the interface ToolLoopTransformer. The code for the LLM call I showed in the beginning has the transformer configured using withToolLoopTransformers(..).

Below is the code for the interface to implement

public class PIIToolLoopTransformer implements ToolLoopTransformer {
    private static final Logger logger = LoggerFactory.getLogger(PIIToolLoopTransformer.class);

    public PIIToolLoopTransformer() {
    }

    @NotNull
    @Override
    public List<Message> transformBeforeLlmCall(@NotNull BeforeLlmCallContext context) {
        // TODO
    }

    @NotNull
    @Override
    public Message transformAfterLlmCall(@NotNull AfterLlmCallContext context) {
        var response = context.getResponse();
        logger.info("After llm call response: {}", context.getResponse().getContent());
        return response;
    }

    @NotNull
    @Override
    public String transformAfterToolResult(@NotNull AfterToolResultContext context) {
        return ToolLoopTransformer.super.transformAfterToolResult(context);
    }

    @NotNull
    @Override
    public List<Message> transformAfterIteration(@NotNull AfterIterationContext context) {
        return ToolLoopTransformer.super.transformAfterIteration(context);
    }
}

First, you have to check the blackboard for validation problems. The next method checks if the Blackboard has PII validation problems.

private Optional<List<AnalyzeResult>> getPiiAnalyzeResult(Blackboard blackboard) {
    Object piiAnalyzeResult = blackboard.get(PII_ANALYZE_RESULT_KEY);
    // Use pattern matching, for instanceof to safely cast and check for null
    if (!(piiAnalyzeResult instanceof List<?> rawList)) {
        logger.info("No PII analyze result found or it is not a list.");
        return Optional.empty();
    }
    // Ensure the list contains the expected types (optional but recommended for safety)
    List<AnalyzeResult> analyzeResultList = rawList.stream()
            .filter(AnalyzeResult.class::isInstance)
            .map(AnalyzeResult.class::cast)
            .toList();

    if (analyzeResultList.isEmpty()) {
        return Optional.empty();
    }

    // Log the found PII entities
    analyzeResultList.forEach(result ->
            logger.info("Entity: {}, Start: {}, End: {}", result.entityType(), result.start(), result.end())
    );

    return Optional.of(analyzeResultList);
}

Next, I implement the transformBeforeLlmCall. First, obtain the blackboard and check for any PII messages.

logger.info("Before llm is called, the blackboard is checked for PII entities.");
// Read the PII analyze result from the blackboard, if none is found, return the original history
AgentProcess agentProcess = AgentProcess.get();
if (agentProcess == null) {
    logger.error("AgentProcess is null, this is unexpected.");
    throw new RuntimeException("There is no reference to the agent process");
}
Blackboard blackboard = agentProcess.getBlackboard();

Optional<List<AnalyzeResult>> optionalPiiAnalyzeResult = getPiiAnalyzeResult(blackboard);
if (optionalPiiAnalyzeResult.isEmpty()) {
    return context.getHistory();
}
var piiAnalyzeResult = optionalPiiAnalyzeResult.get();

Notice how I return the history in the context if everything is fine and no PII data exists. Next, we fetch the user message from the history. Obtain the text and, using the analysis result, replace the found entities with placeholders. After that, replace the user message in the history with the new message containing the placeholders.

// Find the last message in the conversation and only continue if it is a user message
var history = context.getHistory();
if (history.isEmpty() || !(history.getLast() instanceof UserMessage)) {
    return history;
}
var lastMessage = history.getLast();

// Replace PII entities with placeholders in the form of <ENTITY_TYPE>
var text = lastMessage.getContent();
StringBuilder piiTransformedStringBuilder = new StringBuilder(text);

// Sort results by start index descending to avoid shifting issues when replacing text
var sortedPiiAnalyzeResult = piiAnalyzeResult.stream()
        .sorted(Comparator.comparingInt(AnalyzeResult::start).reversed())
        .toList();

for (var result : sortedPiiAnalyzeResult) {
    piiTransformedStringBuilder.replace(result.start(), result.end(), "<" + result.entityType() + ">");
}

// Replace the message in the history with the transformed one
Message transformedMessage = new UserMessage(piiTransformedStringBuilder.toString(),
        lastMessage.getRole().getDisplayName(),
        lastMessage.getTimestamp());
history.set(history.size() - 1, transformedMessage);

Now everything is in place. If things are not clear, check the unit tests in the project. They test the different components of the solution.

Time to test the system. I run the application and execute the following command.

execute “Create a linkedin post for this blog post from Jettro Coenradie: https://coenradie.com/posts/welcome-to-coenradie-com/"

The PII Guardrail and transformer should replace my name with a placeholder before sending it to the LLM. The next image shows the log from OpenAI.

The PII person transformer replace my name with the PERSON placeholder.

Famous last words

When working with PII data, be warned that with these tools, you can only influence what is sent to the LLM. If you have observability running, you might still log all the PII data to, for instance, Langfuse.

Another option could be to change the UserInput in the first action of Embabel.

If you want to have a look at the example, you can find the code in this repository.

[embed]GitHub - jettro/embabel-agent-blog-promotor: A repository to accompany my experiments with the new… A repository to accompany my experiments with the new Embabel project for agents - jettro/embabel-agent-blog-promotorgithub.com