Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 26 additions & 38 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,21 @@

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.base.Strings.nullToEmpty;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.lang.String.format;

import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.annotations.Experimental;
import com.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.adk.plugins.Plugin;
import com.google.adk.telemetry.Instrumentation;
import com.google.adk.telemetry.Instrumentation.AgentInvocation;
import com.google.adk.telemetry.Tracing;
import com.google.adk.utils.AgentEnums.AgentOrigin;
import com.google.adk.workflow.BaseNode;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.errorprone.annotations.DoNotCall;
Expand All @@ -49,21 +52,12 @@
import org.jspecify.annotations.Nullable;

/** Base class for all agents. */
public abstract class BaseAgent {
public abstract class BaseAgent extends BaseNode {

// Pattern for valid agent names.
private static final String IDENTIFIER_REGEX = "^_?[a-zA-Z0-9]*([. _-][a-zA-Z0-9]+)*$";
private static final Pattern IDENTIFIER_PATTERN = Pattern.compile(IDENTIFIER_REGEX);

/** The agent's name. Must be a unique identifier within the agent tree. */
private final String name;

/**
* One line description about the agent's capability. The system can use this for decision-making
* when delegating control to different agents.
*/
private final String description;

/**
* The parent agent in the agent tree. Note that one agent cannot be added to two different
* parents' sub-agents lists.
Expand All @@ -79,7 +73,7 @@ public abstract class BaseAgent {
* Creates a new BaseAgent.
*
* @param name Unique agent name. Cannot be "user" (reserved).
* @param description Agent purpose.
* @param description Agent purpose; null is treated as an empty string.
* @param subAgents Agents managed by this agent.
* @param beforeAgentCallback Callbacks before agent execution. Invoked in order until one doesn't
* return null.
Expand All @@ -88,16 +82,14 @@ public abstract class BaseAgent {
*/
public BaseAgent(
String name,
String description,
@Nullable String description,
@Nullable List<? extends BaseAgent> subAgents,
@Nullable List<? extends BeforeAgentCallback> beforeAgentCallback,
@Nullable List<? extends AfterAgentCallback> afterAgentCallback) {
validateAgentName(name);
this.name = name;
this.description = description;
super(validateAgentName(name), nullToEmpty(description));
this.parentAgent = null;
this.subAgents = (subAgents != null) ? ImmutableList.copyOf(subAgents) : ImmutableList.of();
validateSubAgents(this.name, this.subAgents);
validateSubAgents(name, this.subAgents);
this.beforeAgentCallback =
(beforeAgentCallback != null)
? ImmutableList.copyOf(beforeAgentCallback)
Expand Down Expand Up @@ -128,10 +120,12 @@ public Completable close() {
* Validates the agent name.
*
* @param name The agent name to validate.
* @return the validated name.
* @throws IllegalArgumentException if the agent name is null, empty, or does not match the
* identifier pattern.
*/
private static void validateAgentName(String name) {
@CanIgnoreReturnValue
private static String validateAgentName(String name) {
if (isNullOrEmpty(name)) {
throw new IllegalArgumentException("Agent name cannot be null or empty.");
}
Expand All @@ -141,6 +135,7 @@ private static void validateAgentName(String name) {
}
checkArgument(
!name.equals(Role.USER), "Agent name cannot be 'user'; reserved for end-user input.");
return name;
}

/**
Expand Down Expand Up @@ -173,24 +168,6 @@ private static void validateSubAgents(
}
}

/**
* Gets the agent's unique name.
*
* @return the unique name of the agent.
*/
public final String name() {
return name;
}

/**
* Gets the one-line description of the agent's capability.
*
* @return the description of the agent.
*/
public final String description() {
return description;
}

/**
* Retrieves the parent agent in the agent tree.
*
Expand Down Expand Up @@ -312,7 +289,7 @@ private InvocationContext createInvocationContext(InvocationContext parentContex
* @return stream of agent-generated events.
*/
public Flowable<Event> runAsync(InvocationContext parentContext) {
return run(parentContext, this::runAsyncImpl);
return runWithCallbacks(parentContext, this::runAsyncImpl);
}

/**
Expand All @@ -322,7 +299,7 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
* @param runImplementation The agent-specific logic to run.
* @return stream of agent-generated events.
*/
private Flowable<Event> run(
private Flowable<Event> runWithCallbacks(
InvocationContext parentContext,
Function<InvocationContext, Flowable<Event>> runImplementation) {
return Flowable.using(
Expand Down Expand Up @@ -468,7 +445,7 @@ private Maybe<Event> callCallback(
* @return stream of agent-generated events.
*/
public Flowable<Event> runLive(InvocationContext parentContext) {
return run(parentContext, this::runLiveImpl);
return runWithCallbacks(parentContext, this::runLiveImpl);
}

/**
Expand Down Expand Up @@ -510,6 +487,17 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) {
.build();
}

/**
* Runs this agent as a workflow node through its {@link #runAsync} lifecycle, including its
* callbacks. The node input is not used.
*/
@Experimental
@Override
public Flowable<Event> runNode(
com.google.adk.agents.Context context, @Nullable Object nodeInput) {
return runAsync(context.invocationContext());
}

/**
* Agent-specific asynchronous logic.
*
Expand Down
2 changes: 2 additions & 0 deletions core/src/main/java/com/google/adk/agents/LlmAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -1029,10 +1029,12 @@ public List<? extends OnToolErrorCallback> canonicalOnToolErrorCallbacks() {
return onToolErrorCallback;
}

@Override
public Optional<Schema> inputSchema() {
return inputSchema;
}

@Override
public Optional<Schema> outputSchema() {
return outputSchema;
}
Expand Down
50 changes: 50 additions & 0 deletions core/src/test/java/com/google/adk/agents/BaseAgentTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.google.adk.testing.TestBaseAgent;
import com.google.adk.testing.TestCallback;
import com.google.adk.testing.TestUtils;
import com.google.adk.workflow.Node;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
Expand Down Expand Up @@ -112,6 +113,13 @@ public void constructor_setsNameAndDescription() {
assertThat(agent.description()).isEqualTo(description);
}

@Test
public void constructor_nullDescription_setsEmptyDescription() {
TestBaseAgent agent = new TestBaseAgent("testName", null, null, ImmutableList.of(), null, null);

assertThat(agent.description()).isEmpty();
}

@Test
public void findAgent_returnsCorrectAgent() {
TestBaseAgent subSubAgent =
Expand Down Expand Up @@ -399,6 +407,34 @@ public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback()
assertThat(afterCallback2.wasCalled()).isFalse();
}

@Test
public void runNode_runsAgentWithItsCallbacks() {
var runAsyncImpl = TestCallback.<Void>returningEmpty();
Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output"));
Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output"));
var beforeCallback = TestCallback.<Content>returningEmpty();
var afterCallback = TestCallback.returning(afterCallbackContent);
TestBaseAgent agent =
new TestBaseAgent(
TEST_AGENT_NAME,
TEST_AGENT_DESCRIPTION,
ImmutableList.of(beforeCallback.asBeforeAgentCallback()),
ImmutableList.of(afterCallback.asAfterAgentCallback()),
runAsyncImpl.asRunAsyncImplSupplier(runAsyncImplContent));
Context context =
new CallbackContext(TestUtils.createInvocationContext(agent), /* eventActions= */ null);
Node node = agent;

List<Event> results =
node.runNode(context, /* nodeInput= */ "unused").cast(Event.class).toList().blockingGet();

assertThat(results).hasSize(2);
assertThat(results.get(0).content()).hasValue(runAsyncImplContent);
assertThat(results.get(1).content()).hasValue(afterCallbackContent);
assertThat(beforeCallback.wasCalled()).isTrue();
assertThat(runAsyncImpl.wasCalled()).isTrue();
}

@Test
public void canonicalCallbacks_returnsEmptyListWhenNull() {
TestBaseAgent agent =
Expand Down Expand Up @@ -637,6 +673,20 @@ public void runLive_firstAfterCallbackReturnsContent_skipsSecondAfterCallback()
assertThat(afterCallback2.wasCalled()).isFalse();
}

@Test
public void constructor_nullName_throwsIllegalArgumentException() {
assertThrows(
IllegalArgumentException.class,
() -> new TestBaseAgent(null, "description", null, null, null));
}

@Test
public void constructor_emptyName_throwsIllegalArgumentException() {
assertThrows(
IllegalArgumentException.class,
() -> new TestBaseAgent("", "description", null, null, null));
}

@Test
public void constructor_invalidName_throwsIllegalArgumentException() {
assertThrows(
Expand Down
Loading