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
45 changes: 32 additions & 13 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import com.google.adk.telemetry.Tracing;
import com.google.adk.utils.AgentEnums.AgentOrigin;
import com.google.adk.workflow.BaseNode;
import com.google.adk.workflow.NodeExecutionState;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.errorprone.annotations.DoNotCall;
Expand Down Expand Up @@ -269,16 +270,20 @@ public ImmutableList<? extends AfterAgentCallback> canonicalAfterAgentCallbacks(
* instance.
*
* @param parentContext Parent context to copy.
* @return new context with updated branch name.
* @param keepBranch Whether to keep the parent's branch instead of appending this agent's name.
* @return the new context for this agent.
*/
private InvocationContext createInvocationContext(InvocationContext parentContext) {
private InvocationContext createInvocationContext(
InvocationContext parentContext, boolean keepBranch) {
InvocationContext.Builder builder = parentContext.toBuilder();
builder.agent(this);
// Check for branch to be truthy (not None, not empty string),
parentContext
.branch()
.filter(s -> !s.isEmpty())
.ifPresent(branch -> builder.branch(branch + "." + name()));
if (!keepBranch) {
// Check for branch to be truthy (not None, not empty string),
parentContext
.branch()
.filter(s -> !s.isEmpty())
.ifPresent(branch -> builder.branch(branch + "." + name()));
}
return builder.build();
}

Expand All @@ -289,24 +294,26 @@ private InvocationContext createInvocationContext(InvocationContext parentContex
* @return stream of agent-generated events.
*/
public Flowable<Event> runAsync(InvocationContext parentContext) {
return runWithCallbacks(parentContext, this::runAsyncImpl);
return runWithCallbacks(parentContext, /* keepBranch= */ false, this::runAsyncImpl);
}

/**
* Runs the agent with the given implementation.
*
* @param parentContext Parent context to inherit.
* @param keepBranch Whether to keep the parent's branch instead of appending this agent's name.
* @param runImplementation The agent-specific logic to run.
* @return stream of agent-generated events.
*/
private Flowable<Event> runWithCallbacks(
InvocationContext parentContext,
boolean keepBranch,
Function<InvocationContext, Flowable<Event>> runImplementation) {
return Flowable.using(
() -> {
Context otelContext = Context.current();
return Instrumentation.recordAgentInvocation(
createInvocationContext(parentContext), this, otelContext);
createInvocationContext(parentContext, keepBranch), this, otelContext);
},
agentInvocation -> {
InvocationContext invocationContext = agentInvocation.getCtx();
Expand Down Expand Up @@ -445,7 +452,7 @@ private Maybe<Event> callCallback(
* @return stream of agent-generated events.
*/
public Flowable<Event> runLive(InvocationContext parentContext) {
return runWithCallbacks(parentContext, this::runLiveImpl);
return runWithCallbacks(parentContext, /* keepBranch= */ false, this::runLiveImpl);
}

/**
Expand Down Expand Up @@ -488,14 +495,26 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) {
}

/**
* Runs this agent as a workflow node through its {@link #runAsync} lifecycle, including its
* callbacks. The node input is not used.
* Runs this agent as a workflow node, on the node's branch so the nodes after it see its events.
* It runs the callbacks and {@link #runAsyncImpl} that {@link #runAsync} runs, but does not call
* an override of {@link #runAsync}. The node input is not used, and each event's author is copied
* into the node state so that the agent's events keep their own authors instead of an inherited
* author or the node's name.
*
* @throws IllegalStateException if {@code context} is not the context of a workflow node.
*/
@Experimental
@Override
public Flowable<Event> runNode(
com.google.adk.agents.Context context, @Nullable Object nodeInput) {
return runAsync(context.invocationContext());
NodeExecutionState nodeState = context.requireNodeState();
return runWithCallbacks(context.invocationContext(), /* keepBranch= */ true, this::runAsyncImpl)
.doOnNext(
event -> {
if (!isNullOrEmpty(event.author())) {
nodeState.eventAuthor(event.author());
}
});
}

/**
Expand Down
96 changes: 94 additions & 2 deletions core/src/main/java/com/google/adk/agents/Context.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,16 @@

package com.google.adk.agents;

import com.google.adk.annotations.Experimental;
import com.google.adk.artifacts.ListArtifactsResponse;
import com.google.adk.events.EventActions;
import com.google.adk.events.ToolConfirmation;
import com.google.adk.memory.SearchMemoryResponse;
import com.google.adk.sessions.State;
import com.google.adk.workflow.NodeExecutionState;
import com.google.adk.workflow.Route;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
Expand All @@ -31,12 +35,13 @@
import org.jspecify.annotations.Nullable;

/**
* The context passed to callbacks and tools during an invocation.
* The context passed to callbacks, tools, and workflow nodes during an invocation.
*
* <p>Agent and model callbacks receive a {@link CallbackContext}, and tools and tool callbacks a
* {@link com.google.adk.tools.ToolContext}; both extend this class, so a helper that serves both,
* or a function tool's {@code toolContext} parameter, can take a {@code Context}. {@link
* #requestConfirmation()} needs a function call ID, which only a tool call has.
* #requestConfirmation()} requires a function call ID, which only a tool call has, while {@link
* #requireNodeState()} and the other workflow node methods require the context of a workflow node.
*/
public class Context extends ReadonlyContext {

Expand All @@ -45,6 +50,7 @@ public class Context extends ReadonlyContext {

private final State state;
private final @Nullable String eventId;
private final @Nullable NodeExecutionState nodeState;
private Optional<String> functionCallId = Optional.empty();
private Optional<ToolConfirmation> toolConfirmation = Optional.empty();

Expand All @@ -59,10 +65,32 @@ public class Context extends ReadonlyContext {
InvocationContext invocationContext,
@Nullable EventActions eventActions,
@Nullable String eventId) {
this(invocationContext, eventActions, eventId, /* nodeState= */ null);
}

/**
* Creates the context of the workflow node run whose state is {@code nodeState}. Only the
* workflow runtime creates a node state.
*/
@Experimental
public Context(InvocationContext invocationContext, NodeExecutionState nodeState) {
this(
invocationContext,
/* eventActions= */ null,
/* eventId= */ null,
Preconditions.checkNotNull(nodeState));
}

private Context(
InvocationContext invocationContext,
@Nullable EventActions eventActions,
@Nullable String eventId,
@Nullable NodeExecutionState nodeState) {
super(invocationContext);
this.eventActions = eventActions != null ? eventActions : EventActions.builder().build();
this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta());
this.eventId = eventId;
this.nodeState = nodeState;
}

/** Returns the delta-aware state of the current context. */
Expand Down Expand Up @@ -223,4 +251,68 @@ public Single<SearchMemoryResponse> searchMemory(String query) {
.searchMemory(
invocationContext.session().appName(), invocationContext.session().userId(), query);
}

/**
* Returns the state of the workflow node run this context belongs to.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final NodeExecutionState requireNodeState() {
Preconditions.checkState(
nodeState != null, "Only the context of a workflow node has a node state.");
return nodeState;
}

/**
* Returns the author of the events this workflow node emits, or empty to use the node's name.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final String eventAuthor() {
return requireNodeState().eventAuthor();
}

/**
* Sets the author of the events this workflow node emits.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final void eventAuthor(String eventAuthor) {
requireNodeState().eventAuthor(eventAuthor);
}

/**
* Returns this workflow node's output, or empty if it has produced none.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final Optional<Object> output() {
return requireNodeState().output();
}

/**
* Returns the routes this workflow node selected, or empty if it made no routing decision. An
* empty list is a decision that selects no route.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final Optional<ImmutableList<Route>> routes() {
return requireNodeState().selectedRoutes();
}

/**
* Selects this workflow node's routes, which decide the outgoing edges to follow, replacing any
* earlier selection.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final void routes(List<? extends Route> routes) {
requireNodeState().selectRoutes(routes);
}
}
Loading
Loading