diff --git a/core/src/main/java/com/google/adk/agents/CallbackContext.java b/core/src/main/java/com/google/adk/agents/CallbackContext.java index f7bbdcdb..808c737a 100644 --- a/core/src/main/java/com/google/adk/agents/CallbackContext.java +++ b/core/src/main/java/com/google/adk/agents/CallbackContext.java @@ -31,6 +31,7 @@ public class CallbackContext extends ReadonlyContext { protected EventActions eventActions; private final State state; + private final String eventId; /** * Initializes callback context. @@ -39,9 +40,22 @@ public class CallbackContext extends ReadonlyContext { * @param eventActions Callback event actions. */ public CallbackContext(InvocationContext invocationContext, EventActions eventActions) { + this(invocationContext, eventActions, null); + } + + /** + * Initializes callback context. + * + * @param invocationContext Current invocation context. + * @param eventActions Callback event actions. + * @param eventId The ID of the event associated with this context. + */ + public CallbackContext( + InvocationContext invocationContext, EventActions eventActions, String eventId) { super(invocationContext); this.eventActions = eventActions != null ? eventActions : EventActions.builder().build(); this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta()); + this.eventId = eventId; } /** Returns the delta-aware state of the current callback. */ @@ -55,6 +69,11 @@ public EventActions eventActions() { return eventActions; } + /** Returns the ID of the event associated with this context. */ + public String eventId() { + return eventId; + } + /** * Lists the filenames of the artifacts attached to the current session. * diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 8e654485..6dfbf586 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -237,7 +237,8 @@ private Flowable callLlm( private Single> handleBeforeModelCallback( InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) { Event callbackEvent = modelResponseEvent.toBuilder().build(); - CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions()); + CallbackContext callbackContext = + new CallbackContext(context, callbackEvent.actions(), callbackEvent.id()); Maybe pluginResult = context.pluginManager().beforeModelCallback(callbackContext, llmRequestBuilder); @@ -274,7 +275,8 @@ private Maybe handleOnModelErrorCallback( Event modelResponseEvent, Throwable throwable) { Event callbackEvent = modelResponseEvent.toBuilder().build(); - CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions()); + CallbackContext callbackContext = + new CallbackContext(context, callbackEvent.actions(), callbackEvent.id()); Exception ex = throwable instanceof Exception e ? e : new Exception(throwable); Maybe pluginResult = @@ -308,7 +310,8 @@ private Maybe handleOnModelErrorCallback( private Single handleAfterModelCallback( InvocationContext context, LlmResponse llmResponse, Event modelResponseEvent) { Event callbackEvent = modelResponseEvent.toBuilder().build(); - CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions()); + CallbackContext callbackContext = + new CallbackContext(context, callbackEvent.actions(), callbackEvent.id()); Maybe pluginResult = context.pluginManager().afterModelCallback(callbackContext, llmResponse); diff --git a/core/src/main/java/com/google/adk/tools/ToolContext.java b/core/src/main/java/com/google/adk/tools/ToolContext.java index b421a8e5..5192d19f 100644 --- a/core/src/main/java/com/google/adk/tools/ToolContext.java +++ b/core/src/main/java/com/google/adk/tools/ToolContext.java @@ -35,8 +35,9 @@ private ToolContext( InvocationContext invocationContext, EventActions eventActions, Optional functionCallId, - Optional toolConfirmation) { - super(invocationContext, eventActions); + Optional toolConfirmation, + @Nullable String eventId) { + super(invocationContext, eventActions, eventId); this.functionCallId = functionCallId; this.toolConfirmation = toolConfirmation; } @@ -125,7 +126,8 @@ public Builder toBuilder() { return new Builder(invocationContext) .actions(eventActions) .functionCallId(functionCallId.orElse(null)) - .toolConfirmation(toolConfirmation.orElse(null)); + .toolConfirmation(toolConfirmation.orElse(null)) + .eventId(eventId()); } @Override @@ -148,6 +150,7 @@ public static final class Builder { private EventActions eventActions = EventActions.builder().build(); // Default empty actions private Optional functionCallId = Optional.empty(); private Optional toolConfirmation = Optional.empty(); + private String eventId; private Builder(InvocationContext invocationContext) { this.invocationContext = invocationContext; @@ -171,8 +174,15 @@ public Builder toolConfirmation(ToolConfirmation toolConfirmation) { return this; } + @CanIgnoreReturnValue + public Builder eventId(String eventId) { + this.eventId = eventId; + return this; + } + public ToolContext build() { - return new ToolContext(invocationContext, eventActions, functionCallId, toolConfirmation); + return new ToolContext( + invocationContext, eventActions, functionCallId, toolConfirmation, eventId); } } }