feat: add streaming chunk callback support to generation middleware - #171
Conversation
Add `onChunk` callback to `GenerateParams` and `withStreamCallback` to `ModelParams`, enabling middleware to observe, transform, or replace streaming chunks. This mirrors the JS SDK's `onChunk` middleware hook. - Thread `onChunk` through `GenerateParams` and into `ModelParams` - Add `withStreamCallback` to `ModelParams` for middleware wrapping - Wire streaming callback through `GenerateAction` tool loop - Add `StreamingMiddlewareTest` covering intercept, transform, and multi-middleware streaming scenarios
There was a problem hiding this comment.
Code Review
This pull request refactors the streaming implementation to propagate streaming callbacks through the entire middleware chain, enabling middleware to observe or transform response chunks. Key changes include adding onChunk to GenerateParams, streamCallback to ModelParams, and refactoring Genkit.generateStream to delegate to a unified generateInternal method. Comprehensive tests demonstrate middleware capabilities such as re-chunking and content filtering. Feedback suggests moving the supportsStreaming() validation deeper into the call chain to correctly handle middleware-modified model names and eliminate redundant checks.
| } | ||
|
|
||
| final String spanPath = "/generate/" + modelName; | ||
| final java.util.function.Consumer<ModelResponseChunk> sc = mparams.getStreamCallback(); |
There was a problem hiding this comment.
Adding the supportsStreaming() check here ensures that it is performed on the final resolved model, after any middleware has had a chance to modify the model name. This is more robust than checking it at the entry point of generateStream (lines 1587-1590), which is now redundant and can be removed.
final java.util.function.Consumer<ModelResponseChunk> sc = mparams.getStreamCallback();
if (sc != null && !model.supportsStreaming()) {
throw new GenkitException("Model " + modelName + " does not support streaming");
}| // The callback flows through the full middleware chain: | ||
| // generateInternal → GenerateParams.onChunk → wrapGenerate hooks | ||
| // → rawGenerate → ModelParams.streamCallback → wrapModel hooks → model.run(ctx, req, cb) | ||
| return generateInternal(options, streamCallback); |
There was a problem hiding this comment.
The supportsStreaming() check at the beginning of this method (lines 1587-1590) is now redundant if the validation is moved into the model call chain (see suggestion in buildWrappedModelCall). Removing it also avoids a redundant model lookup and ensures that middleware-modified model names are correctly validated.
Add
onChunkcallback toGenerateParamsandwithStreamCallbacktoModelParams, enabling middleware to observe, transform, or replace streaming chunks. This mirrors the JS SDK'sonChunkmiddleware hook.onChunkthroughGenerateParamsand intoModelParamswithStreamCallbacktoModelParamsfor middleware wrappingGenerateActiontool loopStreamingMiddlewareTestcovering intercept, transform, and multi-middleware streaming scenarios