Skip to content

Commit 5ab8b14

Browse files
google-genai-botcopybara-github
authored andcommitted
feat!: Added serviceAccountJson as a parameter for toolset
BREAKING CHANGE: This change requires users to update their configurations to provide a service account JSON file. This enables authentication with cloud services. PiperOrigin-RevId: 784598108
1 parent 5a7ab20 commit 5ab8b14

File tree

4 files changed

+238
-43
lines changed

4 files changed

+238
-43
lines changed

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/ApplicationIntegrationToolset.java

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,24 +5,28 @@
55
import com.fasterxml.jackson.databind.JsonNode;
66
import com.fasterxml.jackson.databind.ObjectMapper;
77
import com.fasterxml.jackson.databind.node.ObjectNode;
8+
import com.google.adk.agents.ReadonlyContext;
89
import com.google.adk.tools.BaseTool;
10+
import com.google.adk.tools.BaseToolset;
911
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.DefaultHttpExecutor;
1012
import com.google.adk.tools.applicationintegrationtoolset.IntegrationConnectorTool.HttpExecutor;
13+
import io.reactivex.rxjava3.core.Flowable;
1114
import java.util.ArrayList;
1215
import java.util.Iterator;
1316
import java.util.List;
1417
import java.util.Map;
1518
import org.jspecify.annotations.Nullable;
1619

