Skip to content

Commit 4be33a9

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add tools and toolset to use SkillSource in ADK agents
This change introduces a set of tools and a toolset to enable ADK agents to interact with and use skills loaded via `SkillSource`. Key changes: - ListSkillsTool: Lists available skills. - LoadSkillTool: Loads skill instructions. - LoadSkillResourceTool: Loads skill resources. - SkillToolset: Groups the above tools. - Integrates SkillSource into LlmAgent and BaseLlmFlow. - Tests for all new tools. PiperOrigin-RevId: 920071248
1 parent 5ee51fd commit 4be33a9

20 files changed

Lines changed: 1686 additions & 38 deletions

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

Lines changed: 48 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@
3838
import com.google.adk.models.LlmRequest;
3939
import com.google.adk.models.LlmResponse;
4040
import com.google.adk.telemetry.Tracing;
41+
import com.google.adk.tools.BaseTool;
42+
import com.google.adk.tools.BaseToolset;
43+
import com.google.adk.tools.LlmRequestProcessor;
4144
import com.google.adk.tools.ToolContext;
4245
import com.google.common.collect.ImmutableList;
4346
import com.google.common.collect.Iterables;
@@ -58,6 +61,7 @@
5861
import java.util.Optional;
5962
import java.util.Set;
6063
import java.util.concurrent.atomic.AtomicReference;
64+
import java.util.function.BiFunction;
6165
import org.slf4j.Logger;
6266
import org.slf4j.LoggerFactory;
6367

@@ -96,20 +100,8 @@ private Flowable<Event> preprocess(
96100
Context currentContext = Context.current();
97101
LlmAgent agent = (LlmAgent) context.agent();
98102

99-
RequestProcessor toolsProcessor =
100-
(ctx, req) -> {
101-
LlmRequest.Builder builder = req.toBuilder();
102-
return agent
103-
.canonicalTools(new ReadonlyContext(ctx))
104-
.concatMapCompletable(
105-
tool -> tool.processLlmRequest(builder, ToolContext.builder(ctx).build()))
106-
.andThen(
107-
Single.fromCallable(
108-
() -> RequestProcessingResult.create(builder.build(), ImmutableList.of())));
109-
};
110-
111103
Iterable<RequestProcessor> allProcessors =
112-
Iterables.concat(requestProcessors, ImmutableList.of(toolsProcessor));
104+
Iterables.concat(requestProcessors, ImmutableList.of(getRequestProcessorFromTools(agent)));
113105

114106
return Flowable.fromIterable(allProcessors)
115107
.concatMap(
@@ -121,6 +113,49 @@ private Flowable<Event> preprocess(
121113
result -> result.events() != null ? result.events() : ImmutableList.of()));
122114
}
123115

