Merge pull request #1003 from holly-cummins/guardrail-api

cescoffier · web-flow · commit 31600d716b58 · 2024-10-22T08:35:14.000+02:00
Move inner classes to top level classes for API concision
diff --git a/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/GuardrailWithAugmentationTest.java b/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/GuardrailWithAugmentationTest.java
@@ -33,9 +33,11 @@
 import dev.langchain4j.service.UserMessage;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.InputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrails;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrails;
 import io.quarkus.test.QuarkusUnitTest;
diff --git a/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/InputAndOutputGuardrailsTest.java b/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/InputAndOutputGuardrailsTest.java
@@ -27,9 +27,11 @@
 import dev.langchain4j.service.UserMessage;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.InputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrails;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrails;
 import io.quarkus.test.QuarkusUnitTest;
diff --git a/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/InputGuardrailValidationTest.java b/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/InputGuardrailValidationTest.java
@@ -30,6 +30,7 @@
 import dev.langchain4j.service.UserMessage;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.InputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrails;
 import io.quarkiverse.langchain4j.runtime.aiservice.GuardrailException;
@@ -179,7 +180,7 @@ public static class MemoryCheck implements InputGuardrail {
         AtomicInteger spy = new AtomicInteger(0);
 
         @Override
-        public InputGuardrailResult validate(InputGuardrail.InputGuardrailParams params) {
+        public InputGuardrailResult validate(InputGuardrailParams params) {
             spy.incrementAndGet();
             if (params.memory().messages().isEmpty()) {
                 assertThat(params.userMessage().singleText()).isEqualTo("foo");
diff --git a/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/OutputGuardrailRepromptingRetryDisabledTest.java b/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/OutputGuardrailRepromptingRetryDisabledTest.java
@@ -29,6 +29,7 @@
 import dev.langchain4j.service.SystemMessage;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrails;
 import io.quarkiverse.langchain4j.runtime.aiservice.GuardrailException;
diff --git a/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/OutputGuardrailRepromptingTest.java b/core/deployment/src/test/java/io/quarkiverse/langchain4j/test/guardrails/OutputGuardrailRepromptingTest.java
@@ -30,6 +30,7 @@
 import dev.langchain4j.service.SystemMessage;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrails;
 import io.quarkiverse.langchain4j.runtime.aiservice.GuardrailException;
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/InputGuardrail.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/InputGuardrail.java
@@ -3,8 +3,6 @@
 import java.util.Arrays;
 
 import dev.langchain4j.data.message.UserMessage;
-import dev.langchain4j.memory.ChatMemory;
-import dev.langchain4j.rag.AugmentationResult;
 import io.smallrye.common.annotation.Experimental;
 
 /**
@@ -14,7 +12,7 @@
  * Implementation should be exposed as a CDI bean, and the class name configured in {@link InputGuardrails#value()} annotation.
  */
 @Experimental("This feature is experimental and the API is subject to change")
-public interface InputGuardrail extends Guardrail<InputGuardrail.InputGuardrailParams, InputGuardrailResult> {
+public interface InputGuardrail extends Guardrail<InputGuardrailParams, InputGuardrailResult> {
 
     /**
      * Validates the {@code user message} that will be sent to the LLM.
@@ -42,17 +40,6 @@ default InputGuardrailResult validate(InputGuardrailParams params) {
         return validate(params.userMessage());
     }
 
-    /**
-     * Represents the parameter passed to {@link #validate(InputGuardrailParams)}.
-     *
-     * @param userMessage the user message, cannot be {@code null}
-     * @param memory the memory, can be {@code null} or empty
-     * @param augmentationResult the augmentation result, can be {@code null}
-     */
-    record InputGuardrailParams(UserMessage userMessage, ChatMemory memory,
-            AugmentationResult augmentationResult) implements GuardrailParams {
-    }
-
     /**
      * @return The result of a successful input guardrail validation.
      */
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/InputGuardrailParams.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/InputGuardrailParams.java
@@ -0,0 +1,16 @@
+package io.quarkiverse.langchain4j.guardrails;
+
+import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.memory.ChatMemory;
+import dev.langchain4j.rag.AugmentationResult;
+
+/**
+ * Represents the parameter passed to {@link InputGuardrail#validate(InputGuardrailParams)}.
+ *
+ * @param userMessage the user message, cannot be {@code null}
+ * @param memory the memory, can be {@code null} or empty
+ * @param augmentationResult the augmentation result, can be {@code null}
+ */
+public record InputGuardrailParams(UserMessage userMessage, ChatMemory memory,
+        AugmentationResult augmentationResult) implements GuardrailParams {
+}
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/OutputGuardrail.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/OutputGuardrail.java
@@ -3,8 +3,6 @@
 import java.util.Arrays;
 
 import dev.langchain4j.data.message.AiMessage;
-import dev.langchain4j.memory.ChatMemory;
-import dev.langchain4j.rag.AugmentationResult;
 import io.smallrye.common.annotation.Experimental;
 
 /**
@@ -18,7 +16,7 @@
  * The maximum number of retries is configurable using {@code quarkus.langchain4j.guardrails.max-retries}, defaulting to 3.
  */
 @Experimental("This feature is experimental and the API is subject to change")
-public interface OutputGuardrail extends Guardrail<OutputGuardrail.OutputGuardrailParams, OutputGuardrailResult> {
+public interface OutputGuardrail extends Guardrail<OutputGuardrailParams, OutputGuardrailResult> {
 
     /**
      * Validates the response from the LLM.
@@ -45,17 +43,6 @@ default OutputGuardrailResult validate(OutputGuardrailParams params) {
         return validate(params.responseFromLLM());
     }
 
-    /**
-     * Represents the parameter passed to {@link #validate(OutputGuardrailParams)}.
-     *
-     * @param responseFromLLM the response from the LLM
-     * @param memory the memory, can be {@code null} or empty
-     * @param augmentationResult the augmentation result, can be {@code null}
-     */
-    record OutputGuardrailParams(AiMessage responseFromLLM, ChatMemory memory,
-            AugmentationResult augmentationResult) implements GuardrailParams {
-    }
-
     /**
      * @return The result of a successful output guardrail validation.
      */
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/OutputGuardrailParams.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/OutputGuardrailParams.java
@@ -0,0 +1,16 @@
+package io.quarkiverse.langchain4j.guardrails;
+
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.memory.ChatMemory;
+import dev.langchain4j.rag.AugmentationResult;
+
+/**
+ * Represents the parameter passed to {@link OutputGuardrail#validate(OutputGuardrailParams)}.
+ *
+ * @param responseFromLLM the response from the LLM
+ * @param memory the memory, can be {@code null} or empty
+ * @param augmentationResult the augmentation result, can be {@code null}
+ */
+public record OutputGuardrailParams(AiMessage responseFromLLM, ChatMemory memory,
+        AugmentationResult augmentationResult) implements GuardrailParams {
+}
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/AiServiceMethodImplementationSupport.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/AiServiceMethodImplementationSupport.java
@@ -62,6 +62,7 @@
 import io.quarkiverse.langchain4j.audit.Audit;
 import io.quarkiverse.langchain4j.audit.AuditService;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.runtime.ContextLocals;
 import io.quarkiverse.langchain4j.runtime.QuarkusServiceOutputParser;
 import io.quarkiverse.langchain4j.runtime.ResponseSchemaUtil;
@@ -305,7 +306,7 @@ private List<ChatMessage> messagesToSend(ChatMessage augmentedUserMessage,
 
         response = GuardrailsSupport.invokeOutputGuardrails(methodCreateInfo, chatMemory, context.chatModel, response,
                 toolSpecifications,
-                new OutputGuardrail.OutputGuardrailParams(response.content(), chatMemory, augmentationResult));
+                new OutputGuardrailParams(response.content(), chatMemory, augmentationResult));
 
         // everything worked as expected so let's commit the messages
         chatMemory.commit();
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/GuardrailsSupport.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/GuardrailsSupport.java
@@ -19,8 +19,10 @@
 import io.quarkiverse.langchain4j.guardrails.GuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.GuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.InputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.InputGuardrailResult;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
+import io.quarkiverse.langchain4j.guardrails.OutputGuardrailParams;
 import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
 
 public class GuardrailsSupport {
@@ -30,7 +32,7 @@ public static void invokeInputGuardrails(AiServiceMethodCreateInfo methodCreateI
         InputGuardrailResult result;
         try {
             result = invokeInputGuardRails(methodCreateInfo,
-                    new InputGuardrail.InputGuardrailParams(userMessage, chatMemory, augmentationResult));
+                    new InputGuardrailParams(userMessage, chatMemory, augmentationResult));
         } catch (Exception e) {
             throw new GuardrailException(e.getMessage(), e);
         }
@@ -44,7 +46,7 @@ public static Response<AiMessage> invokeOutputGuardrails(AiServiceMethodCreateIn
             ChatLanguageModel chatModel,
             Response<AiMessage> response,
             List<ToolSpecification> toolSpecifications,
-            OutputGuardrail.OutputGuardrailParams output) {
+            OutputGuardrailParams output) {
         int attempt = 0;
         int max = methodCreateInfo.getGuardrailsMaxRetry();
         if (max <= 0) {
@@ -80,7 +82,7 @@ public static Response<AiMessage> invokeOutputGuardrails(AiServiceMethodCreateIn
                     chatMemory.add(response.content());
                 }
                 attempt++;
-                output = new OutputGuardrail.OutputGuardrailParams(response.content(), output.memory(),
+                output = new OutputGuardrailParams(response.content(), output.memory(),
                         output.augmentationResult());
             } else {
                 break;
@@ -95,7 +97,7 @@ public static Response<AiMessage> invokeOutputGuardrails(AiServiceMethodCreateIn
 
     @SuppressWarnings("unchecked")
     private static OutputGuardrailResult invokeOutputGuardRails(AiServiceMethodCreateInfo methodCreateInfo,
-            OutputGuardrail.OutputGuardrailParams params) {
+            OutputGuardrailParams params) {
         if (methodCreateInfo.getOutputGuardrailsClassNames().isEmpty()) {
             return OutputGuardrailResult.success();
         }
@@ -122,7 +124,7 @@ private static OutputGuardrailResult invokeOutputGuardRails(AiServiceMethodCreat
 
     @SuppressWarnings("unchecked")
     private static InputGuardrailResult invokeInputGuardRails(AiServiceMethodCreateInfo methodCreateInfo,
-            InputGuardrail.InputGuardrailParams params) {
+            InputGuardrailParams params) {
         if (methodCreateInfo.getInputGuardrailsClassNames().isEmpty()) {
             return InputGuardrailResult.success();
         }
diff --git a/docs/modules/ROOT/pages/guardrails.adoc b/docs/modules/ROOT/pages/guardrails.adoc
@@ -35,7 +35,7 @@ import dev.langchain4j.rag.AugmentationResult;
  * <p>
  * Implementation should be exposed as a CDI bean, and the class name configured in {@link InputGuardrails#value()} annotation.
  */
-public interface InputGuardrail extends Guardrail<InputGuardrail.InputGuardrailParams, InputGuardrailResult> {
+public interface InputGuardrail extends Guardrail<InputGuardrailParams, InputGuardrailResult> {
 
     /**
      * Validates the {@code user message} that will be sent to the LLM.
@@ -64,7 +64,7 @@ public interface InputGuardrail extends Guardrail<InputGuardrail.InputGuardrailP
     }
 
     /**
-     * Represents the parameter passed to {@link #validate(InputGuardrailParams)}.
+     * Represents the parameter passed to {@link InputGuardrail#validate(InputGuardrailParams)}.
      *
      * @param userMessage the user message, cannot be {@code null}
      * @param memory the memory, can be {@code null} or empty
@@ -216,7 +216,7 @@ import dev.langchain4j.rag.AugmentationResult;
  * <p>
  * The maximum number of retries is configurable using {@code quarkus.langchain4j.guardrails.max-retries}, defaulting to 3.
  */
-public interface OutputGuardrail extends Guardrail<OutputGuardrail.OutputGuardrailParams, OutputGuardrailResult> {
+public interface OutputGuardrail extends Guardrail<OutputGuardrailParams, OutputGuardrailResult> {
 
     /**
      * Validates the response from the LLM.
@@ -244,7 +244,7 @@ public interface OutputGuardrail extends Guardrail<OutputGuardrail.OutputGuardra
     }
 
     /**
-     * Represents the parameter passed to {@link #validate(OutputGuardrailParams)}.
+     * Represents the parameter passed to {@link OutputGuardrail#validate(OutputGuardrailParams)}.
      *
      * @param responseFromLLM the response from the LLM
      * @param memory the memory, can be {@code null} or empty
