exo/lm-lib
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Enhance JsonAgentResponseHandler with JSON sanitization and add comprehensive unit tests

Browse Source
This commit is contained in:
extio-de
2025-12-09 21:03:40 +01:00
parent dff1b46228
commit e7c8dc3e2a
2 changed files with 440 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java
View File

@@ -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<String, List<String>>();
@@ -103,12 +112,8 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
private void parseProperty(final LinkedHashMap<String, List<String>> items, final Entry<String, JsonNode> 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();
}
}

369
src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java Normal file
View File

@@ -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<String> 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<String> names = context.getStringValues("name");
final List<String> 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<Conversation.Turn> 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());
}
}