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

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

Browse Source
This commit is contained in:
extio-de
2026-03-31 22:36:15 +02:00
parent 00a223bd64
commit 438197403c
5 changed files with 182 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/main/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandler.java
View File

@@ -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<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" +
@@ -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");

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

@@ -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<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)
@@ -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) {

18
src/main/java/de/extio/lmlib/agent/responsehandler/ToolCallingAgentResponseHandler.java
View File

@@ -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);
}

49
src/test/java/de/extio/lmlib/agent/responsehandler/CsvAgentResponseHandlerTest.java
View File

@@ -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));
}
}

49
src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java
View File

@@ -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<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));
}
}