Skip to content

Commit 0f95125

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: 781819195
1 parent 2aa474d commit 0f95125

File tree

8 files changed

+2291
-195
lines changed

8 files changed

+2291
-195
lines changed

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

Lines changed: 213 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,32 @@
11
package com.google.adk.tools.applicationintegrationtoolset;
22

3+
import static com.google.common.base.Strings.isNullOrEmpty;
4+
import static java.nio.charset.StandardCharsets.UTF_8;
5+
36
import com.fasterxml.jackson.databind.JsonNode;
47
import com.fasterxml.jackson.databind.ObjectMapper;
8+
import com.fasterxml.jackson.databind.node.ArrayNode;
59
import com.fasterxml.jackson.databind.node.ObjectNode;
610
import com.google.adk.tools.BaseTool;
711
import com.google.adk.tools.ToolContext;
812
import com.google.auth.oauth2.GoogleCredentials;
13+
import com.google.common.collect.ImmutableList;
914
import com.google.common.collect.ImmutableMap;
15+
import com.google.common.collect.Streams;
1016
import com.google.genai.types.FunctionDeclaration;
1117
import com.google.genai.types.Schema;
1218
import io.reactivex.rxjava3.core.Single;
19+
import java.io.ByteArrayInputStream;
1320
import java.io.IOException;
21+
import java.io.InputStream;
1422
import java.net.URI;
1523
import java.net.http.HttpClient;
1624
import java.net.http.HttpRequest;
1725
import java.net.http.HttpResponse;
1826
import java.util.Iterator;
27+
import java.util.List;
1928
import java.util.Map;
29+
import java.util.Objects;
2030
import java.util.Optional;
2131
import org.jspecify.annotations.Nullable;
2232

@@ -26,39 +36,152 @@ public class ApplicationIntegrationTool extends BaseTool {
2636
private final String openApiSpec;
2737
private final String pathUrl;
2838
private final HttpExecutor httpExecutor;
39+
private final String connectionName;
40+
private final String serviceName;
41+
private final String host;
42+
private String entity;
43+
private String operation;
44+
private String action;
2945
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
3046

3147
interface HttpExecutor {
3248
<T> HttpResponse<T> send(HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler)
3349
throws IOException, InterruptedException;
50+
51+
String getToken() throws IOException;
3452
}
3553

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

3972
@Override
4073
public <T> HttpResponse<T> send(
4174
HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler)
4275
throws IOException, InterruptedException {
4376
return client.send(request, responseBodyHandler);
4477
}
78+
79+
@Override
80+
public String getToken() throws IOException {
81+
GoogleCredentials credentials;
82+
83+
if (this.serviceAccountJson != null && !this.serviceAccountJson.trim().isEmpty()) {
84+
try (InputStream is = new ByteArrayInputStream(this.serviceAccountJson.getBytes(UTF_8))) {
85+
credentials =
86+
GoogleCredentials.fromStream(is)
87+
.createScoped("https://www.googleapis.com/auth/cloud-platform");
88+
} catch (IOException e) {
89+
throw new IOException("Failed to load credentials from service_account_json.", e);
90+
}
91+
} else {
92+
try {
93+
credentials =
94+
GoogleCredentials.getApplicationDefault()
95+
.createScoped("https://www.googleapis.com/auth/cloud-platform");
96+
} catch (IOException e) {
97+
throw new IOException(
98+
"Please provide a service account or configure Application Default Credentials. To"
99+
+ " set up ADC, see"
100+
+ " https://cloud.google.com/docs/authentication/external/set-up-adc.",
101+
e);
102+
}
103+
}
104+
105+
credentials.refreshIfExpired();
106+
return credentials.getAccessToken().getTokenValue();
107+
}
45108
}
46109

