Embabel: Adding Guardrails to the Agent

Table of Contents

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 Building Agents with Embabel: A Hands-On Introduction. I keep that post up to date — recently I upgraded the sample to version 0.3.5 and added observability.

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: Building AI agents safely

Guardrails in Embabel
#

Embabel supports guardrails out of the box. There are two types: UserInputGuardRail and AssistantMessageGuardRail. The names make clear when they 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)
    ).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 a GuardRailViolationException.

In Embabel, the guardrail does not support transformations — you need an additional component for that. More on that later. The first focus is the guardrail itself.

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 to 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;

/**
 * Checks for PII in user input using the Presidio Analyzer API.
 * Returns WARNING-level validation errors for each detected PII entity.
 */
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 the analysis result is stored on the Blackboard. The blackboard is provided by the framework and is easily accessible across components. This allows the downstream transformer to retrieve the Presidio result without calling the API a second time.

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 configuration properties and the PII types we want to extract:

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

Next, the Spring Boot configuration to initialise the REST client:

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 declarative HTTP client interface:

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 to an LLM, implement the ToolLoopTransformer interface. The LLM call shown earlier configures it via withToolLoopTransformers(transformer).

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.

(https://github.com/jettro/embabel-agent-blog-promotor)

Originally published on jettro.dev

A short recap on guardrails #

Guardrails in Embabel #

Implement a PII detection guardrail #

Interact with Microsoft Presidio #

Use the blackboard for transformations #

Famous last words #