diff --git a/core/src/main/java/com/google/adk/agents/Callbacks.java b/core/src/main/java/com/google/adk/agents/Callbacks.java index 1c26897b..c37a9533 100644 --- a/core/src/main/java/com/google/adk/agents/Callbacks.java +++ b/core/src/main/java/com/google/adk/agents/Callbacks.java @@ -66,6 +66,33 @@ public interface AfterModelCallback extends AfterModelCallbackBase { Maybe call(CallbackContext callbackContext, LlmResponse llmResponse); } + /** An extended callback interface for streaming LLM responses. */ + public interface StreamingAfterModelCallback extends AfterModelCallback { + /** + * Invoked repeatedly for each significant chunk during a live stream. + * + * @param callbackContext Callback context. + * @param llmResponse LLM response chunk. + * @return modified response chunk, or empty to keep original. + */ + Maybe callInStream(CallbackContext callbackContext, LlmResponse llmResponse); + + /** + * Invoked once at turn completion. + * + * @param callbackContext Callback context. + * @param finalResponse The final, aggregated LLM response. + * @return modified final response, or empty to keep original. + */ + Maybe callAtStreamEnd(CallbackContext callbackContext, LlmResponse finalResponse); + + /** Defaults to the end-of-stream behavior. */ + @Override + default Maybe call(CallbackContext callbackContext, LlmResponse llmResponse) { + return callAtStreamEnd(callbackContext, llmResponse); + } + } + /** * Helper interface to allow for sync afterModelCallback. The function is wrapped into an async * one before being processed further.