Feature request: Support Planners in ADK Java

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

@@ -51,6 +51,7 @@
 import com.google.adk.models.BaseLlm;
 import com.google.adk.models.LlmRegistry;
 import com.google.adk.models.Model;
+import com.google.adk.planners.BasePlanner;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.BaseTool.ToolArgsConfig;
 import com.google.adk.tools.BaseTool.ToolConfig;
@@ -101,6 +102,7 @@ public enum IncludeContents {
   private final IncludeContents includeContents;
 
   private final boolean planning;
+  private final Optional<BasePlanner> planner;
   private final Optional<Integer> maxSteps;
   private final boolean disallowTransferToParent;
   private final boolean disallowTransferToPeers;
@@ -134,6 +136,7 @@ protected LlmAgent(Builder builder) {
     this.includeContents =
         builder.includeContents != null ? builder.includeContents : IncludeContents.DEFAULT;
     this.planning = builder.planning != null && builder.planning;
+    this.planner = Optional.ofNullable(builder.planner);
     this.maxSteps = Optional.ofNullable(builder.maxSteps);
     this.disallowTransferToParent = builder.disallowTransferToParent;
     this.disallowTransferToPeers = builder.disallowTransferToPeers;
@@ -174,6 +177,7 @@ public static class Builder {
     private BaseExampleProvider exampleProvider;
     private IncludeContents includeContents;
     private Boolean planning;
+    private BasePlanner planner;
     private Integer maxSteps;
     private Boolean disallowTransferToParent;
     private Boolean disallowTransferToPeers;
@@ -310,6 +314,12 @@ public Builder planning(boolean planning) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder planner(BasePlanner planner) {
+      this.planner = planner;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder maxSteps(int maxSteps) {
       this.maxSteps = maxSteps;
@@ -767,6 +777,10 @@ public boolean planning() {
     return planning;
   }
 
+  public Optional<BasePlanner> planner() {
+    return planner;
+  }
+
   public Optional<Integer> maxSteps() {
     return maxSteps;
   }

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

@@ -0,0 +1,166 @@
+package com.google.adk.flows.llmflows;
+
+import com.google.adk.agents.CallbackContext;
+import com.google.adk.agents.InvocationContext;
+import com.google.adk.agents.LlmAgent;
+import com.google.adk.agents.ReadonlyContext;
+import com.google.adk.events.Event;
+import com.google.adk.models.LlmRequest;
+import com.google.adk.models.LlmResponse;
+import com.google.adk.planners.BasePlanner;
+import com.google.adk.planners.BuiltInPlanner;
+import com.google.common.collect.ImmutableList;
+import com.google.genai.types.Content;
+import com.google.genai.types.Part;
+import io.reactivex.rxjava3.core.Single;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+public class NLPlanning {
+
+  static class NlPlanningRequestProcessor implements RequestProcessor {
+
+    @Override
+    public Single<RequestProcessingResult> processRequest(
+        InvocationContext context, LlmRequest llmRequest) {
+
+      if (!(context.agent() instanceof LlmAgent)) {
+        throw new IllegalArgumentException(
+            "Agent in InvocationContext is not an instance of LlmAgent.");
+      }
+
+      Optional<BasePlanner> plannerOpt = getPlanner(context);
+      if (plannerOpt.isEmpty()) {
+        return Single.just(RequestProcessor.RequestProcessingResult.create(llmRequest, ImmutableList.of()));
+      }
+
+      BasePlanner planner = plannerOpt.get();
+
+      // Apply thinking configuration for built-in planners
+      if (planner instanceof BuiltInPlanner) {
+        ((BuiltInPlanner) planner).applyThinkingConfig(llmRequest);
+      }
+
+      // Build and append planning instruction
+      Optional<String> planningInstruction =
+          planner.generatePlanningInstruction(new ReadonlyContext(context), llmRequest);
+
+      LlmRequest.Builder b = llmRequest.toBuilder();
+      planningInstruction.ifPresent(s -> b.appendInstructions(ImmutableList.of(s)));
+      llmRequest = b.build();
+
+      // Remove thought annotations from request
+      llmRequest = removeThoughtFromRequest(llmRequest);
+
+      return Single.just(RequestProcessor.RequestProcessingResult.create(llmRequest, ImmutableList.of()));
+    }
+  }
+
+  static class NlPlanningResponseProcessor implements ResponseProcessor {
+
+    @Override
+    public Single<ResponseProcessingResult> processResponse(
+        InvocationContext context, LlmResponse llmResponse) {
+
+      if (!(context.agent() instanceof LlmAgent)) {
+        throw new IllegalArgumentException(
+            "Agent in InvocationContext is not an instance of LlmAgent.");
+      }
+
+      // Validate response structure
+      if (llmResponse == null || llmResponse.content().isEmpty()) {
+        return Single.just(
+            ResponseProcessor.ResponseProcessingResult.create(
+                llmResponse, ImmutableList.of(), Optional.empty()));
+      }
+
+      Optional<BasePlanner> plannerOpt = getPlanner(context);
+      if (plannerOpt.isEmpty()) {
+        return Single.just(
+            ResponseProcessor.ResponseProcessingResult.create(
+                llmResponse, ImmutableList.of(), Optional.empty()));
+      }
+
+      BasePlanner planner = plannerOpt.get();
+      LlmResponse.Builder responseBuilder = llmResponse.toBuilder();
+
+      // Process the planning response
+      CallbackContext callbackContext = new CallbackContext(context, null);
+      Optional<List<Part>> processedParts =
+          planner.processPlanningResponse(
+              callbackContext, llmResponse.content().get().parts().orElse(List.of()));
+
+      // Update response with processed parts
+      if (processedParts.isPresent()) {
+        Content.Builder contentBuilder = llmResponse.content().get().toBuilder();
+        contentBuilder.parts(processedParts.get());
+        responseBuilder.content(contentBuilder.build());
+      }
+
+      ImmutableList.Builder<Event> eventsBuilder = ImmutableList.builder();
+
+      // Generate state update event if there are deltas
+      if (callbackContext.state().hasDelta()) {
+        Event stateUpdateEvent =
+            Event.builder()
+                .invocationId(context.invocationId())
+                .author(context.agent().name())
+                .branch(context.branch())
+                .actions(callbackContext.eventActions())
+                .build();
+
+        eventsBuilder.add(stateUpdateEvent);
+      }
+
+      return Single.just(
+          ResponseProcessor.ResponseProcessingResult.create(
+              responseBuilder.build(), eventsBuilder.build(), Optional.empty()));
+    }
+  }
+
+  /**
+   * Retrieves the planner from the invocation context.
+   *
+   * @param invocationContext the current invocation context
+   * @return optional planner instance, or empty if none available
+   */
+  private static Optional<BasePlanner> getPlanner(InvocationContext invocationContext) {
+    if (!(invocationContext.agent() instanceof LlmAgent agent)) {
+      return Optional.empty();
+    }
+
+    return agent.planner();
+  }
+
+  /**
+   * Removes thought annotations from all parts in the LLM request.
+   *
+   * <p>This method iterates through all content parts and sets the thought field to false,
+   * effectively removing thought markings from the request.
+   *
+   * @param llmRequest the LLM request to process
+   */
+  private static LlmRequest removeThoughtFromRequest(LlmRequest llmRequest) {
+    if (llmRequest.contents() == null || llmRequest.contents().isEmpty()) {
+      return llmRequest;
+    }
+
+    // Process each content and update its parts
+    List<Content> updatedContents = llmRequest.contents().stream().map(content -> {
+      if (content.parts().isEmpty()) {
+        return content;
+      }
+
+      // Update all parts to set thought to false
+      List<Part> updatedParts = content.parts().get().stream().map(part -> part.toBuilder().thought(false).build()).collect(Collectors.toList());
+
+      // Return updated content with modified parts
+      return content.toBuilder().parts(updatedParts).build();
+    }).collect(Collectors.toList());
+
+    // Return updated LlmRequest with modified contents
+    return llmRequest.toBuilder().contents(updatedContents).build();
+  }
+}

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

@@ -31,10 +31,12 @@ public class SingleFlow extends BaseLlmFlow {
           new Identity(),
           new Contents(),
           new Examples(),
+          new NLPlanning.NlPlanningRequestProcessor(),
           CodeExecution.requestProcessor);
 
   protected static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS =
-      ImmutableList.of(CodeExecution.responseProcessor);
+      ImmutableList.of(
+          new NLPlanning.NlPlanningResponseProcessor(), CodeExecution.responseProcessor);
 
   public SingleFlow() {
     this(/* maxSteps= */ Optional.empty());

core/src/main/java/com/google/adk/planners/BasePlanner.java

@@ -0,0 +1,28 @@
+package com.google.adk.planners;
+
+import com.google.adk.agents.CallbackContext;
+import com.google.adk.agents.ReadonlyContext;
+import com.google.adk.models.LlmRequest;
+import com.google.genai.types.Part;
+import java.util.List;
+import java.util.Optional;
+
+public interface BasePlanner {
+  /**
+   * Generates system instruction text for LLM planning requests.
+   *
+   * @param context readonly invocation context
+   * @param request the LLM request being prepared
+   * @return planning instruction text, or empty if no instruction needed
+   */
+  Optional<String> generatePlanningInstruction(ReadonlyContext context, LlmRequest request);
+
+  /**
+   * Processes and transforms LLM response parts for planning workflow.
+   *
+   * @param context callback context for the current invocation
+   * @param responseParts list of response parts from the LLM
+   * @return processed response parts, or empty if no processing required
+   */
+  Optional<List<Part>> processPlanningResponse(CallbackContext context, List<Part> responseParts);
+}

core/src/main/java/com/google/adk/planners/BuiltInPlanner.java

@@ -0,0 +1,56 @@
+package com.google.adk.planners;
+
+import com.google.adk.agents.CallbackContext;
+import com.google.adk.agents.ReadonlyContext;
+import com.google.adk.models.LlmRequest;
+import com.google.genai.types.GenerateContentConfig;
+import com.google.genai.types.Part;
+import com.google.genai.types.ThinkingConfig;
+import java.util.List;
+import java.util.Optional;
+
+public class BuiltInPlanner implements BasePlanner {
+  private ThinkingConfig cognitiveConfig;
+
+  private BuiltInPlanner() {}
+
+  private BuiltInPlanner(ThinkingConfig cognitiveConfig) {
+    this.cognitiveConfig = cognitiveConfig;
+  }
+
+  public static BuiltInPlanner buildPlanner(ThinkingConfig cognitiveConfig) {
+    return new BuiltInPlanner(cognitiveConfig);
+  }
+
+  @Override
+  public Optional<String> generatePlanningInstruction(ReadonlyContext context, LlmRequest request) {
+    return Optional.empty();
+  }
+
+  @Override
+  public Optional<List<Part>> processPlanningResponse(
+      CallbackContext context, List<Part> responseParts) {
+    return Optional.empty();
+  }
+
+  /**
+   * Configures the LLM request with thinking capabilities. This method modifies the request to
+   * include the thinking configuration, enabling the model's native cognitive processing features.
+   *
+   * @param request the LLM request to configure
+   */
+  public LlmRequest applyThinkingConfig(LlmRequest request) {
+    if (this.cognitiveConfig != null) {
+      // Ensure config exists
+      GenerateContentConfig.Builder configBuilder =
+          request.config().map(GenerateContentConfig::toBuilder).orElse(GenerateContentConfig.builder());
+
+      // Apply thinking configuration
+      request =
+          request.toBuilder()
+              .config(configBuilder.thinkingConfig(this.cognitiveConfig).build())
+              .build();
+    }
+    return request;
+  }
+}