Skip to content

Commit dfd2944

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: Adding OnModelErrorCallback
This Changelist (CL) introduces new error handling mechanisms for LLM model calls and tool invocations within the Agent Development Kit. Specifically, it adds: * **`OnModelErrorCallback`**: An interface for handling errors that occur during an LLM model call. * **`OnToolErrorCallback`**: An interface for handling errors during a tool invocation. These callbacks allow for asynchronous or synchronous handling of errors, providing an opportunity to override the error with a custom response or result. The `LlmAgent` class is updated to support these new callbacks, and they are integrated into `BaseLlmFlow` and `Functions` to be triggered upon model or tool failures, respectively. Plugin-defined error callbacks are given precedence over agent-defined ones. Additionally, the CL introduces `canonical` methods (`canonicalBeforeAgentCallbacks`, `canonicalAfterAgentCallbacks`, etc.) to `BaseAgent` and `LlmAgent` to simplify accessing callback lists, and updates Javadoc for existing callback interfaces. New tests are included to verify the functionality and precedence of the error callbacks. Also, this CL renames many test variables with "unused" prefixes. PiperOrigin-RevId: 859056190
1 parent 6dd51cc commit dfd2944

File tree

9 files changed

+912
-164
lines changed

9 files changed

+912
-164
lines changed

‎core/src/main/java/com/google/adk/agents/BaseAgent.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,24 @@ public Optional<List<? extends AfterAgentCallback>> afterAgentCallback() {
179179
return afterAgentCallback;
180180
}
181181

182+
/**
183+
* The resolved beforeAgentCallback field as a list.
184+
*
185+
* <p>This method is only for use by Agent Development Kit.
186+
*/
187+
public List<? extends BeforeAgentCallback> canonicalBeforeAgentCallbacks() {
188+
return beforeAgentCallback.orElse(ImmutableList.of());
189+
}
190+
191+
/**
192+
* The resolved afterAgentCallback field as a list.
193+
*
194+
* <p>This method is only for use by Agent Development Kit.
195+
*/
196+
public List<? extends AfterAgentCallback> canonicalAfterAgentCallbacks() {
197+
return afterAgentCallback.orElse(ImmutableList.of());
198+
}
199+
182200
/**
183201
* Creates a shallow copy of the parent context with the agent properly being set to this
184202
* instance.

‎core/src/main/java/com/google/adk/agents/CallbackUtil.java‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,12 +54,11 @@ public final class CallbackUtil {
5454
builder.add(beforeAgentCallbackInstance);
5555
} else if (callback instanceof BeforeAgentCallbackSync beforeAgentCallbackSyncInstance) {
5656
builder.add(
57-
(BeforeAgentCallback)
58-
(callbackContext) ->
59-
Maybe.fromOptional(beforeAgentCallbackSyncInstance.call(callbackContext)));
57+
(callbackContext) ->
58+
Maybe.fromOptional(beforeAgentCallbackSyncInstance.call(callbackContext)));
6059
} else {
6160
logger.warn(
62-
"Invalid beforeAgentCallback callback type: %s. Ignoring this callback.",
61+
"Invalid beforeAgentCallback callback type: {}. Ignoring this callback.",
6362
callback.getClass().getName());
6463
}
6564
}
@@ -87,12 +86,11 @@ public final class CallbackUtil {
8786
builder.add(afterAgentCallbackInstance);
8887
} else if (callback instanceof AfterAgentCallbackSync afterAgentCallbackSyncInstance) {
8988
builder.add(
90-
(AfterAgentCallback)
91-
(callbackContext) ->
92-
Maybe.fromOptional(afterAgentCallbackSyncInstance.call(callbackContext)));
89+
(callbackContext) ->
90+
Maybe.fromOptional(afterAgentCallbackSyncInstance.call(callbackContext)));
9391
} else {
9492
logger.warn(
95-
"Invalid afterAgentCallback callback type: %s. Ignoring this callback.",
93+
"Invalid afterAgentCallback callback type: {}. Ignoring this callback.",
9694
callback.getClass().getName());
9795
}
9896
}

‎core/src/main/java/com/google/adk/agents/Callbacks.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,8 +75,36 @@ public interface AfterModelCallbackSync extends AfterModelCallbackBase {
7575
Optional<LlmResponse> call(CallbackContext callbackContext, LlmResponse llmResponse);
7676
}
7777

78+
interface OnModelErrorCallbackBase {}
79+
80+
/** Async callback interface for handling errors that occur during an LLM model call. */
81+
@FunctionalInterface
82+
public interface OnModelErrorCallback extends OnModelErrorCallbackBase {
83+
/**
84+
* Async callback when model call fails.
85+
*
86+
* @param callbackContext Callback context.
87+
* @param llmRequest LLM request.
88+
* @param error The exception that occurred.
89+
* @return response override, or empty to continue with error.
90+
*/
91+
Maybe<LlmResponse> call(
92+
CallbackContext callbackContext, LlmRequest llmRequest, Exception error);
93+
}
94+
95+
/**
96+
* Helper interface to allow for sync onModelErrorCallback. The function is wrapped into an async
97+
* one before being processed further.
98+
*/
99+
@FunctionalInterface
100+
public interface OnModelErrorCallbackSync extends OnModelErrorCallbackBase {
101+
Optional<LlmResponse> call(
102+
CallbackContext callbackContext, LlmRequest llmRequest, Exception error);
103+
}
104+
78105
interface BeforeAgentCallbackBase {}
79106

