feat: Add support for configuring agent callbacks in YAML

Doris26 · copybara-github · commit 27c01724d6e9 · 2025-09-09T16:03:22.000-07:00
This change enables developers to define various lifecycle callbacks for `LlmAgent` directly within the agent's YAML configuration file. It introduces new sections like `before_agent_callbacks`, `after_agent_callbacks`, `before_model_callbacks`, `after_model_callbacks`, `before_tool_callbacks`, and `after_tool_callbacks` in the YAML schema. The `LlmAgentConfig` and `LlmAgent` classes are updated to parse these sections, and the `ComponentRegistry` is enhanced to resolve the named callback implementations provided in the application. This allows for more flexible and declarative configuration of agent behavior.

PiperOrigin-RevId: 805098543
diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java
@@ -976,6 +976,9 @@ public static LlmAgent fromConfig(LlmAgentConfig config, String configAbsPath)
       builder.generateContentConfig(config.generateContentConfig());
     }
 
+    // Resolve callbacks if configured
+    setCallbacksFromConfig(config, builder);
+
     // Build and return the agent
     LlmAgent agent = builder.build();
     logger.info(
@@ -986,6 +989,93 @@ public static LlmAgent fromConfig(LlmAgentConfig config, String configAbsPath)
     return agent;
   }
 
