From 7cf067400b79eb93d2e2de48f259d85e425dc616 Mon Sep 17 00:00:00 2001 From: extio-de Date: Wed, 2 Jul 2025 21:45:44 +0200 Subject: [PATCH] feat: Allow AgentContext update consumer in BaseAgent and add beforeStream to TextAgentResponseHandler This commit introduces the following changes: - **BaseAgent:** Added a mechanism to allow consumers to update the AgentContext after post-processing a split. This provides more flexibility for downstream components to react to the processed data. - **TextAgentResponseHandler:** Introduced a `beforeStream` consumer that allows updating the AgentContext before the stream begins. This enables pre-processing or initialization steps within the response handler. The constructor has been updated to accept this new consumer. Also fixed a bug where the chunkKey and updateKey were not being removed after processing. These changes enhance the extensibility and configurability of the agent processing pipeline. --- .../java/de/extio/lmlib/agent/BaseAgent.java | 3 +++ .../TextAgentResponseHandler.java | 16 ++++++++++++---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/main/java/de/extio/lmlib/agent/BaseAgent.java b/src/main/java/de/extio/lmlib/agent/BaseAgent.java index 86e0fd7..066a006 100644 --- a/src/main/java/de/extio/lmlib/agent/BaseAgent.java +++ b/src/main/java/de/extio/lmlib/agent/BaseAgent.java @@ -142,6 +142,9 @@ public interface BaseAgent { } this.postProcess(split.context()); + if (split.context().getAgentContextUpdateConsumer() != null) { + split.context().getAgentContextUpdateConsumer().accept(split.context()); + } finishedSplits.add(split); } diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java index 18c108b..f878132 100755 --- a/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java @@ -22,18 +22,21 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler { private final Function transformer; + private final Consumer beforeStream; + private final Function chunkTransformer; private final Consumer afterChunkUpdate; public TextAgentResponseHandler(final String key) { - this(key, null, null, null); + this(key, null, null, null, null); } - public TextAgentResponseHandler(final String key, final Function transformer, final Function chunkTransformer, final Consumer afterChunkUpdate) { + public TextAgentResponseHandler(final String key, final Function transformer, final Consumer beforeStream, final Function chunkTransformer, final Consumer afterChunkUpdate) { this.key = key; this.chunkKey = key.concat("_chunk"); this.transformer = transformer; + this.beforeStream = beforeStream; this.chunkTransformer = chunkTransformer; this.afterChunkUpdate = afterChunkUpdate; } @@ -52,6 +55,9 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler { @Override public void beforeStream(final AgentContext context) { context.setStringValue(this.key, ""); + if (this.beforeStream != null) { + this.beforeStream.accept(context); + } } @Override @@ -60,11 +66,13 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler { if (this.chunkTransformer != null) { chunk = this.chunkTransformer.apply(chunk); } - context.setStringValue(this.chunkKey, chunk); - context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "").concat(chunk)); context.setStringValue(UPDATE_KEY, this.key); if (afterChunkUpdate != null) { + context.setStringValue(this.chunkKey, chunk); + context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "").concat(chunk)); afterChunkUpdate.accept(context); + context.getContext().remove(UPDATE_KEY); + context.getContext().remove(this.chunkKey); } return true; }