[Gemini] Add Java GeminiModelHandler class (#39245)

* [Gemini] Add Java GeminiModelHandler class

* Remove model scratch files

* address review comments

* checkstyle issue

* unit test fix

* Fix weird coder UX

* fix altered base class signatures

Author: jrmccluskey

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

@@ -0,0 +1,42 @@
+/*
+ * 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',
+)
+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"
+
+
+  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/GeminiImageResponse.java

@@ -0,0 +1,33 @@
+/*
+ * 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 org.apache.beam.sdk.ml.inference.remote.BaseResponse;
+
+/** Wrapper for image (byte array) responses from Gemini. */
+public class GeminiImageResponse implements BaseResponse {
+  public final byte[] imageBytes;
+
+  public GeminiImageResponse(byte[] imageBytes) {
+    this.imageBytes = imageBytes;
+  }
+
+  public byte[] getImageBytes() {
+    return imageBytes;
+  }
+}

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

@@ -0,0 +1,67 @@
+/*
+ * 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<GeminiStringInput, GeminiStringResponse>
+      generateFromString() {
+    return (modelName, batch, client) -> {
+      List<GeminiStringResponse> results = new ArrayList<>();
+      for (GeminiStringInput input : batch) {
+        GenerateContentResponse response =
+            client.models.generateContent(
+                modelName, input.getText(), GenerateContentConfig.builder().build());
+        String text = response.text();
+        results.add(new GeminiStringResponse(text != null ? text : ""));
+      }
+      return results;
+    };
+  }
+
+  /** Generates images from string prompts using the generateImages API. */
+  public static GeminiRequestFunction<GeminiStringInput, GeminiImageResponse>
+      generateImageFromString() {
+    return (modelName, batch, client) -> {
+      List<GeminiImageResponse> results = new ArrayList<>();
+      for (GeminiStringInput input : batch) {
+        GenerateImagesResponse response =
+            client.models.generateImages(
+                modelName, input.getText(), 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(new GeminiImageResponse(imageBytes));
+        } else {
+          results.add(new GeminiImageResponse(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,86 @@
+/*
+ * 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.BaseInput;
+import org.apache.beam.sdk.ml.inference.remote.BaseModelHandler;
+import org.apache.beam.sdk.ml.inference.remote.BaseResponse;
+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 extends BaseInput, OutputT extends BaseResponse>
+    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,60 @@
+/*
+ * 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.BaseInput;
+import org.apache.beam.sdk.ml.inference.remote.BaseModelParameters;
+import org.apache.beam.sdk.ml.inference.remote.BaseResponse;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+@AutoValue
+public abstract class GeminiModelParameters<InputT extends BaseInput, OutputT extends BaseResponse>
+    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 extends BaseInput, OutputT extends BaseResponse>
+      Builder<InputT, OutputT> builder() {
+    return new AutoValue_GeminiModelParameters.Builder<InputT, OutputT>();
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder<InputT extends BaseInput, OutputT extends BaseResponse> {
+    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,31 @@
+/*
+ * 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;
+import org.apache.beam.sdk.ml.inference.remote.BaseInput;
+import org.apache.beam.sdk.ml.inference.remote.BaseResponse;
+
+/** Functional interface for custom request functions to the Gemini API. */
+@FunctionalInterface
+public interface GeminiRequestFunction<InputT extends BaseInput, OutputT extends BaseResponse>
+    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/GeminiStringInput.java

@@ -0,0 +1,33 @@
+/*
+ * 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 org.apache.beam.sdk.ml.inference.remote.BaseInput;
+
+/** Wrapper for string inputs to Gemini. */
+public class GeminiStringInput implements BaseInput {
+  public final String text;
+
+  public GeminiStringInput(String text) {
+    this.text = text;
+  }
+
+  public String getText() {
+    return text;
+  }
+}

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

@@ -0,0 +1,33 @@
+/*
+ * 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 org.apache.beam.sdk.ml.inference.remote.BaseResponse;
+
+/** Wrapper for string responses from Gemini. */
+public class GeminiStringResponse implements BaseResponse {
+  public final String text;
+
+  public GeminiStringResponse(String text) {
+    this.text = text;
+  }
+
+  public String getText() {
+    return text;
+  }
+}