47-
public ApplicationIntegrationTool(
48-
String openApiSpec, String pathUrl, String toolName, String toolDescription) {
49-
// Chain to the internal constructor, providing real dependencies.
50-
this(openApiSpec, pathUrl, toolName, toolDescription, new DefaultHttpExecutor());
110+
/** A private static factory method to choose the correct HttpExecutor. */
111+
private static HttpExecutor createExecutor(String serviceAccountJson) {
112+
if (isNullOrEmpty(serviceAccountJson)) {
113+
return new DefaultHttpExecutor();
114+
} else {
115+
return new DefaultHttpExecutor(serviceAccountJson);
116+
}
51117
}
52118

119+
private static final ImmutableList<String> EXCLUDE_FIELDS =
120+
ImmutableList.of("connectionName", "serviceName", "host", "entity", "operation", "action");
121+
122+
private static final ImmutableList<String> OPTIONAL_FIELDS =
123+
ImmutableList.of("pageSize", "pageToken", "filter", "sortByColumns");
124+
125+
/** Constructor for Application Integration Tool for integration */
53126
ApplicationIntegrationTool(
54127
String openApiSpec,
55128
String pathUrl,
56129
String toolName,
57130
String toolDescription,
131+
String serviceAccountJson) {
132+
this(
133+
openApiSpec,
134+
pathUrl,
135+
toolName,
136+
toolDescription,
137+
null,
138+
null,
139+
null,
140+
serviceAccountJson,
141+
createExecutor(serviceAccountJson));
142+
}
143+
144+
/**
145+
* Constructor for Application Integration Tool with connection name, service name, host, entity,
146+
* operation, and action
147+
*/
148+
ApplicationIntegrationTool(
149+
String openApiSpec,
150+
String pathUrl,
151+
String toolName,
152+
String toolDescription,
153+
String connectionName,
154+
String serviceName,
155+
String host,
156+
String serviceAccountJson) {
157+
this(
158+
openApiSpec,
159+
pathUrl,
160+
toolName,
161+
toolDescription,
162+
connectionName,
163+
serviceName,
164+
host,
165+
serviceAccountJson,
166+
createExecutor(serviceAccountJson));
167+
}
168+
169+
ApplicationIntegrationTool(
170+
String openApiSpec,
171+
String pathUrl,
172+
String toolName,
173+
String toolDescription,
174+
@Nullable String connectionName,
175+
@Nullable String serviceName,
176+
@Nullable String host,
177+
@Nullable String serviceAccountJson,
58178
HttpExecutor httpExecutor) {
59179
super(toolName, toolDescription);
60180
this.openApiSpec = openApiSpec;
61181
this.pathUrl = pathUrl;
182+
this.connectionName = connectionName;
183+
this.serviceName = serviceName;
184+
this.host = host;
62185
this.httpExecutor = httpExecutor;
63186
}
64187

@@ -67,19 +190,10 @@ Schema toGeminiSchema(String openApiSchema, String operationId) throws Exception
67190
return Schema.fromJson(resolvedSchemaString);
68191
}
69192

