feat: Introduced ConnectionClient and IntegrationClient to get OPENAPISPEC of connection

PiperOrigin-RevId: 775484780

Author: google-genai-bot

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

@@ -2,11 +2,14 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.ToolContext;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Streams;
 import com.google.genai.types.FunctionDeclaration;
 import com.google.genai.types.Schema;
 import io.reactivex.rxjava3.core.Single;
@@ -16,7 +19,9 @@
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 
@@ -26,11 +31,19 @@ public class ApplicationIntegrationTool extends BaseTool {
   private final String openApiSpec;
   private final String pathUrl;
   private final HttpExecutor httpExecutor;
+  private final String connectionName;
+  private final String serviceName;
+  private final String host;
+  private String entity;
+  private String operation;
+  private String action;
   private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
 
   interface HttpExecutor {
     <T> HttpResponse<T> send(HttpRequest request, HttpResponse.BodyHandler<T> responseBodyHandler)
         throws IOException, InterruptedException;
+
+    String getToken() throws IOException;
   }
 
   static class DefaultHttpExecutor implements HttpExecutor {
@@ -42,23 +55,75 @@ public <T> HttpResponse<T> send(
         throws IOException, InterruptedException {
       return client.send(request, responseBodyHandler);
     }
+
+    @Override
+    public String getToken() throws IOException {
+      GoogleCredentials credentials =
+          GoogleCredentials.getApplicationDefault()
+              .createScoped("https://www.googleapis.com/auth/cloud-platform");
+      credentials.refreshIfExpired();
+      return credentials.getAccessToken().getTokenValue();
+    }
   }
 
-  public ApplicationIntegrationTool(
+  private static final ImmutableList<String> EXCLUDE_FIELDS =
+      ImmutableList.of("connectionName", "serviceName", "host", "entity", "operation", "action");
+
+  private static final ImmutableList<String> OPTIONAL_FIELDS =
+      ImmutableList.of("pageSize", "pageToken", "filter", "sortByColumns");
+
+  /** Constructor for Application Integration Tool for integration */
+  ApplicationIntegrationTool(
       String openApiSpec, String pathUrl, String toolName, String toolDescription) {
-    // Chain to the internal constructor, providing real dependencies.
-    this(openApiSpec, pathUrl, toolName, toolDescription, new DefaultHttpExecutor());
+    this(
+        openApiSpec,
+        pathUrl,
+        toolName,
+        toolDescription,
+        null,
+        null,
+        null,
+        new DefaultHttpExecutor());
+  }
+
+  /**
+   * Constructor for Application Integration Tool with connection name, service name, host, entity,
+   * operation, and action
+   */
+  ApplicationIntegrationTool(
+      String openApiSpec,
+      String pathUrl,
+      String toolName,
+      String toolDescription,
+      String connectionName,
+      String serviceName,
+      String host) {
+    this(
+        openApiSpec,
+        pathUrl,
+        toolName,
+        toolDescription,
+        connectionName,
+        serviceName,
+        host,
+        new DefaultHttpExecutor());
   }
 
   ApplicationIntegrationTool(
       String openApiSpec,
       String pathUrl,
       String toolName,
       String toolDescription,
+      @Nullable String connectionName,
+      @Nullable String serviceName,
+      @Nullable String host,
       HttpExecutor httpExecutor) {
     super(toolName, toolDescription);
     this.openApiSpec = openApiSpec;
     this.pathUrl = pathUrl;
+    this.connectionName = connectionName;
+    this.serviceName = serviceName;
+    this.host = host;
     this.httpExecutor = httpExecutor;
   }
 
@@ -67,19 +132,10 @@ Schema toGeminiSchema(String openApiSchema, String operationId) throws Exception
     return Schema.fromJson(resolvedSchemaString);
   }
 
-  @Nullable String extractTriggerIdFromPath(String path) {
-    String prefix = "triggerId=api_trigger/";
-    int startIndex = path.indexOf(prefix);
-    if (startIndex == -1) {
-      return null;
-    }
-    return path.substring(startIndex + prefix.length());
-  }
-
   @Override
   public Optional<FunctionDeclaration> declaration() {
     try {
-      String operationId = extractTriggerIdFromPath(pathUrl);
+      String operationId = getOperationIdFromPathUrl(openApiSpec, pathUrl);
       Schema parametersSchema = toGeminiSchema(openApiSpec, operationId);
       String operationDescription = getOperationDescription(openApiSpec, operationId);
 
@@ -98,6 +154,18 @@ public Optional<FunctionDeclaration> declaration() {
 
   @Override
   public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
+    if (this.connectionName != null) {
+      args.put("connectionName", this.connectionName);
+      args.put("serviceName", this.serviceName);
+      args.put("host", this.host);
+      if (!Objects.equals(this.entity, "")) {
+        args.put("entity", this.entity);
+        args.put("operation", this.operation);
+      } else if (!Objects.equals(this.action, "")) {
+        args.put("action", this.action);
+      }
+    }
+
     return Single.fromCallable(
         () -> {
           try {
@@ -121,7 +189,7 @@ private String executeIntegration(Map<String, Object> args) throws Exception {
     HttpRequest request =
         HttpRequest.newBuilder()
             .uri(URI.create(url))
-            .header("Authorization", "Bearer " + getAccessToken())
+            .header("Authorization", "Bearer " + httpExecutor.getToken())
             .header("Content-Type", "application/json")
             .POST(HttpRequest.BodyPublishers.ofString(jsonRequestBody))
             .build();
@@ -138,12 +206,47 @@ private String executeIntegration(Map<String, Object> args) throws Exception {
     return response.body();
   }
 
-  String getAccessToken() throws IOException {
-    GoogleCredentials credentials =
-        GoogleCredentials.getApplicationDefault()
-            .createScoped("https://www.googleapis.com/auth/cloud-platform");
-    credentials.refreshIfExpired();
-    return credentials.getAccessToken().getTokenValue();
+  String getOperationIdFromPathUrl(String openApiSchemaString, String pathUrl) throws Exception {
+    JsonNode topLevelNode = OBJECT_MAPPER.readTree(openApiSchemaString);
+    JsonNode specNode = topLevelNode.path("openApiSpec");
+    if (specNode.isMissingNode() || !specNode.isTextual()) {
+      throw new IllegalArgumentException(
+          "Failed to get OpenApiSpec, please check the project and region for the integration.");
+    }
+    JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
+    JsonNode paths = rootNode.path("paths");
+
+    // Iterate through each path in the OpenAPI spec.
+    Iterator<Map.Entry<String, JsonNode>> pathsFields = paths.fields();
+    while (pathsFields.hasNext()) {
+      Map.Entry<String, JsonNode> pathEntry = pathsFields.next();
+      String currentPath = pathEntry.getKey();
+      if (!currentPath.equals(pathUrl)) {
+        continue;
+      }
+      JsonNode pathItem = pathEntry.getValue();
+
+      Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
+      while (methods.hasNext()) {
+        Map.Entry<String, JsonNode> methodEntry = methods.next();
+        JsonNode operationNode = methodEntry.getValue();
+        // Set  values for entity, operation, and action
+        this.entity = "";
+        this.operation = "";
+        this.action = "";
+        if (operationNode.has("x-entity")) {
+          this.entity = operationNode.path("x-entity").asText();
+          this.operation = operationNode.path("x-operation").asText();
+        } else if (operationNode.has("x-action")) {
+          this.action = operationNode.path("x-action").asText();
+        }
+        // Get the operationId from the operationNode
+        if (operationNode.has("operationId")) {
+          return operationNode.path("operationId").asText();
+        }
+      }
+    }
+    throw new Exception("Could not find operationId for pathUrl: " + pathUrl);
   }
 
   private String getResolvedRequestSchemaByOperationId(
@@ -155,7 +258,6 @@ private String getResolvedRequestSchemaByOperationId(
           "Failed to get OpenApiSpec, please check the project and region for the integration.");
     }
     JsonNode rootNode = OBJECT_MAPPER.readTree(specNode.asText());
-
     JsonNode operationNode = findOperationNodeById(rootNode, operationId);
     if (operationNode == null) {
       throw new Exception("Could not find operation with operationId: " + operationId);
@@ -169,19 +271,47 @@ private String getResolvedRequestSchemaByOperationId(
 
     JsonNode resolvedSchema = resolveRefs(requestSchemaNode, rootNode);
 
+    if (resolvedSchema.isObject()) {
+      ObjectNode schemaObject = (ObjectNode) resolvedSchema;
+
+      // 1. Remove excluded fields from the 'properties' object.
+      JsonNode propertiesNode = schemaObject.path("properties");
+      if (propertiesNode.isObject()) {
+        ObjectNode propertiesObject = (ObjectNode) propertiesNode;
+        for (String field : EXCLUDE_FIELDS) {
+          propertiesObject.remove(field);
+        }
+      }
+
+      // 2. Remove optional and excluded fields from the 'required' array.
+      JsonNode requiredNode = schemaObject.path("required");
+      if (requiredNode.isArray()) {
+        // Combine the lists of fields to remove
+        List<String> fieldsToRemove =
+            Streams.concat(OPTIONAL_FIELDS.stream(), EXCLUDE_FIELDS.stream()).toList();
+
+        // To safely remove items from a list while iterating, we must use an Iterator.
+        ArrayNode requiredArray = (ArrayNode) requiredNode;
+        Iterator<JsonNode> elements = requiredArray.elements();
+        while (elements.hasNext()) {
+          JsonNode element = elements.next();
+          if (element.isTextual() && fieldsToRemove.contains(element.asText())) {
+            // This removes the current element from the underlying array.
+            elements.remove();
+          }
+        }
+      }
+    }
     return OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(resolvedSchema);
   }
 
   private @Nullable JsonNode findOperationNodeById(JsonNode rootNode, String operationId) {
     JsonNode paths = rootNode.path("paths");
-    // Iterate through each path in the OpenAPI spec.
     for (JsonNode pathItem : paths) {
-      // Iterate through each HTTP method (e.g., GET, POST) for the current path.
       Iterator<Map.Entry<String, JsonNode>> methods = pathItem.fields();
       while (methods.hasNext()) {
         Map.Entry<String, JsonNode> methodEntry = methods.next();
         JsonNode operationNode = methodEntry.getValue();
-        // Check if the operationId matches the target operationId.
         if (operationNode.path("operationId").asText().equals(operationId)) {
           return operationNode;
         }