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());
|
this.postProcess(split.context());
|
||||||
|
if (split.context().getAgentContextUpdateConsumer() != null) {
|
||||||
|
split.context().getAgentContextUpdateConsumer().accept(split.context());
|
||||||
|
}
|
||||||
|
|
||||||
finishedSplits.add(split);
|
finishedSplits.add(split);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user