diff --git a/core/src/main/java/com/google/adk/agents/CallbackPlugin.java b/core/src/main/java/com/google/adk/agents/CallbackPlugin.java index 791e9455..10e260fd 100644 --- a/core/src/main/java/com/google/adk/agents/CallbackPlugin.java +++ b/core/src/main/java/com/google/adk/agents/CallbackPlugin.java @@ -108,7 +108,7 @@ public static class Builder { // Ensures a unique name for each callback. private static final AtomicInteger callbackId = new AtomicInteger(0); - private final ImmutableList.Builder plugins = ImmutableList.builder(); + private ImmutableList.Builder plugins = ImmutableList.builder(); private final ListMultimap, Object> callbacks = ArrayListMultimap.create(); Builder() {} @@ -174,6 +174,20 @@ public Builder addBeforeModelCallbackSync(Callbacks.BeforeModelCallbackSync call Maybe.fromOptional(callback.call(callbackContext, llmRequest))); } + // (b/476510024): Temporary workaround for ces + @CanIgnoreReturnValue + public Builder clearBeforeModelCallbacks() { + callbacks.removeAll(Callbacks.BeforeModelCallback.class); + ImmutableList.Builder updatedPlugins = ImmutableList.builder(); + for (BasePlugin plugin : plugins.build()) { + if (!plugin.getName().startsWith("BeforeModelCallback_")) { + updatedPlugins.add(plugin); + } + } + plugins = updatedPlugins; + return this; + } + @CanIgnoreReturnValue public Builder addAfterModelCallback(Callbacks.AfterModelCallback callback) { callbacks.put(Callbacks.AfterModelCallback.class, callback); diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index 3913b746..fe6763d1 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -274,6 +274,13 @@ public Builder disallowTransferToPeers(boolean disallowTransferToPeers) { return this; } + // (b/476510024): Temporary workaround for ces + @CanIgnoreReturnValue + public Builder clearBeforeModelCallbacks() { + callbackPluginBuilder.clearBeforeModelCallbacks(); + return this; + } + @CanIgnoreReturnValue public Builder beforeModelCallback(BeforeModelCallback beforeModelCallback) { callbackPluginBuilder.addBeforeModelCallback(beforeModelCallback);