Skip to content

Commit e0f7137

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: add the AgentExecutor config
PiperOrigin-RevId: 876156542
1 parent f422b47 commit e0f7137

4 files changed

Lines changed: 163 additions & 17 deletions

File tree

a2a/src/main/java/com/google/adk/a2a/AgentExecutor.java renamed to a2a/src/main/java/com/google/adk/a2a/executor/AgentExecutor.java

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
package com.google.adk.a2a;
1+
package com.google.adk.a2a.executor;
2+
3+
import static java.util.Objects.requireNonNull;
24

35
import com.google.adk.a2a.converters.EventConverter;
46
import com.google.adk.a2a.converters.PartConverter;
57
import com.google.adk.agents.BaseAgent;
6-
import com.google.adk.agents.RunConfig;
78
import com.google.adk.apps.App;
89
import com.google.adk.artifacts.BaseArtifactService;
910
import com.google.adk.events.Event;
@@ -44,12 +45,10 @@ public class AgentExecutor implements io.a2a.server.agentexecution.AgentExecutor
4445

4546
private static final Logger logger = LoggerFactory.getLogger(AgentExecutor.class);
4647
private static final String USER_ID_PREFIX = "A2A_USER_";
47-
private static final RunConfig DEFAULT_RUN_CONFIG =
48-
RunConfig.builder().setStreamingMode(RunConfig.StreamingMode.NONE).setMaxLlmCalls(20).build();
4948

5049
private final Map<String, Disposable> activeTasks = new ConcurrentHashMap<>();
5150
private final Runner.Builder runnerBuilder;
52-
private final RunConfig runConfig;
51+
private final AgentExecutorConfig agentExecutorConfig;
5352

