Skip to content

Commit 01ddc6e

Browse files
leweiileweii
committed
Move the planning_re_act_planner python implementation to java version.
1 parent 2bfc95d commit 01ddc6e

File tree

8 files changed

+544
-1
lines changed

8 files changed

+544
-1
lines changed

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import com.google.adk.models.BaseLlm;
5252
import com.google.adk.models.LlmRegistry;
5353
import com.google.adk.models.Model;
54+
import com.google.adk.planners.BasePlanner;
5455
import com.google.adk.tools.BaseTool;
5556
import com.google.adk.tools.BaseToolset;
5657
import com.google.adk.utils.ComponentRegistry;
@@ -98,6 +99,7 @@ public enum IncludeContents {
9899
private final IncludeContents includeContents;
99100

100101
private final boolean planning;
102+
private final Optional<BasePlanner> planner;
101103
private final Optional<Integer> maxSteps;
102104
private final boolean disallowTransferToParent;
103105
private final boolean disallowTransferToPeers;
@@ -131,6 +133,7 @@ protected LlmAgent(Builder builder) {
131133
this.includeContents =
132134
builder.includeContents != null ? builder.includeContents : IncludeContents.DEFAULT;
133135
this.planning = builder.planning != null && builder.planning;
136+
this.planner = Optional.ofNullable(builder.planner);
134137
this.maxSteps = Optional.ofNullable(builder.maxSteps);
135138
this.disallowTransferToParent = builder.disallowTransferToParent;
136139
this.disallowTransferToPeers = builder.disallowTransferToPeers;
@@ -180,6 +183,7 @@ public static class Builder {
180183
private BaseExampleProvider exampleProvider;
181184
private IncludeContents includeContents;
182185
private Boolean planning;
186+
private BasePlanner planner;
183187
private Integer maxSteps;
184188
private Boolean disallowTransferToParent;
185189
private Boolean disallowTransferToPeers;
@@ -306,6 +310,12 @@ public Builder planning(boolean planning) {
306310
return this;
307311
}
308312

313+
@CanIgnoreReturnValue
314+
public Builder planner(BasePlanner planner) {
315+
this.planner = planner;
316+
return this;
317+
}
318+
309319
@CanIgnoreReturnValue
310320
public Builder maxSteps(int maxSteps) {
311321
this.maxSteps = maxSteps;
@@ -764,6 +774,10 @@ public boolean planning() {
764774
return planning;
765775
}
766776

777+
public Optional<BasePlanner> planner() {
778+
return planner;
779+
}
780+
767781
public Optional<Integer> maxSteps() {
768782
return maxSteps;
769783
}

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

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
package com.google.adk.flows.llmflows;
2+
3+
import com.google.adk.agents.CallbackContext;
4+
import com.google.adk.agents.InvocationContext;
5+
import com.google.adk.agents.LlmAgent;
6+
import com.google.adk.agents.ReadonlyContext;
7+
import com.google.adk.events.Event;
8+
import com.google.adk.models.LlmRequest;
9+
import com.google.adk.models.LlmResponse;
10+
import com.google.adk.planners.BasePlanner;
11+
import com.google.adk.planners.BuiltInPlanner;
12+
import com.google.common.collect.ImmutableList;
13+
import com.google.genai.types.Content;
14+
import com.google.genai.types.Part;
15+
import io.reactivex.rxjava3.core.Single;
16+
17+
import java.util.List;
18+
import java.util.Optional;
19+
import java.util.stream.Collectors;
20+
21+
public class NLPlanning {
22+
23+
static class NlPlanningRequestProcessor implements RequestProcessor {
24+
25+
@Override
26+
public Single<RequestProcessingResult> processRequest(
27+
InvocationContext context, LlmRequest llmRequest) {
28+
29+
if (!(context.agent() instanceof LlmAgent)) {
30+
throw new IllegalArgumentException(
31+
"Agent in InvocationContext is not an instance of LlmAgent.");
32+
}
33+
34+
Optional<BasePlanner> plannerOpt = getPlanner(context);
35+
if (plannerOpt.isEmpty()) {
36+
return Single.just(RequestProcessor.RequestProcessingResult.create(llmRequest, ImmutableList.of()));
37+
}
38+
39+
BasePlanner planner = plannerOpt.get();
40+
41+
// Apply thinking configuration for built-in planners
42+
if (planner instanceof BuiltInPlanner) {
43+
llmRequest = ((BuiltInPlanner) planner).applyThinkingConfig(llmRequest);
44+
}
45+
46+
// Build and append planning instruction
47+
Optional<String> planningInstruction =
48+
planner.generatePlanningInstruction(new ReadonlyContext(context), llmRequest);
49+
50+
LlmRequest.Builder b = llmRequest.toBuilder();
51+
planningInstruction.ifPresent(s -> b.appendInstructions(ImmutableList.of(s)));
52+
llmRequest = b.build();
53+
54+
// Remove thought annotations from request
55+
llmRequest = removeThoughtFromRequest(llmRequest);
56+
57+
return Single.just(RequestProcessor.RequestProcessingResult.create(llmRequest, ImmutableList.of()));
58+
}
59+
}
60+
61+
static class NlPlanningResponseProcessor implements ResponseProcessor {
62+
63+
@Override
64+
public Single<ResponseProcessingResult> processResponse(
65+
InvocationContext context, LlmResponse llmResponse) {
66+
67+
if (!(context.agent() instanceof LlmAgent)) {
68+
throw new IllegalArgumentException(
69+
"Agent in InvocationContext is not an instance of LlmAgent.");
70+
}
71+
72+
// Validate response structure
73+
if (llmResponse == null || llmResponse.content().isEmpty()) {
74+
return Single.just(
75+
ResponseProcessor.ResponseProcessingResult.create(
76+
llmResponse, ImmutableList.of(), Optional.empty()));
77+
}
78+
79+
Optional<BasePlanner> plannerOpt = getPlanner(context);
80+
if (plannerOpt.isEmpty()) {
81+
return Single.just(
82+
ResponseProcessor.ResponseProcessingResult.create(
83+
llmResponse, ImmutableList.of(), Optional.empty()));
84+
}
85+
86+
BasePlanner planner = plannerOpt.get();
87+
LlmResponse.Builder responseBuilder = llmResponse.toBuilder();
88+
89+
// Process the planning response
90+
CallbackContext callbackContext = new CallbackContext(context, null);
91+
Optional<List<Part>> processedParts =
92+
planner.processPlanningResponse(
93+
callbackContext, llmResponse.content().get().parts().orElse(List.of()));
94+
95+
// Update response with processed parts
96+
if (processedParts.isPresent()) {
97+
Content.Builder contentBuilder = llmResponse.content().get().toBuilder();
98+
contentBuilder.parts(processedParts.get());
99+
responseBuilder.content(contentBuilder.build());
100+
}
101+
102+
ImmutableList.Builder<Event> eventsBuilder = ImmutableList.builder();
103+
104+
// Generate state update event if there are deltas
105+
if (callbackContext.state().hasDelta()) {
106+
Event stateUpdateEvent =
107+
Event.builder()
108+
.invocationId(context.invocationId())
109+
.author(context.agent().name())
110+
.branch(context.branch())
111+
.actions(callbackContext.eventActions())
112+
.build();
113+
114+
eventsBuilder.add(stateUpdateEvent);
115+
}
116+
117+
return Single.just(
118+
ResponseProcessor.ResponseProcessingResult.create(
119+
responseBuilder.build(), eventsBuilder.build(), Optional.empty()));
120+
}
121+
}
122+
123+
/**
124+
* Retrieves the planner from the invocation context.
125+
*
126+
* @param invocationContext the current invocation context
127+
* @return optional planner instance, or empty if none available
128+
*/
129+
private static Optional<BasePlanner> getPlanner(InvocationContext invocationContext) {
130+
if (!(invocationContext.agent() instanceof LlmAgent agent)) {
131+
return Optional.empty();
132+
}
133+
134+
return agent.planner();
135+
}
136+
137+
/**
138+
* Removes thought annotations from all parts in the LLM request.
139+
*
140+
* <p>This method iterates through all content parts and sets the thought field to false,
141+
* effectively removing thought markings from the request.
142+
*
143+
* @param llmRequest the LLM request to process
144+
*/
145+
private static LlmRequest removeThoughtFromRequest(LlmRequest llmRequest) {
146+
if (llmRequest.contents() == null || llmRequest.contents().isEmpty()) {
147+
return llmRequest;
148+
}
149+
150+
// Process each content and update its parts
151+
List<Content> updatedContents = llmRequest.contents().stream().map(content -> {
152+
if (content.parts().isEmpty()) {
153+
return content;
154+
}
155+
156+
// Update all parts to set thought to false
157+
List<Part> updatedParts = content.parts().get().stream().map(part -> part.toBuilder().thought(false).build()).collect(Collectors.toList());
158+
159+
// Return updated content with modified parts
160+
return content.toBuilder().parts(updatedParts).build();
161+
}).collect(Collectors.toList());
162+
163+
// Return updated LlmRequest with modified contents
164+
return llmRequest.toBuilder().contents(updatedContents).build();
165+
}
166+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,12 @@ public class SingleFlow extends BaseLlmFlow {
3131
new Identity(),
3232
new Contents(),
3333
new Examples(),
34+
new NLPlanning.NlPlanningRequestProcessor(),
3435
CodeExecution.requestProcessor);
3536

3637
protected static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS =
37-
ImmutableList.of(CodeExecution.responseProcessor);
38+
ImmutableList.of(
39+
new NLPlanning.NlPlanningResponseProcessor(), CodeExecution.responseProcessor);
3840

3941
public SingleFlow() {
4042
this(/* maxSteps= */ Optional.empty());

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package com.google.adk.planners;
2+
3+
import com.google.adk.agents.CallbackContext;
4+
import com.google.adk.agents.ReadonlyContext;
5+
import com.google.adk.models.LlmRequest;
6+
import com.google.genai.types.Part;
7+
import java.util.List;
8+
import java.util.Optional;
9+
10+
public interface BasePlanner {
11+
/**
12+
* Generates system instruction text for LLM planning requests.
13+
*
14+
* @param context readonly invocation context
15+
* @param request the LLM request being prepared
16+
* @return planning instruction text, or empty if no instruction needed
17+
*/
18+
Optional<String> generatePlanningInstruction(ReadonlyContext context, LlmRequest request);
19+
20+
/**
21+
* Processes and transforms LLM response parts for planning workflow.
22+
*
23+
* @param context callback context for the current invocation
24+
* @param responseParts list of response parts from the LLM
25+
* @return processed response parts, or empty if no processing required
26+
*/
27+
Optional<List<Part>> processPlanningResponse(CallbackContext context, List<Part> responseParts);
28+
}

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package com.google.adk.planners;
2+
3+
import com.google.adk.agents.CallbackContext;
4+
import com.google.adk.agents.ReadonlyContext;
5+
import com.google.adk.models.LlmRequest;
6+
import com.google.genai.types.GenerateContentConfig;
7+
import com.google.genai.types.Part;
8+
import com.google.genai.types.ThinkingConfig;
9+
import java.util.List;
10+
import java.util.Optional;
11+
12+
public class BuiltInPlanner implements BasePlanner {
13+
private ThinkingConfig cognitiveConfig;
14+
15+
private BuiltInPlanner() {}
16+
17+
private BuiltInPlanner(ThinkingConfig cognitiveConfig) {
18+
this.cognitiveConfig = cognitiveConfig;
19+
}
20+
21+
public static BuiltInPlanner buildPlanner(ThinkingConfig cognitiveConfig) {
22+
return new BuiltInPlanner(cognitiveConfig);
23+
}
24+
25+
@Override
26+
public Optional<String> generatePlanningInstruction(ReadonlyContext context, LlmRequest request) {
27+
return Optional.empty();
28+
}
29+
30+
@Override
31+
public Optional<List<Part>> processPlanningResponse(
32+
CallbackContext context, List<Part> responseParts) {
33+
return Optional.empty();
34+
}
35+
36+
/**
37+
* Configures the LLM request with thinking capabilities. This method modifies the request to
38+
* include the thinking configuration, enabling the model's native cognitive processing features.
39+
*
40+
* @param request the LLM request to configure
41+
*/
42+
public LlmRequest applyThinkingConfig(LlmRequest request) {
43+
if (this.cognitiveConfig != null) {
44+
// Ensure config exists
45+
GenerateContentConfig.Builder configBuilder =
46+
request.config().map(GenerateContentConfig::toBuilder).orElse(GenerateContentConfig.builder());
47+
48+
// Apply thinking configuration
49+
request =
50+
request.toBuilder()
51+
.config(configBuilder.thinkingConfig(this.cognitiveConfig).build())
52+
.build();
53+
}
54+
return request;
55+
}
56+
}

0 commit comments

Comments
 (0)