diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java index 9738c7e..7883722 100755 --- a/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java @@ -66,9 +66,18 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { jsonNode = this.objectMapper.readTree(response); } catch (final Exception e) { - LOGGER.warn("Cannot parse json response: {}", completion.response()); - addJsonResponseErrorPrompt(context); - return false; + LOGGER.debug("Initial JSON parse failed, attempting to sanitize: {}", e.getMessage()); + try { + final var response = completion.response().substring(firstCurly, lastCurly + 1); + final var sanitized = sanitizeJson(response); + jsonNode = this.objectMapper.readTree(sanitized); + LOGGER.debug("Successfully parsed after sanitization"); + } + catch (final Exception e2) { + LOGGER.warn("Cannot parse json response even after sanitization: {}", completion.response()); + addJsonResponseErrorPrompt(context); + return false; + } } final var items = new LinkedHashMap>(); @@ -103,12 +112,8 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { private void parseProperty(final LinkedHashMap> items, final Entry nestedProperty) { if (nestedProperty.getValue().isObject()) { - final var nestedValuesIter = nestedProperty.getValue().values(); - while (nestedValuesIter.hasNext()) { - final var v = nestedValuesIter.next(); - if (v.isTextual() || v.isNumber() || v.isBoolean()) { - items.computeIfAbsent(this.prefixFields + nestedProperty.getKey(), k -> new ArrayList<>()).add(v.asText()); - } + for (final var entry : nestedProperty.getValue().properties()) { + this.parseProperty(items, entry); } } else { @@ -121,4 +126,61 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { context.getConversation().replaceTurn(new Conversation.Turn(turn.type(), turn.text() + "\n\nThe previous response could not be fully processed or validated. Please make sure to format the response in valid JSON syntax with properly escaped characters.")); } + private String sanitizeJson(final String json) { + final var result = new StringBuilder(json.length() + 10); + boolean inString = false; + boolean escaped = false; + + for (int i = 0; i < json.length(); i++) { + final char c = json.charAt(i); + + if (escaped) { + result.append(c); + escaped = false; + continue; + } + + if (c == '\\' && inString) { + escaped = true; + result.append(c); + continue; + } + + if (c == '"') { + boolean unescapedQuote = false; + if (inString) { + for (int j = i + 1; j < json.length(); j++) { + final var nextChar = json.charAt(j); + if (Character.isWhitespace(nextChar)) { + continue; + } + if (nextChar == ':' || nextChar == ',' || nextChar == '}' || nextChar == ']') { + break; + } + result.append('\\'); + unescapedQuote = true; + break; + } + } + + inString = unescapedQuote || !inString; + result.append(c); + continue; + } + + if (inString && (c == '}' || c == ']') && (i == 0 || json.charAt(i - 1) != '\\')) { + int insertPos = result.length(); + while (insertPos > 0 && Character.isWhitespace(result.charAt(insertPos - 1))) { + insertPos--; + } + result.insert(insertPos, '"'); + inString = false; + } + + result.append(c); + } + + return result.toString(); + } + } diff --git a/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java b/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java new file mode 100644 index 0000000..759e5eb --- /dev/null +++ b/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java @@ -0,0 +1,369 @@ +package de.extio.lmlib.agent.responsehandler; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; + +import org.junit.jupiter.api.Test; + +import de.extio.lmlib.agent.AgentContext; +import de.extio.lmlib.client.Completion; +import de.extio.lmlib.client.Conversation; + +public class JsonAgentResponseHandlerTest { + + @Test + public void testSimpleJsonParsing() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "name": "John Doe", + "age": "30", + "city": "New York" + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("John Doe", context.getStringValue("name")); + assertEquals("30", context.getStringValue("age")); + assertEquals("New York", context.getStringValue("city")); + } + + @Test + public void testJsonWithPrefix() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler("test_"); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "field1": "value1", + "field2": "value2" + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("value1", context.getStringValue("test_field1")); + assertEquals("value2", context.getStringValue("test_field2")); + } + + @Test + public void testJsonWithPreamble() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + Here is the result: + + { + "status": "success", + "message": "Operation completed" + } + + Additional notes after the JSON. + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("success", context.getStringValue("status")); + assertEquals("Operation completed", context.getStringValue("message")); + } + + @Test + public void testJsonArray() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "items": ["apple", "banana", "cherry"] + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + final List items = context.getStringValues("items"); + assertNotNull(items); + assertEquals(3, items.size()); + assertEquals("apple", items.get(0)); + assertEquals("banana", items.get(1)); + assertEquals("cherry", items.get(2)); + } + + @Test + public void testJsonNestedObject() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "user": { + "name": "Jane", + "email": "jane@example.com" + } + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("Jane", context.getStringValue("name")); + assertEquals("jane@example.com", context.getStringValue("email")); + } + + @Test + public void testJsonArrayOfObjects() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "users": [ + {"name": "Alice", "age": "25"}, + {"name": "Bob", "age": "30"} + ] + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + final List names = context.getStringValues("name"); + final List ages = context.getStringValues("age"); + assertNotNull(names); + assertNotNull(ages); + assertEquals(2, names.size()); + assertEquals(2, ages.size()); + assertEquals("Alice", names.get(0)); + assertEquals("Bob", names.get(1)); + assertEquals("25", ages.get(0)); + assertEquals("30", ages.get(1)); + } + + @Test + public void testSanitizeMissingClosingQuote() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + // With lenient JSON parsing enabled, this should work + final String jsonResponse = """ + { + "field": "this is not right + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("this is not right", context.getStringValue("field")); + } + + @Test + public void testSanitizeUnescapedQuotesInString() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "field": '"yipii" said the clown' + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + final String value = context.getStringValue("field"); + assertNotNull(value); + assertTrue(value.contains("yipii")); + assertTrue(value.contains("clown")); + } + + @Test + public void testSanitizeUnescapedQuotesInString2() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "field": ""yipii" said the clown" + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + final String value = context.getStringValue("field"); + assertNotNull(value); + assertTrue(value.contains("yipii")); + assertTrue(value.contains("clown")); + } + + @Test + public void testLenientJsonFeatures() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "greeting": 'Hello "World"', + "message": 'She said "hi" to me', + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertNotNull(context.getStringValue("greeting")); + assertNotNull(context.getStringValue("message")); + } + + @Test + public void testJsonWithEscapedCharacters() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "path": "C:\\\\Users\\\\Documents", + "quote": "He said \\"hello\\"" + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertNotNull(context.getStringValue("path")); + assertNotNull(context.getStringValue("quote")); + } + + @Test + public void testInvalidJsonCannotBeSanitized() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + this is completely broken [[[ + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertFalse(result); + + // Verify error message was added to conversation + final List turns = context.getConversation().getConversation(); + final Conversation.Turn lastTurn = turns.get(turns.size() - 1); + assertTrue(lastTurn.text().contains("valid JSON syntax")); + } + + @Test + public void testJsonWithNoCurlyBraces() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = "This is just plain text with no JSON"; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertFalse(result); + } + + @Test + public void testJsonWithAllowedFeatures() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + // Test various lenient JSON features that should be accepted + final String jsonResponse = """ + { + // This is a comment + field1: 'single quotes', + field2: "trailing comma", + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("single quotes", context.getStringValue("field1")); + assertEquals("trailing comma", context.getStringValue("field2")); + } + + @Test + public void testJsonWithNumbers() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = """ + { + "count": 42, + "price": 19.99, + "active": true + } + """; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertEquals("42", context.getStringValue("count")); + assertEquals("19.99", context.getStringValue("price")); + assertEquals("true", context.getStringValue("active")); + } + + @Test + public void testEmptyJson() { + final JsonAgentResponseHandler handler = new JsonAgentResponseHandler(); + final AgentContext context = new AgentContext(); + context.setConversation(Conversation.create("test")); + + final String jsonResponse = "{}"; + + final Completion completion = new Completion(jsonResponse, null, null, null); + final boolean result = handler.handle(completion, context); + + assertTrue(result); + assertTrue(context.getContext().isEmpty()); + } + +}