feat: Implement reasoning streaming and chunk handling in Csv and Json agent response handlers

src/main/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandler.java

@@ -2,17 +2,20 @@ package de.extio.lmlib.agent.responsehandler;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.regex.Pattern;
+import java.util.function.BiConsumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Chunk;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
 import de.extio.lmlib.util.TextUtils;
 
-public class CsvAgentResponseHandler implements AgentResponseHandler {
+public class CsvAgentResponseHandler implements StreamedAgentResponseHandler {
 	
 	private static final Logger LOGGER = LoggerFactory.getLogger(CsvAgentResponseHandler.class);
 	
@@ -26,11 +29,19 @@ public class CsvAgentResponseHandler implements AgentResponseHandler {
 	
 	private final String errorPromptSuffix;
 
+	private final String reasoningKey;
+
+	private final BiConsumer<AgentContext, Chunk> afterChunkUpdate;
+	
 	public CsvAgentResponseHandler(final List<String> headings, final String keyPrefix) {
 		this(headings, keyPrefix, ";");
 	}
 	
 	public CsvAgentResponseHandler(final List<String> headings, final String keyPrefix, final String fieldSeparator) {
+		this(headings, keyPrefix, fieldSeparator, null);
+	}
+
+	public CsvAgentResponseHandler(final List<String> headings, final String keyPrefix, final String fieldSeparator, final BiConsumer<AgentContext, Chunk> afterChunkUpdate) {
 		if (headings == null || headings.isEmpty()) {
 			throw new IllegalArgumentException("Headings list cannot be null or empty");
 		}
@@ -38,6 +49,8 @@ public class CsvAgentResponseHandler implements AgentResponseHandler {
 		this.keyPrefix = keyPrefix != null ? keyPrefix : "";
 		this.fieldSeparator = fieldSeparator;
 		this.fieldSeparatorQuoted = Pattern.quote(fieldSeparator);
+		this.reasoningKey = this.keyPrefix + "_reasoning";
+		this.afterChunkUpdate = afterChunkUpdate;
 		
 		this.errorPromptSuffix = "\n\nThe previous response could not be fully processed or validated. " +
 				"Please format the response as CSV with '" + fieldSeparator + "' as field separator.\n" +
@@ -69,6 +82,25 @@ public class CsvAgentResponseHandler implements AgentResponseHandler {
 		}
 	}
 
+	@Override
+	public void beforeStream(final AgentContext context) {
+		context.getContext().remove(this.reasoningKey);
+	}
+
+	@Override
+	public boolean handleChunk(final Chunk chunk, final AgentContext context) {
+		boolean handled = false;
+		if (chunk.reasoningContent() != null) {
+			context.setStringValue(this.reasoningKey, Objects.requireNonNullElseGet(context.getStringValue(this.reasoningKey), () -> "").concat(chunk.reasoningContent()));
+			handled = true;
+		}
+		if (this.afterChunkUpdate != null) {
+			this.afterChunkUpdate.accept(context, chunk);
+			return true;
+		}
+		return handled;
+	}
+	
 	private boolean parseCsv(final String csvText, final AgentContext context) {
 		final String[] lines = TextUtils.normalizeModelResponse(csvText.replace("\\r", "\r").replace("\\n", "\n"), false).split("\n");
 		

src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java

@@ -6,6 +6,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.function.BiConsumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -17,10 +18,11 @@ import tools.jackson.databind.ObjectMapper;
 import tools.jackson.databind.json.JsonMapper;
 
 import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Chunk;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
 
-public class JsonAgentResponseHandler implements AgentResponseHandler {
+public class JsonAgentResponseHandler implements StreamedAgentResponseHandler {
 	
 	private static final Logger LOGGER = LoggerFactory.getLogger(JsonAgentResponseHandler.class);
 	
@@ -28,12 +30,22 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
 	
 	private final String prefixFields;
 
+	private final String reasoningKey;
+
+	private final BiConsumer<AgentContext, Chunk> afterChunkUpdate;
+	
 	public JsonAgentResponseHandler() {
-		this(null);
+		this(null, null);
 	}
 	
 	public JsonAgentResponseHandler(final String prefixFields) {
+		this(prefixFields, null);
+	}
+
+	public JsonAgentResponseHandler(final String prefixFields, final BiConsumer<AgentContext, Chunk> afterChunkUpdate) {
 		this.prefixFields = prefixFields == null ? "" : prefixFields;
+		this.reasoningKey = this.prefixFields + "_reasoning";
+		this.afterChunkUpdate = afterChunkUpdate;
 		
 		this.objectMapper = JsonMapper.builder()
 				.configure(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES, false)
@@ -56,6 +68,26 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
 				.build();
 	}
 
+	@Override
+	public void beforeStream(final AgentContext context) {
+		context.getContext().remove(this.reasoningKey);
+	}
+
+	@Override
+	public boolean handleChunk(final Chunk chunk, final AgentContext context) {
+		boolean handled = false;
+		if (chunk.reasoningContent() != null) {
+			final var existingReasoning = context.getStringValue(this.reasoningKey);
+			context.setStringValue(this.reasoningKey, (existingReasoning == null ? "" : existingReasoning).concat(chunk.reasoningContent()));
+			handled = true;
+		}
+		if (this.afterChunkUpdate != null) {
+			this.afterChunkUpdate.accept(context, chunk);
+			return true;
+		}
+		return handled;
+	}
+	
 	@Override
 	public boolean handle(final Completion completion, final AgentContext context) {
 		final int firstCurly = completion.response().indexOf("{");

src/main/java/de/extio/lmlib/agent/responsehandler/ToolCallingAgentResponseHandler.java

@@ -4,11 +4,12 @@ import java.util.ArrayList;
 import java.util.List;
 
 import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Chunk;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.ToolCall;
 import de.extio.lmlib.client.ToolParameters;
 
-public class ToolCallingAgentResponseHandler implements AgentResponseHandler {
+public class ToolCallingAgentResponseHandler implements StreamedAgentResponseHandler {
 
 	@FunctionalInterface
 	public interface ToolCallHandler {
@@ -58,6 +59,21 @@ public class ToolCallingAgentResponseHandler implements AgentResponseHandler {
 		return this.delegate == null || this.delegate.handle(completion, context);
 	}
 
+	@Override
+	public void beforeStream(final AgentContext context) {
+		if (this.delegate instanceof final StreamedAgentResponseHandler streamedAgentResponseHandler) {
+			streamedAgentResponseHandler.beforeStream(context);
+		}
+	}
+
+	@Override
+	public boolean handleChunk(final Chunk chunk, final AgentContext context) {
+		if (this.delegate instanceof final StreamedAgentResponseHandler streamedAgentResponseHandler) {
+			return streamedAgentResponseHandler.handleChunk(chunk, context);
+		}
+		return true;
+	}
+
 	public boolean handleToolCalls(final Completion completion, final AgentContext context, final ToolCallResults toolCallResults) {
 		return this.toolCallHandler.handle(completion, context, toolCallResults);
 	}

src/test/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandlerTest.java

@@ -6,11 +6,13 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import org.junit.jupiter.api.Test;
 
 import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Chunk;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
 
@@ -319,4 +321,51 @@ public class CsvAgentResponseHandlerTest {
 			new CsvAgentResponseHandler(List.of(), "prefix_");
 		});
 	}
+
+	@Test
+	public void testReasoningStreaming() {
+		final CsvAgentResponseHandler handler = new CsvAgentResponseHandler(List.of("a", "b"), "stream");
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		context.setStringValue("stream_reasoning", "stale");
+
+		handler.beforeStream(context);
+		assertEquals(null, context.getStringValue("stream_reasoning"));
+
+		assertTrue(handler.handleChunk(new Chunk(null, "first"), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+
+		assertFalse(handler.handleChunk(new Chunk("1;2", null), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+
+		assertTrue(handler.handleChunk(new Chunk(null, " second"), context));
+		assertEquals("first second", context.getStringValue("stream_reasoning"));
+	}
+
+	@Test
+	public void testAfterChunkUpdateCallback() {
+		final List<String> callbackReasoning = new ArrayList<>();
+		final List<String> callbackContent = new ArrayList<>();
+		final CsvAgentResponseHandler handler = new CsvAgentResponseHandler(List.of("a", "b"), "stream", ";",
+				(context, chunk) -> {
+					callbackReasoning.add(context.getStringValue("stream_reasoning"));
+					callbackContent.add(chunk.content());
+				});
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+
+		handler.beforeStream(context);
+
+		assertTrue(handler.handleChunk(new Chunk(null, "first"), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+		assertEquals(1, callbackReasoning.size());
+		assertEquals("first", callbackReasoning.getFirst());
+		assertEquals(null, callbackContent.getFirst());
+
+		assertTrue(handler.handleChunk(new Chunk("1;2", null), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+		assertEquals(2, callbackReasoning.size());
+		assertEquals("first", callbackReasoning.get(1));
+		assertEquals("1;2", callbackContent.get(1));
+	}
 }

src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java

@@ -5,11 +5,13 @@ 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.ArrayList;
 import java.util.List;
 
 import org.junit.jupiter.api.Test;
 
 import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Chunk;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
 
@@ -366,4 +368,51 @@ public class JsonAgentResponseHandlerTest {
 		assertTrue(context.getContext().isEmpty());
 	}
 
+	@Test
+	public void testReasoningStreaming() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler("stream");
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		context.setStringValue("stream_reasoning", "stale");
+
+		handler.beforeStream(context);
+		assertEquals(null, context.getStringValue("stream_reasoning"));
+
+		assertTrue(handler.handleChunk(new Chunk(null, "first"), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+
+		assertFalse(handler.handleChunk(new Chunk("{\"field\":\"value\"}", null), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+
+		assertTrue(handler.handleChunk(new Chunk(null, " second"), context));
+		assertEquals("first second", context.getStringValue("stream_reasoning"));
+	}
+
+	@Test
+	public void testAfterChunkUpdateCallback() {
+		final List<String> callbackReasoning = new ArrayList<>();
+		final List<String> callbackContent = new ArrayList<>();
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler("stream",
+				(context, chunk) -> {
+					callbackReasoning.add(context.getStringValue("stream_reasoning"));
+					callbackContent.add(chunk.content());
+				});
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+
+		handler.beforeStream(context);
+
+		assertTrue(handler.handleChunk(new Chunk(null, "first"), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+		assertEquals(1, callbackReasoning.size());
+		assertEquals("first", callbackReasoning.getFirst());
+		assertEquals(null, callbackContent.getFirst());
+
+		assertTrue(handler.handleChunk(new Chunk("{\"field\":\"value\"}", null), context));
+		assertEquals("first", context.getStringValue("stream_reasoning"));
+		assertEquals(2, callbackReasoning.size());
+		assertEquals("first", callbackReasoning.get(1));
+		assertEquals("{\"field\":\"value\"}", callbackContent.get(1));
+	}
+	
 }