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:
@@ -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);
|
||||
}
|
||||
|
||||
@@ -22,18 +22,21 @@ public class TextAgentResponseHandler implements StreamedAgentResponseHandler {
|
||||
|
||||
private final Function<String, String> transformer;
|
||||
|
||||
private final Consumer<AgentContext> beforeStream;
|
||||
|
||||
private final Function<String, String> chunkTransformer;
|
||||
|
||||
private final Consumer<AgentContext> afterChunkUpdate;
|
||||
|
||||
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.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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user