-
Notifications
You must be signed in to change notification settings - Fork 152
Feature request: Support Planners in ADK Java #395
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
leweii
wants to merge
9
commits into
google:main
Choose a base branch
from
leweii:feat/planner-support
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
d45af69
Move the planning_re_act_planner python implementation to java version.
leweii 5c3ac37
Move the planning_re_act_planner python implementation to java version.
leweii ef2908a
Update core/src/main/java/com/google/adk/flows/llmflows/NLPlanning.java
leweii ea57eb5
Merge branch 'main' into feat/planner-support
leweii d0b62c2
Merge branch 'main' into feat/planner-support
leweii 5df8bcf
Merge branch 'main' into feat/planner-support
leweii 9e85079
Merge branch 'main' into feat/planner-support
leweii 322378e
Merge branch 'main' into feat/planner-support
leweii 80bc3c7
Merge branch 'main' into feat/planner-support
leweii File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
166 changes: 166 additions & 0 deletions
166
core/src/main/java/com/google/adk/flows/llmflows/NLPlanning.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) { | ||
llmRequest = ((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(); | ||
} | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
28 changes: 28 additions & 0 deletions
28
core/src/main/java/com/google/adk/planners/BasePlanner.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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); | ||
} |
56 changes: 56 additions & 0 deletions
56
core/src/main/java/com/google/adk/planners/BuiltInPlanner.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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; | ||
} | ||
} |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The check for
invocationContext.agent() instanceof LlmAgent
is redundant here. The callers of this method,NlPlanningRequestProcessor.processRequest
andNlPlanningResponseProcessor.processResponse
, already perform this validation and throw anIllegalArgumentException
if the agent is not anLlmAgent
. You can simplify this method by removing the redundant check and casting directly.