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.
This commit is contained in:
extio-de
2025-07-02 21:45:44 +02:00
parent c5e2ace604
commit 7cf067400b
2 changed files with 15 additions and 4 deletions

View File

@@ -142,6 +142,9 @@ public interface BaseAgent {
} }
this.postProcess(split.context()); this.postProcess(split.context());
if (split.context().getAgentContextUpdateConsumer() != null) {
split.context().getAgentContextUpdateConsumer().accept(split.context());
}
finishedSplits.add(split); finishedSplits.add(split);
} }

View File

@@ -22,18 +22,21 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
private final Function<String, String> transformer; private final Function<String, String> transformer;
private final Consumer<AgentContext> beforeStream;
private final Function<String, String> chunkTransformer; private final Function<String, String> chunkTransformer;
private final Consumer<AgentContext> afterChunkUpdate; private final Consumer<AgentContext> afterChunkUpdate;
public TextAgentResponseHandler(final String key) { public TextAgentResponseHandler(final String key) {
this(key, null, null, null); this(key, null, null, null, null);
} }
public TextAgentResponseHandler(final String key, final Function<String, String> transformer, final Function<String, String> chunkTransformer, final Consumer<AgentContext> afterChunkUpdate) { public TextAgentResponseHandler(final String key, final Function<String, String> transformer, final Consumer<AgentContext> beforeStream, final Function<String, String> chunkTransformer, final Consumer<AgentContext> afterChunkUpdate) {
this.key = key; this.key = key;
this.chunkKey = key.concat("_chunk"); this.chunkKey = key.concat("_chunk");
this.transformer = transformer; this.transformer = transformer;
this.beforeStream = beforeStream;
this.chunkTransformer = chunkTransformer; this.chunkTransformer = chunkTransformer;
this.afterChunkUpdate = afterChunkUpdate; this.afterChunkUpdate = afterChunkUpdate;
} }
@@ -52,6 +55,9 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
@Override @Override
public void beforeStream(final AgentContext context) { public void beforeStream(final AgentContext context) {
context.setStringValue(this.key, ""); context.setStringValue(this.key, "");
if (this.beforeStream != null) {
this.beforeStream.accept(context);
}
} }
@Override @Override
@@ -60,11 +66,13 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
if (this.chunkTransformer != null) { if (this.chunkTransformer != null) {
chunk = this.chunkTransformer.apply(chunk); 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); context.setStringValue(UPDATE_KEY, this.key);
if (afterChunkUpdate != null) { if (afterChunkUpdate != null) {
context.setStringValue(this.chunkKey, chunk);
context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "").concat(chunk));
afterChunkUpdate.accept(context); afterChunkUpdate.accept(context);
context.getContext().remove(UPDATE_KEY);
context.getContext().remove(this.chunkKey);
} }
return true; return true;
} }