Skip to content

Commit c6c51f1

Browse files
author
Praveen Kumar
committed
refactor: refactor BaseLlmFlow to make code more readable
1 parent af11445 commit c6c51f1

File tree

1 file changed

+37
-35
lines changed

1 file changed

+37
-35
lines changed

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

Lines changed: 37 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -454,45 +454,16 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
454454

455455
String eventIdForSendData = Event.generateEventId();
456456
LlmAgent agent = (LlmAgent) invocationContext.agent();
457-
BaseLlm llm =
458-
agent.resolvedModel().model().isPresent()
459-
? agent.resolvedModel().model().get()
460-
: LlmRegistry.getLlm(agent.resolvedModel().modelName().get());
457+
BaseLlm llm = agent.resolvedModel().model()
458+
.orElse(LlmRegistry.getLlm(agent.resolvedModel().modelName()
459+
.orElseThrow()));
460+
461461
BaseLlmConnection connection = llm.connect(llmRequestAfterPreprocess);
462462
Completable historySent =
463463
llmRequestAfterPreprocess.contents().isEmpty()
464464
? Completable.complete()
465465
: Completable.defer(
466-
() -> {
467-
Span sendDataSpan =
468-
Telemetry.getTracer().spanBuilder("send_data").startSpan();
469-
try (Scope scope = sendDataSpan.makeCurrent()) {
470-
return connection
471-
.sendHistory(llmRequestAfterPreprocess.contents())
472-
.doOnComplete(
473-
() -> {
474-
try (Scope innerScope = sendDataSpan.makeCurrent()) {
475-
Telemetry.traceSendData(
476-
invocationContext,
477-
eventIdForSendData,
478-
llmRequestAfterPreprocess.contents());
479-
}
480-
})
481-
.doOnError(
482-
error -> {
483-
sendDataSpan.setStatus(
484-
StatusCode.ERROR, error.getMessage());
485-
sendDataSpan.recordException(error);
486-
try (Scope innerScope = sendDataSpan.makeCurrent()) {
487-
Telemetry.traceSendData(
488-
invocationContext,
489-
eventIdForSendData,
490-
llmRequestAfterPreprocess.contents());
491-
}
492-
})
493-
.doFinally(sendDataSpan::end);
494-
}
495-
});
466+
() -> sendHistory(invocationContext, connection, llmRequestAfterPreprocess, eventIdForSendData));
496467

497468
Flowable<LiveRequest> liveRequests = invocationContext.liveRequestQueue().get().get();
498469
Disposable sendTask =
@@ -588,7 +559,38 @@ public void onError(Throwable e) {
588559
});
589560
}
590561

591-
/**
562+
private static Completable sendHistory(final InvocationContext invocationContext, final BaseLlmConnection connection, final LlmRequest llmRequestAfterPreprocess, final String eventIdForSendData) {
563+
Span sendDataSpan =
564+
Telemetry.getTracer().spanBuilder("send_data").startSpan();
565+
try (Scope scope = sendDataSpan.makeCurrent()) {
566+
return connection
567+
.sendHistory(llmRequestAfterPreprocess.contents())
568+
.doOnComplete(
569+
() -> {
570+
try (Scope innerScope = sendDataSpan.makeCurrent()) {
571+
Telemetry.traceSendData(
572+
invocationContext,
573+
eventIdForSendData,
574+
llmRequestAfterPreprocess.contents());
575+
}
576+
})
577+
.doOnError(
578+
error -> {
579+
sendDataSpan.setStatus(
580+
StatusCode.ERROR, error.getMessage());
581+
sendDataSpan.recordException(error);
582+
try (Scope innerScope = sendDataSpan.makeCurrent()) {
583+
Telemetry.traceSendData(
584+
invocationContext,
585+
eventIdForSendData,
586+
llmRequestAfterPreprocess.contents());
587+
}
588+
})
589+
.doFinally(sendDataSpan::end);
590+
}
591+
}
592+
593+
/**
592594
* Builds an {@link Event} from LLM response, request, and base event data.
593595
*
594596
* <p>Populates the event with LLM output and tool function call metadata.

0 commit comments

Comments
 (0)