From b03f26d48e15038656654c2432f8df32f33357b2 Mon Sep 17 00:00:00 2001
From: Eric Deandrea <eric.deandrea@gmail.com>
Date: Wed, 8 Jan 2025 07:58:59 -0500
Subject: [PATCH] Allow customization of invalid json reprompt

---
 .../AbstractJsonExtractorOutputGuardrail.java    | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/AbstractJsonExtractorOutputGuardrail.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/AbstractJsonExtractorOutputGuardrail.java
index 488c832ca..ff25c40ab 100644
--- a/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/AbstractJsonExtractorOutputGuardrail.java
+++ b/core/runtime/src/main/java/io/quarkiverse/langchain4j/guardrails/AbstractJsonExtractorOutputGuardrail.java
@@ -40,9 +40,19 @@ public OutputGuardrailResult validate(AiMessage responseFromLLM) {
             }
         }
 
-        return reprompt("Invalid JSON",
-                "Make sure you return a valid JSON object following "
-                        + "the specified format");
+        return invokeInvalidJson(responseFromLLM, json);
+    }
+
+    protected OutputGuardrailResult invokeInvalidJson(AiMessage aiMessage, String json) {
+        return reprompt(getInvalidJsonMessage(aiMessage, json), getInvalidJsonReprompt(aiMessage, json));
+    }
+
+    protected String getInvalidJsonMessage(AiMessage aiMessage, String json) {
+        return "Invalid JSON";
+    }
+
+    protected String getInvalidJsonReprompt(AiMessage aiMessage, String json) {
+        return "Make sure you return a valid JSON object following the specified format";
     }
 
     protected Object deserialize(String llmResponse) {