[Gemini] Add Java GeminiModelHandler class

sdks/java/ml/inference/gemini/build.gradle

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+plugins {
+  id 'org.apache.beam.module'
+}
+
+applyJavaNature(
+  automaticModuleName: 'org.apache.beam.sdk.ml.inference.gemini',
+  requireJavaVersion: JavaVersion.VERSION_11
+)
+provideIntegrationTestingDependencies()
+enableJavaPerformanceTesting()
+
+description = "Apache Beam :: SDKs :: Java :: ML :: Inference :: Gemini"
+ext.summary = "Gemini model handler for remote inference"
+
+dependencies {
+  implementation project(":sdks:java:ml:inference:remote")
+  implementation "com.google.genai:google-genai:1.59.0"
+  implementation library.java.jackson_databind
+  implementation library.java.jackson_annotations
+  implementation library.java.jackson_core
+
+  testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
+  testImplementation project(path: ":sdks:java:core", configuration: "shadow")
+  testImplementation library.java.slf4j_api
+  testRuntimeOnly library.java.slf4j_simple
+  testImplementation library.java.junit
+  testImplementation library.java.mockito_core
+}

sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.ml.inference.gemini;
+
+import com.google.genai.types.GenerateContentConfig;
+import com.google.genai.types.GenerateContentResponse;
+import com.google.genai.types.GenerateImagesConfig;
+import com.google.genai.types.GenerateImagesResponse;
+import java.util.ArrayList;
+import java.util.List;
+
+/** Common inference functions for Gemini. */
+public class GeminiInferenceFunctions {
+
+  /** Generates content from string prompts using the standard generateContent API. */
+  public static GeminiRequestFunction<String, String> generateFromString() {
+    return (modelName, batch, client) -> {
+      List<String> results = new ArrayList<>();
+      for (String input : batch) {
+        GenerateContentResponse response =
+            client.models.generateContent(
+                modelName, input, GenerateContentConfig.builder().build());
+        String text = response.text();
+        results.add(text != null ? text : "");
+      }
+      return results;
+    };
+  }
+
+  /** Generates images from string prompts using the generateImages API. */
+  public static GeminiRequestFunction<String, byte[]> generateImageFromString() {
+    return (modelName, batch, client) -> {
+      List<byte[]> results = new ArrayList<>();
+      for (String input : batch) {
+        GenerateImagesResponse response =
+            client.models.generateImages(modelName, input, GenerateImagesConfig.builder().build());
+        // Retrieve the base64 string or bytes from the first generated image
+        List<com.google.genai.types.Image> images = response.images();
+        if (images != null && !images.isEmpty()) {
+          byte[] imageBytes = images.get(0).imageBytes().orElse(new byte[0]);
+          results.add(imageBytes);
+        } else {
+          results.add(new byte[0]);
+        }
+      }
+      return results;
+    };
+  }
+}

sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.ml.inference.gemini;
+
+import com.google.genai.Client;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.ml.inference.remote.BaseModelHandler;
+import org.apache.beam.sdk.ml.inference.remote.PredictionResult;
+
+/**
+ * Model handler for Google Gemini API inference requests.
+ *
+ * <p>This handler manages communication with Google's Gemini API, including client initialization,
+ * request formatting, and response parsing. It allows executing a custom {@link
+ * GeminiRequestFunction} against a batch of inputs.
+ */
+@SuppressWarnings("nullness")
+public class GeminiModelHandler<InputT, OutputT>
+    implements BaseModelHandler<GeminiModelParameters<InputT, OutputT>, InputT, OutputT> {
+
+  private transient Client client;
+  private GeminiModelParameters<InputT, OutputT> modelParameters;
+
+  @Override
+  public void createClient(GeminiModelParameters<InputT, OutputT> parameters) {
+    if (parameters == null) {
+      throw new NullPointerException("GeminiModelParameters must not be null");
+    }
+    this.modelParameters = parameters;
+
+    // Configure client based on vertex or API key
+    if (parameters.getApiKey() != null) {
+      if (parameters.getProject() != null || parameters.getLocation() != null) {
+        throw new IllegalArgumentException("Project and location must be null if API key is set");
+      }
+      this.client = Client.builder().apiKey(parameters.getApiKey()).build();
+    } else {
+      Client.Builder builder = Client.builder();
+      if (parameters.getProject() != null && parameters.getLocation() != null) {
+        builder.vertexAI(true).project(parameters.getProject()).location(parameters.getLocation());
+      } else if (parameters.getProject() != null || parameters.getLocation() != null) {
+        throw new IllegalArgumentException(
+            "Project and location must both be provided if one is provided");
+      }
+      this.client = builder.build();
+    }
+  }
+
+  @Override
+  public Iterable<PredictionResult<InputT, OutputT>> request(List<InputT> input) {
+    try {
+      GeminiRequestFunction<InputT, OutputT> requestFn = modelParameters.getRequestFn();
+      List<OutputT> responses = requestFn.apply(modelParameters.getModelName(), input, client);
+
+      if (responses.size() != input.size()) {
+        throw new IllegalStateException("Number of responses must match number of inputs");
+      }
+
+      List<PredictionResult<InputT, OutputT>> results = new ArrayList<>();
+      for (int i = 0; i < input.size(); i++) {
+        results.add(PredictionResult.create(input.get(i), responses.get(i)));
+      }
+      return results;
+    } catch (Exception e) {
+      throw new RuntimeException("Error during Gemini inference request", e);
+    }
+  }
+}

sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.ml.inference.gemini;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.sdk.ml.inference.remote.BaseModelParameters;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+@AutoValue
+public abstract class GeminiModelParameters<InputT, OutputT> implements BaseModelParameters {
+
+  public abstract @Nullable String getApiKey();
+
+  public abstract @Nullable String getProject();
+
+  public abstract @Nullable String getLocation();
+
+  public abstract String getModelName();
+
+  public abstract GeminiRequestFunction<InputT, OutputT> getRequestFn();
+
+  public static <InputT, OutputT> Builder<InputT, OutputT> builder() {
+    return new AutoValue_GeminiModelParameters.Builder<InputT, OutputT>();
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder<InputT, OutputT> {
+    public abstract Builder<InputT, OutputT> setApiKey(String apiKey);
+
+    public abstract Builder<InputT, OutputT> setProject(String project);
+
+    public abstract Builder<InputT, OutputT> setLocation(String location);
+
+    public abstract Builder<InputT, OutputT> setModelName(String modelName);
+
+    public abstract Builder<InputT, OutputT> setRequestFn(
+        GeminiRequestFunction<InputT, OutputT> requestFn);
+
+    public abstract GeminiModelParameters<InputT, OutputT> build();
+  }
+}

sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.ml.inference.gemini;
+
+import com.google.genai.Client;
+import java.io.Serializable;
+import java.util.List;
+
+/** Functional interface for custom request functions to the Gemini API. */
+@FunctionalInterface
+public interface GeminiRequestFunction<InputT, OutputT> extends Serializable {
+  List<OutputT> apply(String modelName, List<InputT> batch, Client client) throws Exception;
+}

sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java

@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/** Gemini model handler for remote inference. */
+package org.apache.beam.sdk.ml.inference.gemini;