70-
@Nullable String extractTriggerIdFromPath(String path) {
71-
String prefix = "triggerId=api_trigger/";
72-
int startIndex = path.indexOf(prefix);
73-
if (startIndex == -1) {
74-
return null;
75-
}
76-
return path.substring(startIndex + prefix.length());
77-
}
78-
79193
@Override
80194
public Optional<FunctionDeclaration> declaration() {
81195
try {
82-
String operationId = extractTriggerIdFromPath(pathUrl);
196+
String operationId = getOperationIdFromPathUrl(openApiSpec, pathUrl);
83197
Schema parametersSchema = toGeminiSchema(openApiSpec, operationId);
84198
String operationDescription = getOperationDescription(openApiSpec, operationId);
85199

@@ -98,6 +212,18 @@ public Optional<FunctionDeclaration> declaration() {
98212

99213
@Override
100214
public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
215+
if (this.connectionName != null) {
216+
args.put("connectionName", this.connectionName);
217+
args.put("serviceName", this.serviceName);
218+
args.put("host", this.host);
219+
if (!Objects.equals(this.entity, "")) {
220+
args.put("entity", this.entity);
221+
args.put("operation", this.operation);
222+
} else if (!Objects.equals(this.action, "")) {
223+
args.put("action", this.action);
224+
}
225+
}
226+
101227
return Single.fromCallable(
102228
() -> {
103229
try {
@@ -121,7 +247,7 @@ private String executeIntegration(Map<String, Object> args) throws Exception {
121247
HttpRequest request =
122248
HttpRequest.newBuilder()
123249
.uri(URI.create(url))
124-
.header("Authorization", "Bearer " + getAccessToken())
250+
.header("Authorization", "Bearer " + httpExecutor.getToken())
125251
.header("Content-Type", "application/json")
126252
.POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody))
127253
.build();
@@ -138,12 +264,47 @@ private String executeIntegration(Map<String, Object> args) throws Exception {
138264
return response.body();
139265
}
140266

141-
String getAccessToken() throws IOException {
142-
GoogleCredentials credentials =
143-
GoogleCredentials.getApplicationDefault()
144-
.createScoped("https://www.googleapis.com/auth/cloud-platform");
145-
credentials.refreshIfExpired();
146-
return credentials.getAccessToken().getTokenValue();
267+
String getOperationIdFromPathUrl(String openApiSchemaString, String pathUrl) throws Exception {
268+
JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
269+
JsonNode specNode = topLevelNode.path("openApiSpec");
270+
if (specNode.isMissingNode() || !specNode.isTextual()) {
271+
throw new IllegalArgumentException(
272+
"Failed to get OpenApiSpec, please check the project and region for the integration.");
273+
}
274+
JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
275+
JsonNode paths = rootNode.path("paths");
276+
277+
// Iterate through each path in the OpenAPI spec.
278+
Iterator<Map.Entry<String, JsonNode>> pathsFields = paths.fields();
279+
while (pathsFields.hasNext()) {
280+
Map.Entry<String, JsonNode> pathEntry = pathsFields.next();
281+
String currentPath = pathEntry.getKey();
282+
if (!currentPath.equals(pathUrl)) {
283+
continue;
284+
}
285+
JsonNode pathItem = pathEntry.getValue();
286+
287+
Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
288+
while (methods.hasNext()) {
289+
Map.Entry<String, JsonNode> methodEntry = methods.next();
290+
JsonNode operationNode = methodEntry.getValue();
291+
// Set values for entity, operation, and action
292+
this.entity = "";
293+
this.operation = "";
294+
this.action = "";
295+
if (operationNode.has("x-entity")) {
296+
this.entity = operationNode.path("x-entity").asText();
297+
this.operation = operationNode.path("x-operation").asText();
298+
} else if (operationNode.has("x-action")) {
299+
this.action = operationNode.path("x-action").asText();
300+
}
301+
// Get the operationId from the operationNode
302+
if (operationNode.has("operationId")) {
303+
return operationNode.path("operationId").asText();
304+
}
305+
}
306+
}
307+
throw new Exception("Could not find operationId for pathUrl: " + pathUrl);
147308
}
148309

149310
private String getResolvedRequestSchemaByOperationId(
@@ -155,7 +316,6 @@ private String getResolvedRequestSchemaByOperationId(
155316
"Failed to get OpenApiSpec, please check the project and region for the integration.");
156317
}
157318
JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
158-
159319
JsonNode operationNode = findOperationNodeById(rootNode, operationId);
160320
if (operationNode == null) {
161321
throw new Exception("Could not find operation with operationId: " + operationId);
@@ -169,19 +329,47 @@ private String getResolvedRequestSchemaByOperationId(
169329

170330
JsonNode resolvedSchema = resolveRefs(requestSchemaNode, rootNode);
171331

332+
if (resolvedSchema.isObject()) {
333+
ObjectNode schemaObject = (ObjectNode) resolvedSchema;
334+
335+
// 1. Remove excluded fields from the 'properties' object.
336+
JsonNode propertiesNode = schemaObject.path("properties");
337+
if (propertiesNode.isObject()) {
338+
ObjectNode propertiesObject = (ObjectNode) propertiesNode;
339+
for (String field : EXCLUDE_FIELDS) {
340+
propertiesObject.remove(field);
341+
}
342+
}
343+
344+
// 2. Remove optional and excluded fields from the 'required' array.
345+
JsonNode requiredNode = schemaObject.path("required");
346+
if (requiredNode.isArray()) {
347+
// Combine the lists of fields to remove
348+
List<String> fieldsToRemove =
349+
Streams.concat(OPTIONAL_FIELDS.stream(), EXCLUDE_FIELDS.stream()).toList();
350+
351+
// To safely remove items from a list while iterating, we must use an Iterator.
352+
ArrayNode requiredArray = (ArrayNode) requiredNode;
353+
Iterator<JsonNode> elements = requiredArray.elements();
354+
while (elements.hasNext()) {
355+
JsonNode element = elements.next();
356+
if (element.isTextual() && fieldsToRemove.contains(element.asText())) {
357+
// This removes the current element from the underlying array.
358+
elements.remove();
359+
}
360+
}
361+
}
362+
}
172363
return OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(resolvedSchema);
173364
}
174365

175366
private @Nullable JsonNode findOperationNodeById(JsonNode rootNode, String operationId) {
176367
JsonNode paths = rootNode.path("paths");
177-
// Iterate through each path in the OpenAPI spec.
178368
for (JsonNode pathItem : paths) {
179-
// Iterate through each HTTP method (e.g., GET, POST) for the current path.
180369
Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
181370
while (methods.hasNext()) {
182371
Map.Entry<String, JsonNode> methodEntry = methods.next();
183372
JsonNode operationNode = methodEntry.getValue();
184-
// Check if the operationId matches the target operationId.
185373
if (operationNode.path("operationId").asText().equals(operationId)) {
186374
return operationNode;
187375
}

0 commit comments

Comments
 (0)