Skip to content

Commit ea2b95e

Browse files
tilgalascopybara-github
authored andcommitted
fix: include usage_metadata events in live postprocessing
PiperOrigin-RevId: 870945434
1 parent fe00ef8 commit ea2b95e

File tree

2 files changed

+33
-1
lines changed

2 files changed

+33
-1
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -611,7 +611,8 @@ private Flowable<Event> buildPostprocessingEvents(
611611
if (updatedResponse.content().isEmpty()
612612
&& updatedResponse.errorCode().isEmpty()
613613
&& !updatedResponse.interrupted().orElse(false)
614-
&& !updatedResponse.turnComplete().orElse(false)) {
614+
&& !updatedResponse.turnComplete().orElse(false)
615+
&& updatedResponse.usageMetadata().isEmpty()) {
615616
return processorEvents;
616617
}
617618

core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -575,4 +575,35 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
575575
return Single.just(response);
576576
}
577577
}
578+
579+
@Test
580+
public void postprocess_noResponseProcessors_onlyUsageMetadata_returnsEvent() {
581+
GenerateContentResponseUsageMetadata usageMetadata =
582+
GenerateContentResponseUsageMetadata.builder()
583+
.promptTokenCount(10)
584+
.candidatesTokenCount(20)
585+
.build();
586+
LlmResponse llmResponse = LlmResponse.builder().usageMetadata(usageMetadata).build();
587+
InvocationContext invocationContext =
588+
createInvocationContext(createTestAgent(createTestLlm(llmResponse)));
589+
BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors();
590+
Event baseEvent =
591+
Event.builder()
592+
.invocationId(invocationContext.invocationId())
593+
.author(invocationContext.agent().name())
594+
.build();
595+
596+
List<Event> events =
597+
baseLlmFlow
598+
.postprocess(invocationContext, baseEvent, LlmRequest.builder().build(), llmResponse)
599+
.toList()
600+
.blockingGet();
601+
602+
assertThat(events).hasSize(1);
603+
Event event = getOnlyElement(events);
604+
assertThat(event.content()).isEmpty();
605+
assertThat(event.usageMetadata()).hasValue(usageMetadata);
606+
assertThat(event.author()).isEqualTo(invocationContext.agent().name());
607+
assertThat(event.invocationId()).isEqualTo(invocationContext.invocationId());
608+
}
578609
}

0 commit comments

Comments
 (0)