diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandler.java index e756d56..30620b5 100644 --- a/src/main/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandler.java +++ b/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); @@ -25,12 +28,20 @@ public class CsvAgentResponseHandler implements AgentResponseHandler { private final String fieldSeparatorQuoted; private final String errorPromptSuffix; + + private final String reasoningKey; + + private final BiConsumer afterChunkUpdate; public CsvAgentResponseHandler(final List headings, final String keyPrefix) { this(headings, keyPrefix, ";"); } public CsvAgentResponseHandler(final List headings, final String keyPrefix, final String fieldSeparator) { + this(headings, keyPrefix, fieldSeparator, null); + } + + public CsvAgentResponseHandler(final List headings, final String keyPrefix, final String fieldSeparator, final BiConsumer 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" + @@ -68,6 +81,25 @@ public class CsvAgentResponseHandler implements AgentResponseHandler { return false; } } + + @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"); 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 7177424..2d30bbb 100755 --- a/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java +++ b/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,23 +18,34 @@ 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); private final ObjectMapper objectMapper; private final String prefixFields; + + private final String reasoningKey; + + private final BiConsumer afterChunkUpdate; public JsonAgentResponseHandler() { - this(null); + this(null, null); } public JsonAgentResponseHandler(final String prefixFields) { + this(prefixFields, null); + } + + public JsonAgentResponseHandler(final String prefixFields, final BiConsumer 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) @@ -55,6 +67,26 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { .enable(JsonReadFeature.ALLOW_UNQUOTED_PROPERTY_NAMES) .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) { diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/ToolCallingAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/ToolCallingAgentResponseHandler.java index 285f0c3..ad5c626 100644 --- a/src/main/java/de/extio/lmlib/agent/responsehandler/ToolCallingAgentResponseHandler.java +++ b/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); } diff --git a/src/test/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandlerTest.java b/src/test/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandlerTest.java index c8e392a..73151ad 100644 --- a/src/test/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandlerTest.java +++ b/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 callbackReasoning = new ArrayList<>(); + final List 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)); + } } diff --git a/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java b/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java index 50ad7c2..ac6a9ac 100644 --- a/src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java +++ b/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; @@ -365,5 +367,52 @@ public class JsonAgentResponseHandlerTest { assertTrue(result); 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 callbackReasoning = new ArrayList<>(); + final List 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)); + } }