116+
/**
117+
* Constructs a {@link RequestProcessor} that sequentially applies the {@code processLlmRequest}
118+
* methods of all tools and toolsets associated with this agent to the incoming {@link
119+
* LlmRequest}.
120+
*
121+
* @return A {@link RequestProcessor} that applies tool-specific modifications to LLM requests.
122+
*/
123+
private RequestProcessor getRequestProcessorFromTools(LlmAgent agent) {
124+
return (context, request) -> {
125+
ReadonlyContext readonlyContext = new ReadonlyContext(context);
126+
List<BiFunction<LlmRequest.Builder, ToolContext, Completable>> processors = new ArrayList<>();
127+
128+
for (Object toolOrToolset : agent.toolsUnion()) {
129+
if (toolOrToolset instanceof BaseTool baseTool) {
130+
processors.add(baseTool::processLlmRequest);
131+
} else if (toolOrToolset instanceof BaseToolset baseToolset) {
132+
// First apply the toolset's own request processor if it implements the
133+
// LlmRequestProcessor interface, then unwrap all tools from the toolset
134+
// and apply each individual tool's request processor sequentially.
135+
processors.add(
136+
(builder, ctx) ->
137+
(baseToolset instanceof LlmRequestProcessor processor
138+
? processor.processLlmRequest(builder, ctx)
139+
: Completable.complete())
140+
.andThen(baseToolset.getTools(readonlyContext))
141+
.concatMapCompletable(b -> b.processLlmRequest(builder, ctx)));
142+
} else {
143+
throw new IllegalArgumentException(
144+
"Object in tools list is not of a supported type: "
145+
+ toolOrToolset.getClass().getName());
146+
}
147+
}
148+
149+
LlmRequest.Builder builder = request.toBuilder();
150+
ToolContext toolContext = ToolContext.builder(context).build();
151+
return Flowable.fromIterable(processors)
152+
.concatMapCompletable(f -> f.apply(builder, toolContext))
153+
.andThen(
154+
Single.fromCallable(
155+
() -> RequestProcessingResult.create(builder.build(), ImmutableList.of())));
156+
};
157+
}
158+
124159
/**
125160
* Post-processes the LLM response after receiving it from the LLM. Executes all registered {@link
126161
* ResponseProcessor} instances. Emits events for the model response and any subsequent function

core/src/main/java/com/google/adk/skills/AbstractSkillSource.java

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package com.google.adk.skills;
1818

19+
import static com.google.adk.skills.SkillSourceException.SKILL_FORMAT_ERROR;
20+
import static com.google.adk.skills.SkillSourceException.SKILL_LOAD_ERROR;
1921
import static java.nio.channels.Channels.newReader;
2022
import static java.nio.charset.StandardCharsets.UTF_8;
2123

@@ -82,12 +84,14 @@ private Frontmatter loadFrontmatter(String skillName, PathT skillMdPath)
8284
Frontmatter frontmatter = yamlMapper.readValue(yaml, Frontmatter.class);
8385
if (!frontmatter.name().equals(skillName)) {
8486
throw new SkillSourceException(
85-
"Skill name '%s' does not match directory name '%s'."
86-
.formatted(frontmatter.name(), skillName));
87+
"Skill name in the frontmatter '%s' does not match skill name '%s'."
88+
.formatted(frontmatter.name(), skillName),
89+
SKILL_LOAD_ERROR);
8790
}
8891
return frontmatter;
8992
} catch (IOException e) {
90-
throw new SkillSourceException("Cannot load frontmatter for skill '" + skillName + "'", e);
93+
throw new SkillSourceException(
94+
"Cannot load frontmatter for skill '" + skillName + "'", SKILL_LOAD_ERROR, e);
9195
}
9296
}
9397

@@ -100,7 +104,9 @@ public Single<String> loadInstructions(String skillName) {
100104
return readInstructions(reader);
101105
} catch (IOException e) {
102106
throw new SkillSourceException(
103-
"Failed to load instruction for skill '" + skillName + "'", e);
107+
"Failed to load instruction for skill '" + skillName + "'",
108+
SKILL_LOAD_ERROR,
109+
e);
104110
}
105111
});
106112
}
@@ -140,7 +146,8 @@ private String readFrontmatterYaml(BufferedReader reader)
140146
throws IOException, SkillSourceException {
141147
String line = reader.readLine();
142148
if (line == null || !line.trim().equals(THREE_DASHES)) {
143-
throw new SkillSourceException("Skill file must start with " + THREE_DASHES);
149+
throw new SkillSourceException(
150+
"Skill file must start with " + THREE_DASHES, SKILL_FORMAT_ERROR);
144151
}
145152

146153
StringBuilder sb = new StringBuilder();
@@ -151,14 +158,15 @@ private String readFrontmatterYaml(BufferedReader reader)
151158
sb.append(line).append("\n");
152159
}
153160
throw new SkillSourceException(
154-
"Skill file frontmatter not properly closed with " + THREE_DASHES);
161+
"Skill file frontmatter not properly closed with " + THREE_DASHES, SKILL_FORMAT_ERROR);
155162
}
156163

157164
private String readInstructions(BufferedReader reader) throws IOException, SkillSourceException {
158165
// Skip the frontmatter block
159166
String line = reader.readLine();
160167
if (line == null || !line.trim().equals(THREE_DASHES)) {
161-
throw new SkillSourceException("Skill file must start with " + THREE_DASHES);
168+
throw new SkillSourceException(
169+
"Skill file must start with " + THREE_DASHES, SKILL_FORMAT_ERROR);
162170
}
163171
boolean dashClosed = false;
164172
while ((line = reader.readLine()) != null) {
@@ -169,7 +177,7 @@ private String readInstructions(BufferedReader reader) throws IOException, Skill
169177
}
170178
if (!dashClosed) {
171179
throw new SkillSourceException(
172-
"Skill file frontmatter not properly closed with " + THREE_DASHES);
180+
"Skill file frontmatter not properly closed with " + THREE_DASHES, SKILL_FORMAT_ERROR);
173181
}
174182
// Read the instructions till the end of the file
175183
StringBuilder sb = new StringBuilder();

core/src/main/java/com/google/adk/skills/InMemorySkillSource.java

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package com.google.adk.skills;
1818

19+
import static com.google.adk.skills.SkillSourceException.RESOURCE_NOT_FOUND;
20+
import static com.google.adk.skills.SkillSourceException.SKILL_NOT_FOUND;
1921
import static com.google.common.base.Preconditions.checkState;
2022
import static com.google.common.collect.ImmutableList.toImmutableList;
2123
import static java.nio.charset.StandardCharsets.UTF_8;
@@ -56,7 +58,8 @@ public Single<ImmutableMap<String, Frontmatter>> listFrontmatters() {
5658
public Single<ImmutableList<String>> listResources(String skillName, String resourceDirectory) {
5759
SkillData data = skills.get(skillName);
5860
if (data == null) {
59-
return Single.error(new SkillSourceException("Skill not found: " + skillName));
61+
return Single.error(
62+
new SkillSourceException("Skill not found: " + skillName, SKILL_NOT_FOUND));
6063
}
6164
String prefix =
6265
resourceDirectory.isEmpty()
@@ -67,7 +70,8 @@ public Single<ImmutableList<String>> listResources(String skillName, String reso
6770
&& data.resources().keySet().stream().noneMatch(path -> path.startsWith(prefix))) {
6871
return Single.error(
6972
new SkillSourceException(
70-
"Resource directory not found: " + resourceDirectory + " for skill: " + skillName));
73+
"Resource directory not found: " + resourceDirectory + " for skill: " + skillName,
74+
RESOURCE_NOT_FOUND));
7175
}
7276

7377
return Single.just(
@@ -92,13 +96,16 @@ public Single<ByteSource> loadResource(String skillName, String resourcePath) {
9296
.map(SkillData::resources)
9397
.mapOptional(m -> Optional.ofNullable(m.get(resourcePath)))
9498
.switchIfEmpty(
95-
Single.error(new SkillSourceException("Resource not found: " + resourcePath)));
99+
Single.error(
100+
new SkillSourceException(
101+
"Resource not found: " + resourcePath, RESOURCE_NOT_FOUND)));
96102
}
97103

98104
private Single<SkillData> getSkillData(String skillName) {
99105
SkillData data = skills.get(skillName);
100106
if (data == null) {
101-
return Single.error(new SkillSourceException("Skill not found: " + skillName));
107+
return Single.error(
108+
new SkillSourceException("Skill not found: " + skillName, SKILL_NOT_FOUND));
102109
}
103110
return Single.just(data);
104111
}

core/src/main/java/com/google/adk/skills/LocalSkillSource.java

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616

1717
package com.google.adk.skills;
1818

19+
import static com.google.adk.skills.SkillSourceException.RESOURCE_LOAD_ERROR;
20+
import static com.google.adk.skills.SkillSourceException.RESOURCE_NOT_FOUND;
21+
import static com.google.adk.skills.SkillSourceException.SKILL_LOAD_ERROR;
22+
import static com.google.adk.skills.SkillSourceException.SKILL_NOT_FOUND;
1923
import static com.google.common.collect.ImmutableList.toImmutableList;
2024
import static java.nio.file.Files.isDirectory;
2125

@@ -43,14 +47,16 @@ public LocalSkillSource(Path skillsBasePath) {
4347
public Single<ImmutableList<String>> listResources(String skillName, String resourceDirectory) {
4448
Path skillDir = skillsBasePath.resolve(skillName);
4549
if (!isDirectory(skillDir)) {
46-
return Single.error(new SkillSourceException("Skill not found: " + skillName));
50+
return Single.error(
51+
new SkillSourceException("Skill not found: " + skillName, SKILL_NOT_FOUND));
4752
}
4853
Path resourceDir = skillDir.resolve(resourceDirectory);
4954
if (!isDirectory(resourceDir)) {
5055
return Single.error(
5156
new SkillSourceException(
5257
"Resource directory '%s' not found for skill '%s'"
53-
.formatted(resourceDirectory, skillName)));
58+
.formatted(resourceDirectory, skillName),
59+
RESOURCE_NOT_FOUND));
5460
}
5561

5662
return Single.fromCallable(
@@ -67,7 +73,9 @@ public Single<ImmutableList<String>> listResources(String skillName, String reso
6773
t ->
6874
Single.error(
6975
new SkillSourceException(
70-
"Failed to traverse resource directory: " + resourceDirectory, t)));
76+
"Failed to traverse resource directory: " + resourceDirectory,
77+
RESOURCE_LOAD_ERROR,
78+
t)));
7179
}
7280

7381
@Override
@@ -78,7 +86,9 @@ protected Flowable<SkillMdPath> listSkills() {
7886
t ->
7987
Flowable.error(
8088
new SkillSourceException(
81-
"Failed to list skills in directory: " + skillsBasePath, t)))
89+
"Failed to list skills in directory: " + skillsBasePath,
90+
SKILL_LOAD_ERROR,
91+
t)))
8292
.filter(Files::isDirectory)
8393
.mapOptional(this::findSkillMd)
8494
.map(skillMd -> new SkillMdPath(skillMd.getParent().getFileName().toString(), skillMd));
@@ -88,7 +98,8 @@ protected Flowable<SkillMdPath> listSkills() {
8898
protected Single<Path> findResourcePath(String skillName, String resourcePath) {
8999
Path file = skillsBasePath.resolve(skillName).resolve(resourcePath);
90100
if (!Files.exists(file)) {
91-
return Single.error(new SkillSourceException("Resource not found: " + file));
101+
return Single.error(
102+
new SkillSourceException("Resource not found: " + file, RESOURCE_NOT_FOUND));
92103
}
93104
return Single.just(file);
94105
}
@@ -97,11 +108,13 @@ protected Single<Path> findResourcePath(String skillName, String resourcePath) {
97108
protected Single<Path> findSkillMdPath(String skillName) {
98109
Path skillDir = skillsBasePath.resolve(skillName);
99110
if (!isDirectory(skillDir)) {
100-
return Single.error(new SkillSourceException("Skill directory not found: " + skillName));
111+
return Single.error(
112+
new SkillSourceException("Skill directory not found: " + skillName, SKILL_NOT_FOUND));
101113
}
102114
return Maybe.fromOptional(findSkillMd(skillDir))
103115
.switchIfEmpty(
104-
Single.error(new SkillSourceException("SKILL.md not found in " + skillName)));
116+
Single.error(
117+
new SkillSourceException("SKILL.md not found in " + skillName, SKILL_NOT_FOUND)));
105118
}
106119

107120
@Override

core/src/main/java/com/google/adk/skills/SkillSourceException.java

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,43 @@
2222
*/
2323
public final class SkillSourceException extends Exception {
2424

25-
public SkillSourceException(String message) {
25+
public static final String SKILL_LOAD_ERROR = "SKILL_LOAD_ERROR";
26+
public static final String SKILL_NOT_FOUND = "SKILL_NOT_FOUND";
27+
public static final String SKILL_FORMAT_ERROR = "SKILL_FORMAT_ERROR";
28+
public static final String RESOURCE_LOAD_ERROR = "RESOURCE_LOAD_ERROR";
29+
public static final String RESOURCE_NOT_FOUND = "RESOURCE_NOT_FOUND";
30+
31+
private final String errorCode;
32+
33+
/**
34+
* Constructs a new exception with the specified detail message and error code.
35+
*
36+
* @param message The detail message.
37+
* @param errorCode The specific error code categorizing the failure.
38+
*/
39+
public SkillSourceException(String message, String errorCode) {
2640
super(message);
41+
this.errorCode = errorCode;
2742
}
2843

29-
public SkillSourceException(String message, Throwable cause) {
44+
/**
45+
* Constructs a new exception with the specified detail message, error code, and cause.
46+
*
47+
* @param message The detail message.
48+
* @param errorCode The specific error code categorizing the failure.
49+
* @param cause The cause.
50+
*/
51+
public SkillSourceException(String message, String errorCode, Throwable cause) {
3052
super(message, cause);
53+
this.errorCode = errorCode;
54+
}
55+
56+
/**
57+
* Returns the error code categorizing the failure.
58+
*
59+
* @return The error code string.
60+
*/
61+
public String getErrorCode() {
62+
return errorCode;
3163
}
3264
}

core/src/main/java/com/google/adk/tools/BaseTool.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
import org.slf4j.LoggerFactory;
4646

4747
/** The base class for all ADK tools. */
48-
public abstract class BaseTool {
48+
public abstract class BaseTool implements LlmRequestProcessor {
4949
private final String name;
5050
private final String description;
5151
private final boolean isLongRunning;
@@ -183,6 +183,7 @@ private <I, O> Single<O> runAsync(
183183
* internal list of tools. Override this method for processing the outgoing request.
184184
*/
185185
@CanIgnoreReturnValue
186+
@Override
186187
public Completable processLlmRequest(
187188
LlmRequest.Builder llmRequestBuilder, ToolContext toolContext) {
188189
if (declaration().isEmpty()) {

core/src/main/java/com/google/adk/tools/BaseToolset.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public interface BaseToolset extends AutoCloseable {
2828
* Return all tools in the toolset based on the provided context.
2929
*
3030
* @param readonlyContext Context used to filter tools available to the agent.
31-
* @return A Single emitting a list of tools available under the specified context.
31+
* @return A Flowable emitting tools available under the specified context.
3232
*/
3333
Flowable<BaseTool> getTools(ReadonlyContext readonlyContext);
3434

0 commit comments

Comments
 (0)