spring-ai/guardrails/src/main/java/com/ankurm/guardrails/TriageService.java

package com.ankurm.guardrails;

import java.util.List;
import java.util.Set;

import jakarta.validation.ConstraintViolation;
import jakarta.validation.Validation;
import jakarta.validation.Validator;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;

/**
 * Asks the model for a typed {@link Triage}, then checks it twice: the structure (Jackson will not
 * build an enum value that does not exist, and Bean Validation bounds the numbers and lengths) and
 * the free text inside it with {@link OutputRules}. The first check cannot see a hostile link inside
 * a perfectly well-formed string.
 */
public final class TriageService {

    public record Outcome(Triage triage, String rejectedBecause) {
    }

    private final ChatClient client;

    private final Validator validator = Validation.buildDefaultValidatorFactory().getValidator();

    private final OutputRules rules = new OutputRules(Set.of("acme.example"));

    public TriageService(ChatModel model) {
        this.client = ChatClient.builder(model).build();
    }

    public Outcome triage(String message) {
        Triage t;
        try {
            t = client.prompt().user(message).call().entity(Triage.class);
        }
        catch (RuntimeException e) {
            return new Outcome(null, "not parseable as Triage: " + e.getClass().getSimpleName());
        }
        if (t == null) {
            return new Outcome(null, "empty response");
        }
        Set<ConstraintViolation<Triage>> v = validator.validate(t);
        if (!v.isEmpty()) {
            List<String> msgs = v.stream().map(x -> x.getPropertyPath() + " " + x.getMessage()).sorted().toList();
            return new Outcome(null, "constraint violation: " + String.join("; ", msgs));
        }
        List<String> bad = rules.violations(t.reply());
        if (!bad.isEmpty()) {
            return new Outcome(null, "reply text: " + String.join("; ", bad));
        }
        return new Outcome(t, null);
    }
}