This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 83fcb2f9cdf [Gemini] Add Java GeminiModelHandler class (#39245)
83fcb2f9cdf is described below
commit 83fcb2f9cdfa630f63743bde1c48c88c49812e59
Author: Jack McCluskey <[email protected]>
AuthorDate: Tue Jul 21 13:34:07 2026 -0400
[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
---
sdks/java/ml/inference/gemini/build.gradle | 42 ++++++
.../ml/inference/gemini/GeminiImageResponse.java | 33 +++++
.../inference/gemini/GeminiInferenceFunctions.java | 67 +++++++++
.../ml/inference/gemini/GeminiModelHandler.java | 86 +++++++++++
.../ml/inference/gemini/GeminiModelParameters.java | 60 ++++++++
.../ml/inference/gemini/GeminiRequestFunction.java | 31 ++++
.../sdk/ml/inference/gemini/GeminiStringInput.java | 33 +++++
.../ml/inference/gemini/GeminiStringResponse.java | 33 +++++
.../beam/sdk/ml/inference/gemini/package-info.java | 20 +++
.../ml/inference/gemini/GeminiModelHandlerIT.java | 163 +++++++++++++++++++++
.../inference/gemini/GeminiModelHandlerTest.java | 133 +++++++++++++++++
.../ml/inference/remote/PredictionResultCoder.java | 74 ++++++++++
.../sdk/ml/inference/remote/RemoteInference.java | 34 ++++-
.../ml/inference/remote/RemoteInferenceTest.java | 4 +-
settings.gradle.kts | 1 +
15 files changed, 809 insertions(+), 5 deletions(-)
diff --git a/sdks/java/ml/inference/gemini/build.gradle
b/sdks/java/ml/inference/gemini/build.gradle
new file mode 100644
index 00000000000..f62984eba88
--- /dev/null
+++ b/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
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java
new file mode 100644
index 00000000000..5d74c21347c
--- /dev/null
+++
b/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;
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java
new file mode 100644
index 00000000000..5ba009b7194
--- /dev/null
+++
b/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;
+ };
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java
new file mode 100644
index 00000000000..ecdbf807816
--- /dev/null
+++
b/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);
+ }
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java
new file mode 100644
index 00000000000..6ac4f578a24
--- /dev/null
+++
b/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();
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java
new file mode 100644
index 00000000000..15da8b360c9
--- /dev/null
+++
b/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;
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java
new file mode 100644
index 00000000000..8b0cc7c2689
--- /dev/null
+++
b/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;
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java
new file mode 100644
index 00000000000..94bb4d7168d
--- /dev/null
+++
b/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;
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java
b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java
new file mode 100644
index 00000000000..9dee9da9cc7
--- /dev/null
+++
b/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;
diff --git
a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java
b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java
new file mode 100644
index 00000000000..c1188ad6d04
--- /dev/null
+++
b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java
@@ -0,0 +1,163 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeNotNull;
+import static org.junit.Assume.assumeTrue;
+
+import org.apache.beam.sdk.ml.inference.remote.PredictionResult;
+import org.apache.beam.sdk.ml.inference.remote.RemoteInference;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+/**
+ * Execute Gemini model handler integration test.
+ *
+ * <pre>
+ * ./gradlew :sdks:java:ml:inference:gemini:integrationTest \
+ * --tests org.apache.beam.sdk.ml.inference.gemini.GeminiModelHandlerIT \
+ * --info
+ * </pre>
+ */
+public class GeminiModelHandlerIT {
+
+ public static class GeminiStringContentHandler
+ extends GeminiModelHandler<GeminiStringInput, GeminiStringResponse> {}
+
+ public static class GeminiStringImageHandler
+ extends GeminiModelHandler<GeminiStringInput, GeminiImageResponse> {}
+
+ @Rule public final transient TestPipeline pipeline = TestPipeline.create();
+
+ private String apiKey;
+ private static final String API_KEY_ENV = "GEMINI_API_KEY";
+ private static final String DEFAULT_MODEL = "gemini-3.5-flash";
+
+ @Before
+ public void setUp() {
+ // Get API key
+ apiKey = System.getenv(API_KEY_ENV);
+
+ // Skip tests if API key is not provided
+ assumeNotNull(
+ "Gemini API key not found. Set "
+ + API_KEY_ENV
+ + " environment variable to run integration tests.",
+ apiKey);
+ assumeTrue(
+ "Gemini API key is empty. Set "
+ + API_KEY_ENV
+ + " environment variable to run integration tests.",
+ !apiKey.trim().isEmpty());
+ }
+
+ @Test
+ public void testSentimentAnalysisWithSingleInput() {
+ String input =
+ "Classify the sentiment of the following text as either positive or
negative: This product is absolutely amazing! I love it!";
+
+ PCollection<GeminiStringInput> inputs =
+ pipeline.apply("CreateSingleInput", Create.of(new
GeminiStringInput(input)));
+
+ PCollection<Iterable<PredictionResult<GeminiStringInput,
GeminiStringResponse>>> results =
+ inputs.apply(
+ "SentimentInference",
+ RemoteInference.<GeminiStringInput, GeminiStringResponse>invoke()
+ .handler(GeminiStringContentHandler.class)
+ .withParameters(
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setApiKey(apiKey)
+ .setModelName(DEFAULT_MODEL)
+
.setRequestFn(GeminiInferenceFunctions.generateFromString())
+ .build()));
+
+ // Verify results
+ PAssert.that(results)
+ .satisfies(
+ batches -> {
+ int count = 0;
+ for (Iterable<PredictionResult<GeminiStringInput,
GeminiStringResponse>> batch :
+ batches) {
+ for (PredictionResult<GeminiStringInput, GeminiStringResponse>
result : batch) {
+ count++;
+ assertNotNull("Input should not be null", result.getInput());
+ assertNotNull("Output should not be null",
result.getOutput());
+
+ String sentiment =
result.getOutput().getText().toLowerCase();
+ assertTrue(
+ "Sentiment should be positive or negative, got: " +
sentiment,
+ sentiment.contains("positive") ||
sentiment.contains("negative"));
+ }
+ }
+ assertEquals("Should have exactly 1 result", 1, count);
+ return null;
+ });
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testGenerateImageFromString() {
+ String input = "A beautiful sunset over the mountains, digital art style.";
+
+ PCollection<GeminiStringInput> inputs =
+ pipeline.apply("CreateImageInput", Create.of(new
GeminiStringInput(input)));
+
+ PCollection<Iterable<PredictionResult<GeminiStringInput,
GeminiImageResponse>>> results =
+ inputs.apply(
+ "ImageInference",
+ RemoteInference.<GeminiStringInput, GeminiImageResponse>invoke()
+ .handler(GeminiStringImageHandler.class)
+ .withParameters(
+ GeminiModelParameters.<GeminiStringInput,
GeminiImageResponse>builder()
+ .setApiKey(apiKey)
+ .setModelName("imagen-4.0-generate-001")
+
.setRequestFn(GeminiInferenceFunctions.generateImageFromString())
+ .build()));
+
+ // Verify results
+ PAssert.that(results)
+ .satisfies(
+ batches -> {
+ int count = 0;
+ for (Iterable<PredictionResult<GeminiStringInput,
GeminiImageResponse>> batch :
+ batches) {
+ for (PredictionResult<GeminiStringInput, GeminiImageResponse>
result : batch) {
+ count++;
+ assertNotNull("Input should not be null", result.getInput());
+ assertNotNull("Output should not be null",
result.getOutput());
+ assertTrue(
+ "Output should have generated images",
+ result.getOutput().getImageBytes().length > 0);
+ }
+ }
+ assertEquals("Should have exactly 1 result", 1, count);
+ return null;
+ });
+
+ pipeline.run().waitUntilFinish();
+ }
+}
diff --git
a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java
b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java
new file mode 100644
index 00000000000..125f1e9c3e9
--- /dev/null
+++
b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java
@@ -0,0 +1,133 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import java.util.Arrays;
+import java.util.List;
+import org.apache.beam.sdk.ml.inference.remote.PredictionResult;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class GeminiModelHandlerTest {
+
+ @Test
+ public void testAllParamsSet() {
+ GeminiModelParameters<GeminiStringInput, GeminiStringResponse> parameters =
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setApiKey("test-key")
+ .setProject("test-project")
+ .setLocation("us-central1")
+ .setModelName("gemini-model-123")
+ .setRequestFn(GeminiInferenceFunctions.generateFromString())
+ .build();
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ assertThrows(IllegalArgumentException.class, () ->
handler.createClient(parameters));
+ }
+
+ @Test
+ public void testMissingVertexLocationParam() {
+ GeminiModelParameters<GeminiStringInput, GeminiStringResponse> parameters =
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setProject("test-project")
+ .setModelName("gemini-model-123")
+ .setRequestFn(GeminiInferenceFunctions.generateFromString())
+ .build();
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ assertThrows(IllegalArgumentException.class, () ->
handler.createClient(parameters));
+ }
+
+ @Test
+ public void testMissingVertexProjectParam() {
+ GeminiModelParameters<GeminiStringInput, GeminiStringResponse> parameters =
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setLocation("us-central1")
+ .setModelName("gemini-model-123")
+ .setRequestFn(GeminiInferenceFunctions.generateFromString())
+ .build();
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ assertThrows(IllegalArgumentException.class, () ->
handler.createClient(parameters));
+ }
+
+ @Test
+ public void testNullParameters() {
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ assertThrows(NullPointerException.class, () -> handler.createClient(null));
+ }
+
+ @Test
+ public void testRequest() throws Exception {
+ GeminiModelParameters<GeminiStringInput, GeminiStringResponse> parameters =
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setApiKey("test-key")
+ .setModelName("gemini-model-123")
+ .setRequestFn(
+ (modelName, batch, client) ->
+ Arrays.asList(
+ new GeminiStringResponse("response1"),
+ new GeminiStringResponse("response2")))
+ .build();
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ handler.createClient(parameters);
+
+ List<GeminiStringInput> input =
+ Arrays.asList(new GeminiStringInput("input1"), new
GeminiStringInput("input2"));
+ Iterable<PredictionResult<GeminiStringInput, GeminiStringResponse>>
results =
+ handler.request(input);
+
+ int count = 0;
+ for (PredictionResult<GeminiStringInput, GeminiStringResponse> result :
results) {
+ if (count == 0) {
+ assertEquals("input1", result.getInput().getText());
+ assertEquals("response1", result.getOutput().getText());
+ } else {
+ assertEquals("input2", result.getInput().getText());
+ assertEquals("response2", result.getOutput().getText());
+ }
+ count++;
+ }
+ assertEquals(2, count);
+ }
+
+ @Test
+ public void testRequestMismatchedResponseSize() throws Exception {
+ GeminiModelParameters<GeminiStringInput, GeminiStringResponse> parameters =
+ GeminiModelParameters.<GeminiStringInput,
GeminiStringResponse>builder()
+ .setApiKey("test-key")
+ .setModelName("gemini-model-123")
+ .setRequestFn(
+ (modelName, batch, client) -> Arrays.asList(new
GeminiStringResponse("response1")))
+ .build();
+ GeminiModelHandler<GeminiStringInput, GeminiStringResponse> handler =
+ new GeminiModelHandler<>();
+ handler.createClient(parameters);
+
+ List<GeminiStringInput> input =
+ Arrays.asList(new GeminiStringInput("input1"), new
GeminiStringInput("input2"));
+ assertThrows(RuntimeException.class, () -> handler.request(input));
+ }
+}
diff --git
a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java
new file mode 100644
index 00000000000..0d43870b880
--- /dev/null
+++
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java
@@ -0,0 +1,74 @@
+/*
+ * 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.remote;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.StructuredCoder;
+
+/** A {@link Coder} for {@link PredictionResult}. */
+public class PredictionResultCoder<InputT, OutputT>
+ extends StructuredCoder<PredictionResult<InputT, OutputT>> {
+
+ private final Coder<InputT> inputCoder;
+ private final Coder<OutputT> outputCoder;
+
+ private PredictionResultCoder(Coder<InputT> inputCoder, Coder<OutputT>
outputCoder) {
+ this.inputCoder = inputCoder;
+ this.outputCoder = outputCoder;
+ }
+
+ public static <InputT, OutputT> PredictionResultCoder<InputT, OutputT> of(
+ Coder<InputT> inputCoder, Coder<OutputT> outputCoder) {
+ return new PredictionResultCoder<>(inputCoder, outputCoder);
+ }
+
+ @Override
+ public void encode(PredictionResult<InputT, OutputT> value, OutputStream
outStream)
+ throws CoderException, IOException {
+ if (value == null) {
+ throw new CoderException("cannot encode a null PredictionResult");
+ }
+ inputCoder.encode(value.getInput(), outStream);
+ outputCoder.encode(value.getOutput(), outStream);
+ }
+
+ @Override
+ public PredictionResult<InputT, OutputT> decode(InputStream inStream)
+ throws CoderException, IOException {
+ InputT input = inputCoder.decode(inStream);
+ OutputT output = outputCoder.decode(inStream);
+ return PredictionResult.create(input, output);
+ }
+
+ @Override
+ public List<? extends Coder<?>> getCoderArguments() {
+ return Arrays.asList(inputCoder, outputCoder);
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ verifyDeterministic(this, "Input coder must be deterministic", inputCoder);
+ verifyDeterministic(this, "Output coder must be deterministic",
outputCoder);
+ }
+}
diff --git
a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
index c571911a484..a998f4d0f07 100644
---
a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
+++
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
@@ -21,6 +21,7 @@ import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
import com.google.auto.value.AutoValue;
import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.components.throttling.ReactiveThrottler;
import org.apache.beam.sdk.transforms.BatchElements;
import org.apache.beam.sdk.transforms.DoFn;
@@ -88,6 +89,8 @@ public class RemoteInference {
abstract @Nullable Double overloadRatio();
+ abstract @Nullable Coder<OutputT> outputCoder();
+
abstract Builder<InputT, OutputT> builder();
@AutoValue.Builder
@@ -108,6 +111,8 @@ public class RemoteInference {
abstract Builder<InputT, OutputT> setOverloadRatio(Double overloadRatio);
+ abstract Builder<InputT, OutputT> setOutputCoder(Coder<OutputT>
outputCoder);
+
abstract Invoke<InputT, OutputT> build();
}
@@ -163,6 +168,14 @@ public class RemoteInference {
return builder().setOverloadRatio(overloadRatio).build();
}
+ /**
+ * Configures the coder for the output of the model. If not provided, it
will fallback to using
+ * standard Java serialization for the output element.
+ */
+ public Invoke<InputT, OutputT> withOutputCoder(Coder<OutputT> outputCoder)
{
+ return builder().setOutputCoder(outputCoder).build();
+ }
+
@Override
public PCollection<Iterable<PredictionResult<InputT, OutputT>>> expand(
PCollection<InputT> input) {
@@ -177,9 +190,24 @@ public class RemoteInference {
batchedInput = input.apply("BatchElements",
BatchElements.withDefaults());
}
- return batchedInput
- // Pass the list to the inference function
- .apply("RemoteInference", ParDo.of(new RemoteInferenceFn<InputT,
OutputT>(this)));
+ PCollection<Iterable<PredictionResult<InputT, OutputT>>> result =
+ batchedInput
+ // Pass the list to the inference function
+ .apply("RemoteInference", ParDo.of(new RemoteInferenceFn<InputT,
OutputT>(this)));
+
+ Coder<OutputT> outCoder = outputCoder();
+ if (outCoder != null) {
+ result.setCoder(
+ org.apache.beam.sdk.coders.IterableCoder.of(
+ PredictionResultCoder.of(input.getCoder(), outCoder)));
+ } else {
+ result.setCoder(
+ (org.apache.beam.sdk.coders.Coder)
+ org.apache.beam.sdk.coders.IterableCoder.of(
+
org.apache.beam.sdk.coders.SerializableCoder.of(PredictionResult.class)));
+ }
+
+ return result;
}
/**
diff --git
a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
index 64691aa1fed..5fe8ea1cce1 100644
---
a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
+++
b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
@@ -45,7 +45,7 @@ public class RemoteInferenceTest {
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
// Test input class
- public static class TestInput implements BaseInput {
+ public static class TestInput implements BaseInput, java.io.Serializable {
private final String value;
private TestInput(String value) {
@@ -84,7 +84,7 @@ public class RemoteInferenceTest {
}
// Test output class
- public static class TestOutput implements BaseResponse {
+ public static class TestOutput implements BaseResponse, java.io.Serializable
{
private final String result;
private TestOutput(String result) {
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 44e93029b54..8a116eeb394 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -285,6 +285,7 @@ include(":sdks:java:maven-archetypes:gcp-bom-examples")
include(":sdks:java:maven-archetypes:starter")
include(":sdks:java:ml:inference:remote")
include(":sdks:java:ml:inference:openai")
+include(":sdks:java:ml:inference:gemini")
include(":sdks:java:testing:nexmark")
include(":sdks:java:testing:expansion-service")
include(":sdks:java:testing:jpms-tests")
