Skip to content

Commit 47fac93

Browse files
google-genai-botcopybara-github
authored andcommitted
feat(HITL): Let LlmAgent.tools() return an async type instead of blocking
PiperOrigin-RevId: 843824506
1 parent c5cbc6d commit 47fac93

4 files changed

Lines changed: 61 additions & 37 deletions

File tree

core/src/main/java/com/google/adk/agents/LlmAgent.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -736,8 +736,8 @@ public IncludeContents includeContents() {
736736
return includeContents;
737737
}
738738

739-
public List<BaseTool> tools() {
740-
return canonicalTools().toList().blockingGet();
739+
public Single<List<BaseTool>> tools() {
740+
return canonicalTools().toList();
741741
}
742742

743743
public List<Object> toolsUnion() {

core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.java

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,16 @@ private Maybe<Event> assembleEvent(
106106
InvocationContext invocationContext,
107107
Collection<FunctionCall> functionCalls,
108108
Map<String, ToolConfirmation> toolConfirmations) {
109-
ImmutableMap.Builder<String, BaseTool> toolsBuilder = ImmutableMap.builder();
109+
Single<ImmutableMap<String, BaseTool>> toolsMapSingle;
110110
if (invocationContext.agent() instanceof LlmAgent llmAgent) {
111-
for (BaseTool tool : llmAgent.tools()) {
112-
toolsBuilder.put(tool.name(), tool);
113-
}
111+
toolsMapSingle =
112+
llmAgent
113+
.tools()
114+
.map(
115+
toolList ->
116+
toolList.stream().collect(toImmutableMap(BaseTool::name, tool -> tool)));
117+
} else {
118+
toolsMapSingle = Single.just(ImmutableMap.of());
114119
}
115120

116121
var functionCallEvent =
@@ -124,8 +129,10 @@ private Maybe<Event> assembleEvent(
124129
.build())
125130
.build();
126131

127-
return Functions.handleFunctionCalls(
128-
invocationContext, functionCallEvent, toolsBuilder.buildOrThrow(), toolConfirmations);
132+
return toolsMapSingle.flatMapMaybe(
133+
toolsMap ->
134+
Functions.handleFunctionCalls(
135+
invocationContext, functionCallEvent, toolsMap, toolConfirmations));
129136
}
130137

131138
private ImmutableMap<String, ToolConfirmation> filterRequestConfirmationFunctionResponses(

core/src/main/java/com/google/adk/runner/Runner.java

Lines changed: 43 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -521,35 +521,52 @@ public Flowable<Event> runLive(
521521
try {
522522
InvocationContext invocationContext =
523523
newInvocationContextForLive(session, Optional.of(liveRequestQueue), runConfig);
524+
525+
Single<InvocationContext> invocationContextSingle;
524526
if (invocationContext.agent() instanceof LlmAgent) {
525527
LlmAgent agent = (LlmAgent) invocationContext.agent();
526-
for (BaseTool tool : agent.tools()) {
527-
if (tool instanceof FunctionTool functionTool) {
528-
for (Parameter parameter : functionTool.func().getParameters()) {
529-
if (parameter.getType().equals(LiveRequestQueue.class)) {
530-
invocationContext
531-
.activeStreamingTools()
532-
.put(functionTool.name(), new ActiveStreamingTool(new LiveRequestQueue()));
533-
}
534-
}
535-
}
536-
}
528+
invocationContextSingle =
529+
agent
530+
.tools()
531+
.map(
532+
tools -> {
533+
for (BaseTool tool : tools) {
534+
if (tool instanceof FunctionTool functionTool) {
535+
for (Parameter parameter : functionTool.func().getParameters()) {
536+
if (parameter.getType().equals(LiveRequestQueue.class)) {
537+
invocationContext
538+
.activeStreamingTools()
539+
.put(
540+
functionTool.name(),
541+
new ActiveStreamingTool(new LiveRequestQueue()));
542+
}
543+
}
544+
}
545+
}
546+
return invocationContext;
547+
});
548+
} else {
549+
invocationContextSingle = Single.just(invocationContext);
537550
}
538-
return Telemetry.traceFlowable(
539-
spanContext,
540-
span,
541-
() ->
542-
invocationContext
543-
.agent()
544-
.runLive(invocationContext)
545-
.doOnNext(event -> this.sessionService.appendEvent(session, event))
546-
.onErrorResumeNext(
547-
throwable -> {
548-
span.setStatus(StatusCode.ERROR, "Error in runLive Flowable execution");
549-
span.recordException(throwable);
550-
span.end();
551-
return Flowable.error(throwable);
552-
}));
551+
552+
return invocationContextSingle.flatMapPublisher(
553+
updatedInvocationContext ->
554+
Telemetry.traceFlowable(
555+
spanContext,
556+
span,
557+
() ->
558+
updatedInvocationContext
559+
.agent()
560+
.runLive(updatedInvocationContext)
561+
.doOnNext(event -> this.sessionService.appendEvent(session, event))
562+
.onErrorResumeNext(
563+
throwable -> {
564+
span.setStatus(
565+
StatusCode.ERROR, "Error in runLive Flowable execution");
566+
span.recordException(throwable);
567+
span.end();
568+
return Flowable.error(throwable);
569+
})));
553570
} catch (Throwable t) {
554571
span.setStatus(StatusCode.ERROR, "Error during runLive synchronous setup");
555572
span.recordException(t);

core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -258,8 +258,8 @@ public void fromConfig_withBuiltInTool_loadsTool() throws IOException, Configura
258258

259259
assertThat(agent).isInstanceOf(LlmAgent.class);
260260
LlmAgent llmAgent = (LlmAgent) agent;
261-
assertThat(llmAgent.tools()).hasSize(1);
262-
assertThat(llmAgent.tools().get(0).name()).isEqualTo("google_search");
261+
assertThat(llmAgent.tools().blockingGet()).hasSize(1);
262+
assertThat(llmAgent.tools().blockingGet().get(0).name()).isEqualTo("google_search");
263263
}
264264

265265
@Test
@@ -784,7 +784,7 @@ public void fromConfig_withIncludeContentsAndOtherFields_parsesAllFieldsCorrectl
784784
assertThat(llmAgent.outputKey()).hasValue("testOutput");
785785
assertThat(llmAgent.disallowTransferToParent()).isTrue();
786786
assertThat(llmAgent.disallowTransferToPeers()).isFalse();
787-
assertThat(llmAgent.tools()).hasSize(1);
787+
assertThat(llmAgent.tools().blockingGet()).hasSize(1);
788788
assertThat(llmAgent.model()).isPresent();
789789
}
790790

0 commit comments

Comments
 (0)