Improve response parsing

ballerina/Ballerina.toml

@@ -1,7 +1,7 @@
 [package]
 org = "ballerinax"
 name = "np"
-version = "0.8.1"
+version = "0.8.2"
 authors = ["Ballerina"]
 keywords = ["natural programming", "ai"]
 repository = "https://github.com/ballerina-platform/module-ballerinax-np"
@@ -17,5 +17,5 @@ observabilityIncluded = true
 [[platform.java21.dependency]]
 groupId = "io.ballerina.lib"
 artifactId = "np-native"
-version = "0.8.1"
-path = "../native/build/libs/np-native-0.8.1.jar"
+version = "0.8.2"
+path = "../native/build/libs/np-native-0.8.2-SNAPSHOT.jar"

ballerina/CompilerPlugin.toml

@@ -3,4 +3,4 @@ id = "np-compiler-plugin"
 class = "io.ballerina.lib.np.compilerplugin.CompilerPlugin"
 
 [[dependency]]
-path = "../compiler-plugin/build/libs/np-compiler-plugin-0.8.1.jar"
+path = "../compiler-plugin/build/libs/np-compiler-plugin-0.8.2-SNAPSHOT.jar"

ballerina/Dependencies.toml

@@ -335,7 +335,7 @@ modules = [
 [[package]]
 org = "ballerinax"
 name = "np"
-version = "0.8.1"
+version = "0.8.2"
 dependencies = [
 	{org = "ballerina", name = "http"},
 	{org = "ballerina", name = "jballerina.java"},

ballerina/main.bal

@@ -100,7 +100,17 @@ isolated function callLlmGeneric(Prompt prompt, Context context, typedesc<json>
 }
 
 isolated function parseResponseAsJson(string resp) returns json|error {
-    string processedResponse = re `${"```json|```"}`.replaceAll(resp, "");
+    int startDelimLength = 7;
+    int? startIndex = resp.indexOf("```json");
+    if startIndex is () {
+        startIndex = resp.indexOf("```");
+        startDelimLength = 3;
+    }
+    int? endIndex = resp.lastIndexOf("```");
+
+    string processedResponse = startIndex is () || endIndex is () ? 
+        resp : 
+        resp.substring(startIndex + startDelimLength, endIndex).trim();
     json|error result = trap processedResponse.fromJsonString();
     if result is error {
         return handlepParseResponseError(result);

ballerina/tests/test_utils.bal

@@ -15,6 +15,14 @@ isolated function getExpectedPrompt(string message) returns string {
         return expectedPromptStringForRateBlog4;
     }
 
+    if message.startsWith("What's the output of the Ballerina code below?") {
+        return expectedPromptStringForBalProgram;
+    }
+
+    if message.startsWith("Which country") {
+        return expectedPromptStringForCountry;
+    }
+
     return "INVALID";
 }
 
@@ -35,5 +43,13 @@ isolated function getTheMockLLMResult(string message) returns string {
         return "[{\"name\":\"Virat Kohli\",\"age\":33},{\"name\":\"Kane Williamson\",\"age\":30}";
     }
 
+    if message.startsWith("What's the output of the Ballerina code below?") {
+        return string `The output of the provided Ballerina code calculates the sum of ${"`"}x${"`"} and ${"`"}y${"`"}, which is ${"`"}10 + 20${"`"}. Therefore, the result will be ${"`"}30${"`"}. \n\nHere is the output formatted as a JSON value that satisfies your specified schema:${"\n\n```"}json${"\n"}30${"\n```"}`;
+    }
+
+    if message.startsWith("Which country") {
+        return "```\n\"Sri Lanka\"\n```";
+    }
+
     return "INVALID";
 }

ballerina/tests/test_values.bal

@@ -63,3 +63,28 @@ final string expectedPromptStringForRateBlog4 = string `Tell me name and the age
 
         Schema:
         {"$schema":"https://json-schema.org/draft/2020-12/schema", "type":"array", "items":{"$schema":"https://json-schema.org/draft/2020-12/schema", "type":"object", "properties":{"name":{"type":"string"}}, "required":["name"]}}`;
+
+final string expectedPromptStringForBalProgram = string `What's the output of the Ballerina code below?
+
+    ${"```"}ballerina
+    import ballerina/io;
+
+    public function main() {
+        int x = 10;
+        int y = 20;
+        io:println(x + y);
+    \}
+    ${"```"}.  
+        The output should be a JSON value that satisfies the following JSON schema, 
+        returned within a markdown snippet enclosed within ${"```json"} and ${"```"}
+
+        Schema:
+        {"type":"integer"}`;
+
+final string expectedPromptStringForCountry = string `Which country is known as the pearl of the Indian Ocean?.  
+        The output should be a JSON value that satisfies the following JSON schema, 
+        returned within a markdown snippet enclosed within ${"```json"} and ${"```"}
+
+        Schema:
+        {"type":"string"}`;
+

ballerina/tests/tests.bal

@@ -54,3 +54,25 @@ function testJsonConversionError2() {
     test:assertTrue(rating is error);
     test:assertTrue((<error> rating).message().includes(ERROR_MESSAGE));
 }
+
+@test:Config
+function testJsonContentAfterTextDescription() returns error? {
+    int result = check callLlm(`What's the output of the Ballerina code below?
+
+    ${"```"}ballerina
+    import ballerina/io;
+
+    public function main() {
+        int x = 10;
+        int y = 20;
+        io:println(x + y);
+    \}
+    ${"```"}`);
+    test:assertEquals(result, 30);
+}
+
+@test:Config
+function testJsonContentWithoutJsonAfterBackticks() returns error? {
+    string result = check callLlm(`Which country is known as the pearl of the Indian Ocean?`);
+    test:assertEquals(result, "Sri Lanka");
+}