107+
/** Async callback interface for actions to be performed before an agent starts running. */
80108
@FunctionalInterface
81109
public interface BeforeAgentCallback extends BeforeAgentCallbackBase {
82110
/**
@@ -99,6 +127,7 @@ public interface BeforeAgentCallbackSync extends BeforeAgentCallbackBase {
99127

100128
interface AfterAgentCallbackBase {}
101129

130+
/** Async callback interface for actions to be performed after an agent has finished running. */
102131
@FunctionalInterface
103132
public interface AfterAgentCallback extends AfterAgentCallbackBase {
104133
/**
@@ -121,6 +150,7 @@ public interface AfterAgentCallbackSync extends AfterAgentCallbackBase {
121150

122151
interface BeforeToolCallbackBase {}
123152

153+
/** Async callback interface for actions to be performed before a tool is invoked. */
124154
@FunctionalInterface
125155
public interface BeforeToolCallback extends BeforeToolCallbackBase {
126156
/**
@@ -154,6 +184,7 @@ Optional<Map<String, Object>> call(
154184

155185
interface AfterToolCallbackBase {}
156186

187+
/** Async callback interface for actions to be performed after a tool has been invoked. */
157188
@FunctionalInterface
158189
public interface AfterToolCallback extends AfterToolCallbackBase {
159190
/**
@@ -188,5 +219,42 @@ Optional<Map<String, Object>> call(
188219
Object response);
189220
}
190221

222+
interface OnToolErrorCallbackBase {}
223+
224+
/** Async callback interface for handling errors that occur during a tool invocation. */
225+
@FunctionalInterface
226+
public interface OnToolErrorCallback extends OnToolErrorCallbackBase {
227+
/**
228+
* Async callback when tool call fails.
229+
*
230+
* @param invocationContext Invocation context.
231+
* @param baseTool Tool instance.
232+
* @param input Tool input arguments.
233+
* @param toolContext Tool context.
234+
* @param error The exception that occurred.
235+
* @return override result, or empty to continue with error.
236+
*/
237+
Maybe<Map<String, Object>> call(
238+
InvocationContext invocationContext,
239+
BaseTool baseTool,
240+
Map<String, Object> input,
241+
ToolContext toolContext,
242+
Exception error);
243+
}
244+
245+
/**
246+
* Helper interface to allow for sync onToolErrorCallback. The function is wrapped into an async
247+
* one before being processed further.
248+
*/
249+
@FunctionalInterface
250+
public interface OnToolErrorCallbackSync extends OnToolErrorCallbackBase {
251+
Optional<Map<String, Object>> call(
252+
InvocationContext invocationContext,
253+
BaseTool baseTool,
254+
Map<String, Object> input,
255+
ToolContext toolContext,
256+
Exception error);
257+
}
258+
191259
private Callbacks() {}
192260
}

0 commit comments

Comments
 (0)