Skip to content

Commit 40d5df3

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: Update AgentExecutor so it builds new runner on execute and there is no need to pass the runner instance
PiperOrigin-RevId: 866423638
1 parent 5262d4a commit 40d5df3

2 files changed

Lines changed: 174 additions & 15 deletions

File tree

‎a2a/src/main/java/com/google/adk/a2a/AgentExecutor.java‎

Lines changed: 92 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,13 @@
22

33
import com.google.adk.a2a.converters.EventConverter;
44
import com.google.adk.a2a.converters.PartConverter;
5+
import com.google.adk.agents.BaseAgent;
56
import com.google.adk.agents.RunConfig;
7+
import com.google.adk.apps.App;
8+
import com.google.adk.artifacts.BaseArtifactService;
69
import com.google.adk.events.Event;
10+
import com.google.adk.memory.BaseMemoryService;
11+
import com.google.adk.plugins.Plugin;
712
import com.google.adk.runner.Runner;
813
import com.google.adk.sessions.BaseSessionService;
914
import com.google.adk.sessions.Session;
@@ -21,6 +26,7 @@
2126
import io.reactivex.rxjava3.core.Maybe;
2227
import io.reactivex.rxjava3.disposables.CompositeDisposable;
2328
import io.reactivex.rxjava3.disposables.Disposable;
29+
import java.util.List;
2430
import java.util.Map;
2531
import java.util.Optional;
2632
import java.util.UUID;
@@ -41,29 +47,98 @@ public class AgentExecutor implements io.a2a.server.agentexecution.AgentExecutor
4147
private static final RunConfig DEFAULT_RUN_CONFIG =
4248
RunConfig.builder().setStreamingMode(RunConfig.StreamingMode.NONE).setMaxLlmCalls(20).build();
4349

44-
private final Runner runner;
4550
private final Map<String, Disposable> activeTasks = new ConcurrentHashMap<>();
51+
private final Runner.Builder runnerBuilder;
52+
private final RunConfig runConfig;
4653