1720
/** Application Integration Toolset */
18-
public class ApplicationIntegrationToolset {
21+
public class ApplicationIntegrationToolset implements BaseToolset {
1922
String project;
2023
String location;
2124
@Nullable String integration;
2225
@Nullable List<String> triggers;
2326
@Nullable String connection;
2427
@Nullable Map<String, List<String>> entityOperations;
2528
@Nullable List<String> actions;
29+
String serviceAccountJson;
2630
@Nullable String toolNamePrefix;
2731
@Nullable String toolInstructions;
2832
public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
@@ -35,16 +39,16 @@ public class ApplicationIntegrationToolset {
3539
*
3640
* <p>integrationTool = new ApplicationIntegrationToolset( project="test-project",
3741
* location="us-central1", integration="test-integration",
38-
* triggers=ImmutableList.of("api_trigger/test_trigger",
39-
* "api_trigger/test_trigger_2"),connection=null,enitityOperations=null,actions=null,toolNamePrefix="test-integration-tool",toolInstructions="This
42+
* triggers=ImmutableList.of("api_trigger/test_trigger", "api_trigger/test_trigger_2",
43+
* serviceAccountJson="{....}"),connection=null,enitityOperations=null,actions=null,toolNamePrefix="test-integration-tool",toolInstructions="This
4044
* tool is used to get response from test-integration.");
4145
*
4246
* <p>connectionTool = new ApplicationIntegrationToolset( project="test-project",
4347
* location="us-central1", integration=null, triggers=null, connection="test-connection",
4448
* entityOperations=ImmutableMap.of("Entity1", ImmutableList.of("LIST", "GET", "UPDATE")),
4549
* "Entity2", ImmutableList.of()), actions=ImmutableList.of("ExecuteCustomQuery"),
46-
* toolNamePrefix="test-tool", toolInstructions="This tool is used to list, get and update issues
47-
* in Jira.");
50+
* serviceAccountJson="{....}", toolNamePrefix="test-tool", toolInstructions="This tool is used to
51+
* list, get and update issues in Jira.");
4852
*
4953
* @param project The GCP project ID.
5054
* @param location The GCP location of integration.
@@ -53,6 +57,9 @@ public class ApplicationIntegrationToolset {
5357
* @param connection(Optional) The connection name.
5458
* @param entityOperations(Optional) The entity operations.
5559
* @param actions(Optional) The actions.
60+
* @param serviceAccountJson(Optional) The service account configuration as a dictionary. Required
61+
* if not using default service credential. Used for fetching the Application Integration or
62+
* Integration Connector resource.
5663
* @param toolNamePrefix(Optional) The tool name prefix.
5764
* @param toolInstructions(Optional) The tool instructions.
5865
*/
@@ -64,6 +71,7 @@ public ApplicationIntegrationToolset(
6471
String connection,
6572
Map<String, List<String>> entityOperations,
6673
List<String> actions,
74+
String serviceAccountJson,
6775
String toolNamePrefix,
6876
String toolInstructions) {
6977
this(
@@ -74,9 +82,10 @@ public ApplicationIntegrationToolset(
7482
connection,
7583
entityOperations,
7684
actions,
85+
serviceAccountJson,
7786
toolNamePrefix,
7887
toolInstructions,
79-
new DefaultHttpExecutor());
88+
new DefaultHttpExecutor().createExecutor(serviceAccountJson));
8089
}
8190

8291
ApplicationIntegrationToolset(
@@ -87,6 +96,7 @@ public ApplicationIntegrationToolset(
8796
String connection,
8897
Map<String, List<String>> entityOperations,
8998
List<String> actions,
99+
String serviceAccountJson,
90100
String toolNamePrefix,
91101
String toolInstructions,
92102
HttpExecutor httpExecutor) {
@@ -97,6 +107,7 @@ public ApplicationIntegrationToolset(
97107
this.connection = connection;
98108
this.entityOperations = entityOperations;
99109
this.actions = actions;
110+
this.serviceAccountJson = serviceAccountJson;
100111
this.toolNamePrefix = toolNamePrefix;
101112
this.toolInstructions = toolInstructions;
102113
this.httpExecutor = httpExecutor;
@@ -121,7 +132,7 @@ List<String> getPathUrl(String openApiSchemaString) throws Exception {
121132
return pathUrls;
122133
}
123134

124-
public List<BaseTool> getTools() throws Exception {
135+
private List<BaseTool> getAllTools() throws Exception {
125136
String openApiSchemaString = null;
126137
List<BaseTool> tools = new ArrayList<>();
127138
if (!isNullOrEmpty(this.integration)) {
@@ -142,7 +153,15 @@ public List<BaseTool> getTools() throws Exception {
142153
if (toolName != null) {
143154
tools.add(
144155
new IntegrationConnectorTool(
145-
openApiSchemaString, pathUrl, toolName, "", null, null, null, this.httpExecutor));
156+
openApiSchemaString,
157+
pathUrl,
158+
toolName,
159+
toolInstructions,
160+
null,
161+
null,
162+
null,
163+
this.serviceAccountJson,
164+
this.httpExecutor));
146165
}
147166
}
148167
} else if (!isNullOrEmpty(this.connection)
@@ -182,6 +201,7 @@ public List<BaseTool> getTools() throws Exception {
182201
connectionDetails.name,
183202
connectionDetails.serviceName,
184203
connectionDetails.host,
204+
this.serviceAccountJson,
185205
this.httpExecutor));
186206
}
187207
}
@@ -193,4 +213,19 @@ public List<BaseTool> getTools() throws Exception {
193213

194214
return tools;
195215
}
216+
217+
@Override
218+
public Flowable<BaseTool> getTools(@Nullable ReadonlyContext readonlyContext) {
219+
try {
220+
List<BaseTool> allTools = getAllTools();
221+
return Flowable.fromIterable(allTools);
222+
} catch (Exception e) {
223+
return Flowable.error(e);
224+
}
225+
}
226+
227+
@Override
228+
public void close() throws Exception {
229+
// Nothing to close.
230+
}
196231
}

core/src/main/java/com/google/adk/tools/applicationintegrationtoolset/IntegrationConnectorTool.java

Lines changed: 66 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.google.adk.tools.applicationintegrationtoolset;
22

33
import static com.google.common.base.Strings.isNullOrEmpty;
4+
import static java.nio.charset.StandardCharsets.UTF_8;
45

56
import com.fasterxml.jackson.databind.JsonNode;
67
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -15,7 +16,9 @@
1516
import com.google.genai.types.FunctionDeclaration;
1617
import com.google.genai.types.Schema;
1718
import io.reactivex.rxjava3.core.Single;
19+
import java.io.ByteArrayInputStream;
1820
import java.io.IOException;
21+
import java.io.InputStream;
1922
import java.net.URI;
2023
import java.net.http.HttpClient;
2124
import java.net.http.HttpRequest;
@@ -45,10 +48,27 @@ <T> HttpResponse<T> send(HttpRequest request, HttpResponse.BodyHandler<T> respon
4548
throws IOException, InterruptedException;
4649

4750
String getToken() throws IOException;
51+
52+
public HttpExecutor createExecutor(String serviceAccountJson);
4853
}
4954

5055
static class DefaultHttpExecutor implements HttpExecutor {
5156
private final HttpClient client = HttpClient.newHttpClient();
57+
private final String serviceAccountJson;
58+
59+
/** Default constructor for when no service account is specified. */
60+
DefaultHttpExecutor() {
61+
this(null);
62+
}
63+
64+
/**
65+
* Constructor that accepts an optional service account JSON string.
66+
*
67+
* @param serviceAccountJson The service account key as a JSON string, or null.
68+
*/
69+
DefaultHttpExecutor(@Nullable String serviceAccountJson) {
70+
this.serviceAccountJson = serviceAccountJson;
71+
}
5272

5373
@Override
5474
public <T> HttpResponse<T> send(
@@ -59,12 +79,42 @@ public <T> HttpResponse<T> send(
5979

6080
@Override
6181
public String getToken() throws IOException {
62-
GoogleCredentials credentials =
63-
GoogleCredentials.getApplicationDefault()
64-
.createScoped("https://www.googleapis.com/auth/cloud-platform");
82+
GoogleCredentials credentials;
83+
84+
if (this.serviceAccountJson != null && !this.serviceAccountJson.trim().isEmpty()) {
85+
try (InputStream is = new ByteArrayInputStream(this.serviceAccountJson.getBytes(UTF_8))) {
86+
credentials =
87+
GoogleCredentials.fromStream(is)
88+
.createScoped("https://www.googleapis.com/auth/cloud-platform");
89+
} catch (IOException e) {
90+
throw new IOException("Failed to load credentials from service_account_json.", e);
91+
}
92+
} else {
93+
try {
94+
credentials =
95+
GoogleCredentials.getApplicationDefault()
96+
.createScoped("https://www.googleapis.com/auth/cloud-platform");
97+
} catch (IOException e) {
98+
throw new IOException(
99+
"Please provide a service account or configure Application Default Credentials. To"
100+
+ " set up ADC, see"
101+
+ " https://cloud.google.com/docs/authentication/external/set-up-adc.",
102+
e);
103+
}
104+
}
105+
65106
credentials.refreshIfExpired();
66107
return credentials.getAccessToken().getTokenValue();
67108
}
109+
110+
@Override
111+
public HttpExecutor createExecutor(String serviceAccountJson) {
112+
if (isNullOrEmpty(serviceAccountJson)) {
113+
return new DefaultHttpExecutor();
114+
} else {
115+
return new DefaultHttpExecutor(serviceAccountJson);
116+
}
117+
}
68118
}
69119

70120
private static final ImmutableList<String> EXCLUDE_FIELDS =
@@ -75,7 +125,11 @@ public String getToken() throws IOException {
75125

76126
/** Constructor for Application Integration Tool for integration */
77127
IntegrationConnectorTool(
78-
String openApiSpec, String pathUrl, String toolName, String toolDescription) {
128+
String openApiSpec,
129+
String pathUrl,
130+
String toolName,
131+
String toolDescription,
132+
String serviceAccountJson) {
79133
this(
80134
openApiSpec,
81135
pathUrl,
@@ -84,7 +138,8 @@ public String getToken() throws IOException {
84138
null,
85139
null,
86140
null,
87-
new DefaultHttpExecutor());
141+
serviceAccountJson,
142+
new DefaultHttpExecutor().createExecutor(serviceAccountJson));
88143
}
89144

90145
/**
@@ -98,7 +153,8 @@ public String getToken() throws IOException {
98153
String toolDescription,
99154
String connectionName,
100155
String serviceName,
101-
String host) {
156+
String host,
157+
String serviceAccountJson) {
102158
this(
103159
openApiSpec,
104160
pathUrl,
@@ -107,7 +163,8 @@ public String getToken() throws IOException {
107163
connectionName,
108164
serviceName,
109165
host,
110-
new DefaultHttpExecutor());
166+
serviceAccountJson,
167+
new DefaultHttpExecutor().createExecutor(serviceAccountJson));
111168
}
112169

113170
IntegrationConnectorTool(
@@ -118,6 +175,7 @@ public String getToken() throws IOException {
118175
@Nullable String connectionName,
119176
@Nullable String serviceName,
120177
@Nullable String host,
178+
@Nullable String serviceAccountJson,
121179
HttpExecutor httpExecutor) {
122180
super(toolName, toolDescription);
123181
this.openApiSpec = openApiSpec;
@@ -182,7 +240,7 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
182240
}
183241

184242
private String executeIntegration(Map<String, Object> args) throws Exception {
185-
String url = String.format("https://integrations.googleapis.com%s", pathUrl);
243+
String url = String.format("https://integrations.googleapis.com%s", this.pathUrl);
186244
String jsonRequestBody;
187245
try {
188246
jsonRequestBody = OBJECT_MAPPER.writeValueAsString(args);

0 commit comments

Comments
 (0)