5453
private AgentExecutor(
5554
App app,
@@ -59,7 +58,10 @@ private AgentExecutor(
5958
BaseSessionService sessionService,
6059
BaseMemoryService memoryService,
6160
List<? extends Plugin> plugins,
62-
RunConfig runConfig) {
61+
AgentExecutorConfig agentExecutorConfig) {
62+
requireNonNull(agentExecutorConfig);
63+
this.agentExecutorConfig = agentExecutorConfig;
64+
6365
this.runnerBuilder =
6466
Runner.builder()
6567
.agent(agent)
@@ -73,7 +75,6 @@ private AgentExecutor(
7375
}
7476
// Check that the runner is configured correctly and can be built.
7577
var unused = runnerBuilder.build();
76-
this.runConfig = runConfig == null ? DEFAULT_RUN_CONFIG : runConfig;
7778
}
7879

7980
/** Builder for {@link AgentExecutor}. */
@@ -85,7 +86,13 @@ public static class Builder {
8586
private BaseSessionService sessionService;
8687
private BaseMemoryService memoryService;
8788
private List<? extends Plugin> plugins = ImmutableList.of();
88-
private RunConfig runConfig;
89+
private AgentExecutorConfig agentExecutorConfig;
90+
91+
@CanIgnoreReturnValue
92+
public Builder agentExecutorConfig(AgentExecutorConfig agentExecutorConfig) {
93+
this.agentExecutorConfig = agentExecutorConfig;
94+
return this;
95+
}
8996

9097
@CanIgnoreReturnValue
9198
public Builder app(App app) {
@@ -129,16 +136,17 @@ public Builder plugins(List<? extends Plugin> plugins) {
129136
return this;
130137
}
131138

132-
@CanIgnoreReturnValue
133-
public Builder runConfig(RunConfig runConfig) {
134-
this.runConfig = runConfig;
135-
return this;
136-
}
137-
138139
@CanIgnoreReturnValue
139140
public AgentExecutor build() {
140141
return new AgentExecutor(
141-
app, agent, appName, artifactService, sessionService, memoryService, plugins, runConfig);
142+
app,
143+
agent,
144+
appName,
145+
artifactService,
146+
sessionService,
147+
memoryService,
148+
plugins,
149+
agentExecutorConfig);
142150
}
143151
}
144152

@@ -178,7 +186,8 @@ public void execute(RequestContext ctx, EventQueue eventQueue) {
178186
.flatMapPublisher(
179187
session -> {
180188
updater.startWork();
181-
return runner.runAsync(getUserId(ctx), session.id(), content, runConfig);
189+
return runner.runAsync(
190+
getUserId(ctx), session.id(), content, agentExecutorConfig.runConfig());
182191
})
183192
.subscribe(
184193
event -> {

‎a2a/src/main/java/com/google/adk/a2a/executor/AgentExecutorConfig.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package com.google.adk.a2a.executor;
2+
3+
import com.google.adk.a2a.executor.Callbacks.AfterEventCallback;
4+
import com.google.adk.a2a.executor.Callbacks.AfterExecuteCallback;
5+
import com.google.adk.a2a.executor.Callbacks.BeforeExecuteCallback;
6+
import com.google.adk.agents.RunConfig;
7+
import com.google.auto.value.AutoValue;
8+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
9+
import org.jspecify.annotations.Nullable;
10+
11+
/** Configuration for the {@link AgentExecutor}. */
12+
@AutoValue
13+
public abstract class AgentExecutorConfig {
14+
15+
private static final RunConfig DEFAULT_RUN_CONFIG =
16+
RunConfig.builder().setStreamingMode(RunConfig.StreamingMode.NONE).setMaxLlmCalls(20).build();
17+
18+
public abstract RunConfig runConfig();
19+
20+
public abstract @Nullable BeforeExecuteCallback beforeExecuteCallback();
21+
22+
public abstract @Nullable AfterExecuteCallback afterExecuteCallback();
23+
24+
public abstract @Nullable AfterEventCallback afterEventCallback();
25+
26+
public abstract Builder toBuilder();
27+
28+
public static Builder builder() {
29+
return new AutoValue_AgentExecutorConfig.Builder().runConfig(DEFAULT_RUN_CONFIG);
30+
}
31+
32+
/** Builder for {@link AgentExecutorConfig}. */
33+
@AutoValue.Builder
34+
public abstract static class Builder {
35+
@CanIgnoreReturnValue
36+
public abstract Builder runConfig(RunConfig runConfig);
37+
38+
@CanIgnoreReturnValue
39+
public abstract Builder beforeExecuteCallback(BeforeExecuteCallback beforeExecuteCallback);
40+
41+
@CanIgnoreReturnValue
42+
public abstract Builder afterExecuteCallback(AfterExecuteCallback afterExecuteCallback);
43+
44+
@CanIgnoreReturnValue
45+
public abstract Builder afterEventCallback(AfterEventCallback afterEventCallback);
46+
47+
abstract AgentExecutorConfig autoBuild();
48+
49+
public AgentExecutorConfig build() {
50+
return autoBuild();
51+
}
52+
}
53+
}

‎a2a/src/main/java/com/google/adk/a2a/executor/Callbacks.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package com.google.adk.a2a.executor;
2+
3+
import com.google.adk.events.Event;
4+
import io.a2a.server.agentexecution.RequestContext;
5+
import io.a2a.spec.TaskArtifactUpdateEvent;
6+
import io.a2a.spec.TaskStatusUpdateEvent;
7+
import io.reactivex.rxjava3.core.Maybe;
8+
import io.reactivex.rxjava3.core.Single;
9+
10+
/** Functional interfaces for agent executor lifecycle callbacks. */
11+
public final class Callbacks {
12+
13+
private Callbacks() {}
14+
15+
interface BeforeExecuteCallbackBase {}
16+
17+
/** Async callback interface for actions to be performed before an execution is started. */
18+
@FunctionalInterface
19+
public interface BeforeExecuteCallback extends BeforeExecuteCallbackBase {
20+
/**
21+
* Callback which will be called before an execution is started. It can be used to instrument a
22+
* context or prevent the execution by returning an error.
23+
*
24+
* @param ctx the request context
25+
* @return a {@link Single} that completes with a boolean indicating whether the execution
26+
* should be prevented
27+
*/
28+
Single<Boolean> call(RequestContext ctx);
29+
}
30+
31+
interface AfterExecuteCallbackBase {}
32+
33+
/**
34+
* Async callback interface for actions to be performed after an execution is completed or failed.
35+
*/
36+
@FunctionalInterface
37+
public interface AfterExecuteCallback extends AfterExecuteCallbackBase {
38+
/**
39+
* Callback which will be called after an execution resolved into a completed or failed task.
40+
* This gives an opportunity to enrich the event with additional metadata or log it.
41+
*
42+
* @param ctx the request context
43+
* @param finalUpdateEvent the final update event
44+
* @return a {@link Maybe} that completes when the callback is done
45+
*/
46+
Maybe<TaskStatusUpdateEvent> call(RequestContext ctx, TaskStatusUpdateEvent finalUpdateEvent);
47+
}
48+
49+
interface AfterEventCallbackBase {}
50+
51+
/** Async callback interface for actions to be performed after an event is processed. */
52+
@FunctionalInterface
53+
public interface AfterEventCallback extends AfterEventCallbackBase {
54+
/**
55+
* Callback which will be called after an ADK event is successfully converted to an A2A event.
56+
* This gives an opportunity to enrich the event with additional metadata or abort the execution
57+
* by returning an error. The callback is not invoked for errors originating from ADK or event
58+
* processing.
59+
*
60+
* @param ctx the request context
61+
* @param processedEvent the processed task artifact update event
62+
* @param event the ADK event
63+
* @return a {@link Maybe} that completes when the callback is done
64+
*/
65+
Maybe<TaskArtifactUpdateEvent> call(
66+
RequestContext ctx, TaskArtifactUpdateEvent processedEvent, Event event);
67+
}
68+
}

a2a/src/test/java/com/google/adk/a2a/AgentExecutorTest.java renamed to a2a/src/test/java/com/google/adk/a2a/executor/AgentExecutorTest.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.google.adk.a2a;
1+
package com.google.adk.a2a.executor;
22

33
import static org.junit.Assert.assertThrows;
44

@@ -32,6 +32,7 @@ public void createAgentExecutor_noAgent_succeeds() {
3232
.app(App.builder().name("test_app").rootAgent(testAgent).build())
3333
.sessionService(new InMemorySessionService())
3434
.artifactService(new InMemoryArtifactService())
35+
.agentExecutorConfig(AgentExecutorConfig.builder().build())
3536
.build();
3637
}
3738

@@ -44,6 +45,7 @@ public void createAgentExecutor_withAgentAndApp_throwsException() {
4445
.agent(testAgent)
4546
.app(App.builder().name("test_app").rootAgent(testAgent).build())
4647
.sessionService(new InMemorySessionService())
48+
.agentExecutorConfig(AgentExecutorConfig.builder().build())
4749
.artifactService(new InMemoryArtifactService())
4850
.build();
4951
});
@@ -55,6 +57,20 @@ public void createAgentExecutor_withEmptyAgentAndApp_throwsException() {
5557
IllegalStateException.class,
5658
() -> {
5759
new AgentExecutor.Builder()
60+
.sessionService(new InMemorySessionService())
61+
.artifactService(new InMemoryArtifactService())
62+
.agentExecutorConfig(AgentExecutorConfig.builder().build())
63+
.build();
64+
});
65+
}
66+
67+
@Test
68+
public void createAgentExecutor_noAgentExecutorConfig_throwsException() {
69+
assertThrows(
70+
NullPointerException.class,
71+
() -> {
72+
new AgentExecutor.Builder()
73+
.app(App.builder().name("test_app").rootAgent(testAgent).build())
5874
.sessionService(new InMemorySessionService())
5975
.artifactService(new InMemoryArtifactService())
6076
.build();

0 commit comments

Comments
 (0)