Skip to content

Commit 8a1fffa

Browse files
google-genai-botcopybara-github
authored andcommitted
refactor: Code clean up in BaseAgent
1. **Extracting Common Run Logic:** A new private method `run` is introduced to encapsulate the common setup and execution flow shared by both `runAsync` and `runLive`. This includes creating the `InvocationContext`, setting up tracing, and handling before/after callbacks. 2. **Reducing Duplication:** Both `runAsync` and `runLive` now call this new `run` method, passing their specific implementation (`runAsyncImpl` or `runLiveImpl`) as a function. This eliminates a significant amount of duplicated code. 3. **Minor Cleanups:** * The `createInvocationContext` method now uses `parentContext.branch().filter(...).ifPresent(...)` for a more concise optional handling. * An unnecessary `agentCallbacks == null` check is removed from `callCallback`, as the list should not be null. PiperOrigin-RevId: 868646905
1 parent a364e30 commit 8a1fffa

1 file changed

Lines changed: 20 additions & 44 deletions

File tree

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

Lines changed: 20 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -288,9 +288,10 @@ private InvocationContext createInvocationContext(InvocationContext parentContex
288288
InvocationContext.Builder builder = parentContext.toBuilder();
289289
builder.agent(this);
290290
// Check for branch to be truthy (not None, not empty string),
291-
if (parentContext.branch().filter(s -> !s.isEmpty()).isPresent()) {
292-
builder.branch(parentContext.branch().get() + "." + name());
293-
}
291+
parentContext
292+
.branch()
293+
.filter(s -> !s.isEmpty())
294+
.ifPresent(branch -> builder.branch(branch + "." + name()));
294295
return builder.build();
295296
}
296297

@@ -301,6 +302,19 @@ private InvocationContext createInvocationContext(InvocationContext parentContex
301302
* @return stream of agent-generated events.
302303
*/
303304
public Flowable<Event> runAsync(InvocationContext parentContext) {
305+
return run(parentContext, this::runAsyncImpl);
306+
}
307+
308+
/**
309+
* Runs the agent with the given implementation.
310+
*
311+
* @param parentContext Parent context to inherit.
312+
* @param runImplementation The agent-specific logic to run.
313+
* @return stream of agent-generated events.
314+
*/
315+
private Flowable<Event> run(
316+
InvocationContext parentContext,
317+
Function<InvocationContext, Flowable<Event>> runImplementation) {
304318
Tracer tracer = Tracing.getTracer();
305319
return Flowable.defer(
306320
() -> {
@@ -326,7 +340,7 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
326340

327341
Flowable<Event> beforeEvents = Flowable.fromOptional(beforeEventOpt);
328342
Flowable<Event> mainEvents =
329-
Flowable.defer(() -> runAsyncImpl(invocationContext));
343+
Flowable.defer(() -> runImplementation.apply(invocationContext));
330344
Flowable<Event> afterEvents =
331345
Flowable.defer(
332346
() ->
@@ -382,7 +396,7 @@ private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksT
382396
private Single<Optional<Event>> callCallback(
383397
List<Function<CallbackContext, Maybe<Content>>> agentCallbacks,
384398
InvocationContext invocationContext) {
385-
if (agentCallbacks == null || agentCallbacks.isEmpty()) {
399+
if (agentCallbacks.isEmpty()) {
386400
return Single.just(Optional.empty());
387401
}
388402

@@ -437,45 +451,7 @@ private Single<Optional<Event>> callCallback(
437451
* @return stream of agent-generated events.
438452
*/
439453
public Flowable<Event> runLive(InvocationContext parentContext) {
440-
Tracer tracer = Tracing.getTracer();
441-
return Flowable.defer(
442-
() -> {
443-
InvocationContext invocationContext = createInvocationContext(parentContext);
444-
Span span =
445-
tracer.spanBuilder("invoke_agent " + name()).setParent(Context.current()).startSpan();
446-
Tracing.traceAgentInvocation(span, name(), description(), invocationContext);
447-
Context spanContext = Context.current().with(span);
448-
449-
return Tracing.traceFlowable(
450-
spanContext,
451-
span,
452-
() ->
453-
callCallback(
454-
beforeCallbacksToFunctions(
455-
invocationContext.pluginManager(), beforeAgentCallback),
456-
invocationContext)
457-
.flatMapPublisher(
458-
beforeEventOpt -> {
459-
if (invocationContext.endInvocation()) {
460-
return Flowable.fromOptional(beforeEventOpt);
461-
}
462-
463-
Flowable<Event> beforeEvents = Flowable.fromOptional(beforeEventOpt);
464-
Flowable<Event> mainEvents =
465-
Flowable.defer(() -> runLiveImpl(invocationContext));
466-
Flowable<Event> afterEvents =
467-
Flowable.defer(
468-
() ->
469-
callCallback(
470-
afterCallbacksToFunctions(
471-
invocationContext.pluginManager(),
472-
afterAgentCallback),
473-
invocationContext)
474-
.flatMapPublisher(Flowable::fromOptional));
475-
476-
return Flowable.concat(beforeEvents, mainEvents, afterEvents);
477-
}));
478-
});
454+
return run(parentContext, this::runLiveImpl);
479455
}
480456

481457
/**

0 commit comments

Comments
 (0)