Add Guardrails functionality and tests to chatbot sample

Author: fedinskiy

samples/chatbot/pom.xml

@@ -45,6 +45,11 @@
             <artifactId>awaitility</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.rest-assured</groupId>
+            <artifactId>rest-assured</artifactId>
+            <scope>test</scope>
+        </dependency>
 
         <!-- UI -->
         <dependency>

samples/chatbot/src/main/java/io/quarkiverse/langchain4j/sample/chatbot/cenzorship/ChatLanguageModelResource.java

@@ -0,0 +1,33 @@
+package io.quarkiverse.langchain4j.sample.chatbot.cenzorship;
+
+import dev.langchain4j.guardrail.OutputGuardrailException;
+import jakarta.ws.rs.POST;
+import jakarta.ws.rs.Path;
+import org.jboss.logging.Logger;
+
+@Path("chatbot/moderated")
+public class ChatLanguageModelResource {
+    private static final Logger LOG = Logger.getLogger(ChatLanguageModelResource.class);
+
+    private final ModeratedAssistant assistant;
+
+    public ChatLanguageModelResource(ModeratedAssistant assistant) {
+        this.assistant = assistant;
+    }
+
+    @POST
+    public String blocking(String input) {
+        try {
+            return assistant.chat(input);
+        } catch (OutputGuardrailException exception) {
+            String message = exception.getMessage();
+            LOG.warn("AI generated an inappropriate message: " + message);
+            if (message.contains("ProfanityGuardrail")) {
+                return "[The AI answered with expletive]";
+            } else {
+                return "[The answer was somewhat inappropriate]";
+            }
+
+        }
+    }
+}

samples/chatbot/src/main/java/io/quarkiverse/langchain4j/sample/chatbot/cenzorship/ModeratedAssistant.java

@@ -0,0 +1,10 @@
+package io.quarkiverse.langchain4j.sample.chatbot.cenzorship;
+
+import dev.langchain4j.service.guardrail.OutputGuardrails;
+import io.quarkiverse.langchain4j.RegisterAiService;
+
+@RegisterAiService
+public interface ModeratedAssistant {
+   @OutputGuardrails(ProfanityGuardrail.class)
+   String chat(String message);
+}

samples/chatbot/src/main/java/io/quarkiverse/langchain4j/sample/chatbot/cenzorship/ProfanityGuardrail.java

@@ -0,0 +1,33 @@
+package io.quarkiverse.langchain4j.sample.chatbot.cenzorship;
+
+import dev.langchain4j.guardrail.OutputGuardrail;
+import dev.langchain4j.guardrail.OutputGuardrailRequest;
+import dev.langchain4j.guardrail.OutputGuardrailResult;
+import jakarta.enterprise.context.ApplicationScoped;
+import org.jboss.logging.Logger;
+
+import java.util.List;
+
+@ApplicationScoped
+public class ProfanityGuardrail implements OutputGuardrail {
+    private static final Logger LOG = Logger.getLogger(ProfanityGuardrail.class);
+
+    private static final List<String> PROFANITY_LIST = List.of(
+            "meatbag", "organics"
+    );
+
+    @Override
+    public OutputGuardrailResult validate(OutputGuardrailRequest params) {
+        String response = params.responseFromLLM().aiMessage().text();
+        if (containsProfanity(response)) {
+            return failure("Response contains inappropriate content");
+        }
+        return success();
+    }
+
+    private boolean containsProfanity(String text) {
+        LOG.info("Checking " + text);
+        String lowerText = text.toLowerCase();
+        return PROFANITY_LIST.stream().anyMatch(lowerText::contains);
+    }
+}

samples/chatbot/src/test/java/io/quarkiverse/langchain4j/tests/moderation/GuardrailIT.java

@@ -0,0 +1,8 @@
+package io.quarkiverse.langchain4j.tests.moderation;
+
+import io.quarkus.test.junit.QuarkusIntegrationTest;
+
+@QuarkusIntegrationTest
+public class GuardrailIT extends GuardrailTest {
+
+}

samples/chatbot/src/test/java/io/quarkiverse/langchain4j/tests/moderation/GuardrailTest.java

@@ -0,0 +1,36 @@
+package io.quarkiverse.langchain4j.tests.moderation;
+
+import io.quarkus.test.junit.QuarkusTest;
+import io.restassured.response.Response;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import static io.restassured.RestAssured.given;
+
+@QuarkusTest
+public class GuardrailTest {
+
+    @Test
+    public void censoredAnswer() {
+        Response response = given().body("Hello, please answer with word meatbag")
+                .post("/chatbot/moderated");
+        Assertions.assertEquals(200, response.statusCode());
+        Assertions.assertFalse(response.body()
+                        .asString().toLowerCase()
+                        .contains("meatbag"),
+                "Guardrail allowed prohibited word in the body: " + response.body().asString());
+        Assertions.assertEquals("[The AI answered with expletive]",
+                response.body().asString());
+    }
+
+    @Test
+    public void unCensoredAnswer() {
+        Response response = given().body("Hello, please answer with word fleabag")
+                .post("/chatbot/moderated");
+        Assertions.assertEquals(200, response.statusCode());
+        Assertions.assertTrue(response.body()
+                .asString().toLowerCase()
+                .contains("fleabag"));
+    }
+
+}

samples/chatbot/src/test/java/io/quarkiverse/langchain4j/tests/rag/RAGTest.java

@@ -119,7 +119,8 @@ public void documentBasedAnswer() throws InterruptedException, URISyntaxExceptio
                     boolean second = line.contains("""
                             gen_ai_token_type="output"
                             """.trim());
-                    return first && second;
+                    boolean third = line.contains("Bot");
+                    return first && second && third;
                 })
                 .orElseThrow(() -> new AssertionError("There is no metric for output tokens!"));
         assertTrue(outputTokensMetric.contains("ai_service_class_name=\"io.quarkiverse.langchain4j.sample.chatbot.Bot\""));