+  private static void setCallbacksFromConfig(LlmAgentConfig config, Builder builder)
+      throws ConfigurationException {
+    var beforeAgentCallbacks = config.beforeAgentCallbacks();
+    if (beforeAgentCallbacks != null) {
+      ImmutableList.Builder<Callbacks.BeforeAgentCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : beforeAgentCallbacks) {
+        var callback = ComponentRegistry.resolveBeforeAgentCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid before_agent_callback: " + ref.name());
+      }
+      builder.beforeAgentCallback(list.build());
+    }
+
+    var afterAgentCallbacks = config.afterAgentCallbacks();
+    if (afterAgentCallbacks != null) {
+      ImmutableList.Builder<Callbacks.AfterAgentCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : afterAgentCallbacks) {
+        var callback = ComponentRegistry.resolveAfterAgentCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid after_agent_callback: " + ref.name());
+      }
+      builder.afterAgentCallback(list.build());
+    }
+
+    var beforeModelCallbacks = config.beforeModelCallbacks();
+    if (beforeModelCallbacks != null) {
+      ImmutableList.Builder<Callbacks.BeforeModelCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : beforeModelCallbacks) {
+        var callback = ComponentRegistry.resolveBeforeModelCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid before_model_callback: " + ref.name());
+      }
+      builder.beforeModelCallback(list.build());
+    }
+
+    var afterModelCallbacks = config.afterModelCallbacks();
+    if (afterModelCallbacks != null) {
+      ImmutableList.Builder<Callbacks.AfterModelCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : afterModelCallbacks) {
+        var callback = ComponentRegistry.resolveAfterModelCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid after_model_callback: " + ref.name());
+      }
+      builder.afterModelCallback(list.build());
+    }
+
+    var beforeToolCallbacks = config.beforeToolCallbacks();
+    if (beforeToolCallbacks != null) {
+      ImmutableList.Builder<Callbacks.BeforeToolCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : beforeToolCallbacks) {
+        var callback = ComponentRegistry.resolveBeforeToolCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid before_tool_callback: " + ref.name());
+      }
+      builder.beforeToolCallback(list.build());
+    }
+
+    var afterToolCallbacks = config.afterToolCallbacks();
+    if (afterToolCallbacks != null) {
+      ImmutableList.Builder<Callbacks.AfterToolCallbackBase> list = ImmutableList.builder();
+      for (LlmAgentConfig.CallbackRef ref : afterToolCallbacks) {
+        var callback = ComponentRegistry.resolveAfterToolCallback(ref.name());
+        if (callback.isPresent()) {
+          list.add(callback.get());
+          continue;
+        }
+        throw new ConfigurationException("Invalid after_tool_callback: " + ref.name());
+      }
+      builder.afterToolCallback(list.build());
+    }
+  }
+
   /**
    * Resolves a list of tool configurations into both {@link BaseTool} and {@link BaseToolset}
    * instances.
diff --git a/core/src/main/java/com/google/adk/agents/LlmAgentConfig.java b/core/src/main/java/com/google/adk/agents/LlmAgentConfig.java
@@ -36,6 +36,33 @@ public class LlmAgentConfig extends BaseAgentConfig {
   private IncludeContents includeContents;
   private GenerateContentConfig generateContentConfig;
 
+  // Callback configuration (names resolved via ComponentRegistry)
+  private List<CallbackRef> beforeAgentCallbacks;
+  private List<CallbackRef> afterAgentCallbacks;
+  private List<CallbackRef> beforeModelCallbacks;
+  private List<CallbackRef> afterModelCallbacks;
+  private List<CallbackRef> beforeToolCallbacks;
+  private List<CallbackRef> afterToolCallbacks;
+
+  /** Reference to a callback stored in the ComponentRegistry. */
+  public static class CallbackRef {
+    private String name;
+
+    public CallbackRef() {}
+
+    public CallbackRef(String name) {
+      this.name = name;
+    }
+
+    public String name() {
+      return name;
+    }
+
+    public void setName(String name) {
+      this.name = name;
+    }
+  }
+
   public LlmAgentConfig() {
     super();
     setAgentClass("LlmAgent");
@@ -105,4 +132,52 @@ public GenerateContentConfig generateContentConfig() {
   public void setGenerateContentConfig(GenerateContentConfig generateContentConfig) {
     this.generateContentConfig = generateContentConfig;
   }
+
+  public List<CallbackRef> beforeAgentCallbacks() {
+    return beforeAgentCallbacks;
+  }
+
+  public void setBeforeAgentCallbacks(List<CallbackRef> beforeAgentCallbacks) {
+    this.beforeAgentCallbacks = beforeAgentCallbacks;
+  }
+
+  public List<CallbackRef> afterAgentCallbacks() {
+    return afterAgentCallbacks;
+  }
+
+  public void setAfterAgentCallbacks(List<CallbackRef> afterAgentCallbacks) {
+    this.afterAgentCallbacks = afterAgentCallbacks;
+  }
+
+  public List<CallbackRef> beforeModelCallbacks() {
+    return beforeModelCallbacks;
+  }
+
+  public void setBeforeModelCallbacks(List<CallbackRef> beforeModelCallbacks) {
+    this.beforeModelCallbacks = beforeModelCallbacks;
+  }
+
+  public List<CallbackRef> afterModelCallbacks() {
+    return afterModelCallbacks;
+  }
+
+  public void setAfterModelCallbacks(List<CallbackRef> afterModelCallbacks) {
+    this.afterModelCallbacks = afterModelCallbacks;
+  }
+
+  public List<CallbackRef> beforeToolCallbacks() {
+    return beforeToolCallbacks;
+  }
+
+  public void setBeforeToolCallbacks(List<CallbackRef> beforeToolCallbacks) {
+    this.beforeToolCallbacks = beforeToolCallbacks;
+  }
+
+  public List<CallbackRef> afterToolCallbacks() {
+    return afterToolCallbacks;
+  }
+
+  public void setAfterToolCallbacks(List<CallbackRef> afterToolCallbacks) {
+    this.afterToolCallbacks = afterToolCallbacks;
+  }
 }
diff --git a/core/src/main/java/com/google/adk/utils/ComponentRegistry.java b/core/src/main/java/com/google/adk/utils/ComponentRegistry.java
@@ -20,6 +20,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.adk.agents.BaseAgent;
+import com.google.adk.agents.Callbacks;
 import com.google.adk.agents.LlmAgent;
 import com.google.adk.agents.LoopAgent;
 import com.google.adk.agents.ParallelAgent;
@@ -153,7 +154,7 @@ private void registerAdkToolsetClass(@Nonnull Class<? extends BaseToolset> tools
    * @throws IllegalArgumentException if name is null or empty, or if value is null
    */
   public void register(String name, Object value) {
-    if (name == null || name.trim().isEmpty()) {
+    if (isNullOrEmpty(name) || name.trim().isEmpty()) {
       throw new IllegalArgumentException("Name cannot be null or empty");
     }
     if (value == null) {
@@ -461,4 +462,36 @@ public Set<String> getToolNamesWithPrefix(String prefix) {
         .filter(name -> name.startsWith(prefix))
         .collect(toImmutableSet());
   }
+
+  public static Optional<Callbacks.BeforeAgentCallback> resolveBeforeAgentCallback(String name) {
+    return resolveCallback(name, Callbacks.BeforeAgentCallback.class);
+  }
+
+  public static Optional<Callbacks.AfterAgentCallback> resolveAfterAgentCallback(String name) {
+    return resolveCallback(name, Callbacks.AfterAgentCallback.class);
+  }
+
+  public static Optional<Callbacks.BeforeModelCallback> resolveBeforeModelCallback(String name) {
+    return resolveCallback(name, Callbacks.BeforeModelCallback.class);
+  }
+
+  public static Optional<Callbacks.AfterModelCallback> resolveAfterModelCallback(String name) {
+    return resolveCallback(name, Callbacks.AfterModelCallback.class);
+  }
+
+  public static Optional<Callbacks.BeforeToolCallback> resolveBeforeToolCallback(String name) {
+    return resolveCallback(name, Callbacks.BeforeToolCallback.class);
+  }
+
+  public static Optional<Callbacks.AfterToolCallback> resolveAfterToolCallback(String name) {
+    return resolveCallback(name, Callbacks.AfterToolCallback.class);
+  }
+
+  private static <T> Optional<T> resolveCallback(String name, Class<T> type) {
+    if (isNullOrEmpty(name)) {
+      return Optional.empty();
+    }
+    ComponentRegistry registry = getInstance();
+    return registry.get(name, type);
+  }
 }
diff --git a/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java b/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java