47-
private AgentExecutor(Runner runner) {
48-
this.runner = runner;
54+
private AgentExecutor(
55+
App app,
56+
BaseAgent agent,
57+
String appName,
58+
BaseArtifactService artifactService,
59+
BaseSessionService sessionService,
60+
BaseMemoryService memoryService,
61+
List<? extends Plugin> plugins,
62+
RunConfig runConfig) {
63+
this.runnerBuilder =
64+
Runner.builder()
65+
.agent(agent)
66+
.appName(appName)
67+
.artifactService(artifactService)
68+
.sessionService(sessionService)
69+
.memoryService(memoryService)
70+
.plugins(plugins == null ? ImmutableList.of() : plugins);
71+
if (app != null) {
72+
this.runnerBuilder.app(app);
73+
}
74+
// Check that the runner is configured correctly and can be built.
75+
var unused = runnerBuilder.build();
76+
this.runConfig = runConfig == null ? DEFAULT_RUN_CONFIG : runConfig;
4977
}
5078

5179
/** Builder for {@link AgentExecutor}. */
5280
public static class Builder {
53-
private Runner runner;
81+
private App app;
82+
private BaseAgent agent;
83+
private String appName;
84+
private BaseArtifactService artifactService;
85+
private BaseSessionService sessionService;
86+
private BaseMemoryService memoryService;
87+
private List<? extends Plugin> plugins;
88+
private RunConfig runConfig;
89+
90+
@CanIgnoreReturnValue
91+
public Builder app(App app) {
92+
this.app = app;
93+
return this;
94+
}
95+
96+
@CanIgnoreReturnValue
97+
public Builder agent(BaseAgent agent) {
98+
this.agent = agent;
99+
return this;
100+
}
101+
102+
@CanIgnoreReturnValue
103+
public Builder appName(String appName) {
104+
this.appName = appName;
105+
return this;
106+
}
107+
108+
@CanIgnoreReturnValue
109+
public Builder artifactService(BaseArtifactService artifactService) {
110+
this.artifactService = artifactService;
111+
return this;
112+
}
113+
114+
@CanIgnoreReturnValue
115+
public Builder sessionService(BaseSessionService sessionService) {
116+
this.sessionService = sessionService;
117+
return this;
118+
}
54119

55120
@CanIgnoreReturnValue
56-
public Builder runner(Runner runner) {
57-
this.runner = runner;
121+
public Builder memoryService(BaseMemoryService memoryService) {
122+
this.memoryService = memoryService;
123+
return this;
124+
}
125+
126+
@CanIgnoreReturnValue
127+
public Builder plugins(List<? extends Plugin> plugins) {
128+
this.plugins = plugins;
129+
return this;
130+
}
131+
132+
@CanIgnoreReturnValue
133+
public Builder runConfig(RunConfig runConfig) {
134+
this.runConfig = runConfig;
58135
return this;
59136
}
60137

61138
@CanIgnoreReturnValue
62139
public AgentExecutor build() {
63-
if (runner == null) {
64-
throw new IllegalStateException("Runner must be provided.");
65-
}
66-
return new AgentExecutor(runner);
140+
return new AgentExecutor(
141+
app, agent, appName, artifactService, sessionService, memoryService, plugins, runConfig);
67142
}
68143
}
69144

@@ -96,13 +171,14 @@ public void execute(RequestContext ctx, EventQueue eventQueue) {
96171

97172
EventProcessor p = new EventProcessor();
98173
Content content = PartConverter.messageToContent(message);
174+
Runner runner = runnerBuilder.build();
99175

100176
taskDisposables.add(
101-
prepareSession(ctx, runner.sessionService())
177+
prepareSession(ctx, runner.appName(), runner.sessionService())
102178
.flatMapPublisher(
103179
session -> {
104180
updater.startWork();
105-
return runner.runAsync(getUserId(ctx), session.id(), content, DEFAULT_RUN_CONFIG);
181+
return runner.runAsync(getUserId(ctx), session.id(), content, runConfig);
106182
})
107183
.subscribe(
108184
event -> {
@@ -130,13 +206,14 @@ private String getUserId(RequestContext ctx) {
130206
return USER_ID_PREFIX + ctx.getContextId();
131207
}
132208

133-
private Maybe<Session> prepareSession(RequestContext ctx, BaseSessionService service) {
209+
private Maybe<Session> prepareSession(
210+
RequestContext ctx, String appName, BaseSessionService service) {
134211
return service
135-
.getSession(runner.appName(), getUserId(ctx), ctx.getContextId(), Optional.empty())
212+
.getSession(appName, getUserId(ctx), ctx.getContextId(), Optional.empty())
136213
.switchIfEmpty(
137214
Maybe.defer(
138215
() -> {
139-
return service.createSession(runner.appName(), getUserId(ctx)).toMaybe();
216+
return service.createSession(appName, getUserId(ctx)).toMaybe();
140217
}));
141218
}
142219

‎a2a/src/test/java/com/google/adk/a2a/AgentExecutorTest.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package com.google.adk.a2a;
2+
3+
import static org.junit.Assert.assertThrows;
4+
5+
import com.google.adk.agents.BaseAgent;
6+
import com.google.adk.agents.InvocationContext;
7+
import com.google.adk.apps.App;
8+
import com.google.adk.artifacts.InMemoryArtifactService;
9+
import com.google.adk.events.Event;
10+
import com.google.adk.sessions.InMemorySessionService;
11+
import com.google.common.collect.ImmutableList;
12+
import io.reactivex.rxjava3.core.Flowable;
13+
import org.junit.Before;
14+
import org.junit.Test;
15+
import org.junit.runner.RunWith;
16+
import org.junit.runners.JUnit4;
17+
18+
@RunWith(JUnit4.class)
19+
public final class AgentExecutorTest {
20+
21+
private TestAgent testAgent;
22+
23+
@Before
24+
public void setUp() {
25+
testAgent = new TestAgent();
26+
}
27+
28+
@Test
29+
public void createAgentExecutor_noAgent_succeeds() {
30+
var unused =
31+
new AgentExecutor.Builder()
32+
.app(App.builder().name("test_app").rootAgent(testAgent).build())
33+
.sessionService(new InMemorySessionService())
34+
.artifactService(new InMemoryArtifactService())
35+
.build();
36+
}
37+
38+
@Test
39+
public void createAgentExecutor_withAgentAndApp_throwsException() {
40+
assertThrows(
41+
IllegalStateException.class,
42+
() -> {
43+
new AgentExecutor.Builder()
44+
.agent(testAgent)
45+
.app(App.builder().name("test_app").rootAgent(testAgent).build())
46+
.sessionService(new InMemorySessionService())
47+
.artifactService(new InMemoryArtifactService())
48+
.build();
49+
});
50+
}
51+
52+
@Test
53+
public void createAgentExecutor_withEmptyAgentAndApp_throwsException() {
54+
assertThrows(
55+
IllegalStateException.class,
56+
() -> {
57+
new AgentExecutor.Builder()
58+
.sessionService(new InMemorySessionService())
59+
.artifactService(new InMemoryArtifactService())
60+
.build();
61+
});
62+
}
63+
64+
private static final class TestAgent extends BaseAgent {
65+
private final Flowable<Event> eventsToEmit = Flowable.empty();
66+
67+
TestAgent() {
68+
// BaseAgent constructor: name, description, examples, tools, model
69+
super("test_agent", "test", ImmutableList.of(), null, null);
70+
}
71+
72+
@Override
73+
protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
74+
return eventsToEmit;
75+
}
76+
77+
@Override
78+
protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
79+
return eventsToEmit;
80+
}
81+
}
82+
}

0 commit comments

Comments
 (0)