gemini-code-assist[bot] commented on code in PR #39245: URL: https://github.com/apache/beam/pull/39245#discussion_r3544600731
########## 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) { + 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 { + if (parameters.getProject() == null || parameters.getLocation() == null) { + throw new IllegalArgumentException( + "Project and location must both be provided if API key is not set"); + } + Client.Builder builder = + Client.builder() + .vertexAI(true) + .project(parameters.getProject()) + .location(parameters.getLocation()); + + this.client = builder.build(); + } + } + Review Comment:  The current client initialization logic requires either an explicit API key or both project and location. This prevents users from utilizing the default client initialization, which automatically picks up credentials and configuration from standard environment variables (such as `GOOGLE_API_KEY` or `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION`).\n\nAdditionally, we should add a defensive null check for the `parameters` argument to prevent a `NullPointerException`. ```java @Override\n public void createClient(GeminiModelParameters<InputT, OutputT> parameters) {\n if (parameters == null) {\n throw new IllegalArgumentException(\"Parameters must not be null\");\n }\n this.modelParameters = parameters;\n\n // Configure client based on vertex or API key\n if (parameters.getApiKey() != null) {\n if (parameters.getProject() != null || parameters.getLocation() != null) {\n throw new IllegalArgumentException(\"Project and location must be null if API key is set\");\n }\n this.client = Client.builder().apiKey(parameters.getApiKey()).build();\n } else if (parameters.getProject() != null || parameters.getLocation() != null) {\n if (parameters.getProject() == null || parameters.getLocation() == null) {\n throw new IllegalArgumentException(\n \"Project and location must both be provided if API key is not set\");\n }\n this.client =\n Client.builder()\n .vertexAI (true)\n .project(parameters.getProject())\n .location(parameters.getLocation())\n .build();\n } else {\n // Fallback to default client initialization (picks up environment variables)\n this.client = Client.builder().build();\n }\n } ``` ########## 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) { + 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 { + if (parameters.getProject() == null || parameters.getLocation() == null) { + throw new IllegalArgumentException( + "Project and location must both be provided if API key is not set"); + } + Client.Builder builder = + Client.builder() + .vertexAI(true) + .project(parameters.getProject()) + .location(parameters.getLocation()); + + 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); + } + } Review Comment:  We should add defensive guards to ensure that `input`, `modelParameters`, and `client` are not null before executing the request.\n\nAdditionally, if the `input` list is empty, we can return an empty list immediately to avoid unnecessary remote API calls and potential errors. ```java @Override\n public Iterable<PredictionResult<InputT, OutputT>> request(List<InputT> input) {\n if (input == null) {\n throw new IllegalArgumentException(\"Input list must not be null\");\n }\n if (input.isEmpty()) {\n return java.util.Collections.emptyList();\n }\n if (modelParameters == null) {\n throw new IllegalStateException(\"Model parameters are not initialized. createClient must be called before request.\");\n }\n if (client == null) {\n throw new IllegalStateException(\"Client is not initialized. createClient must be called before request.\");\n }\n try {\n GeminiRequestFunction<InputT, OutputT> requestFn = modelParameters.getRequestFn();\n List<OutputT> responses = requestFn.apply(modelParameters.getModelName(), input, client);\n\n if (responses.size() != input.size()) {\n throw new IllegalStateException(\"Number of responses must match number of inputs\");\n }\n\n List<PredictionResult<Inpu tT, OutputT>> results = new ArrayList<>();\n for (int i = 0; i < input.size(); i++) {\n results.add(PredictionResult.create(input.get(i), responses.get(i)));\n }\n return results;\n } catch (Exception e) {\n throw new RuntimeException(\"Error during Gemini inference request\", e);\n }\n } ``` ########## sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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<String, String> parameters = + GeminiModelParameters.<String, String>builder() + .setApiKey("test-key") + .setProject("test-project") + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler<String, String> handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexLocationParam() { + GeminiModelParameters<String, String> parameters = + GeminiModelParameters.<String, String>builder() + .setProject("test-project") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler<String, String> handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexProjectParam() { + GeminiModelParameters<String, String> parameters = + GeminiModelParameters.<String, String>builder() + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler<String, String> handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingAllParams() { + GeminiModelParameters<String, String> parameters = + GeminiModelParameters.<String, String>builder() + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler<String, String> handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } Review Comment:  With the updated client initialization logic that allows default client creation (picking up environment variables), this test should be updated to assert that no exception is thrown when all parameters are missing. ```java @Test\n public void testDefaultClientInitialization() {\n GeminiModelParameters<String, String> parameters =\n GeminiModelParameters.<String, String>builder()\n .setModelName(\"gemini-model-123\")\n .setRequestFn(GeminiInferenceFunctions.generateFromString())\n .build();\n GeminiModelHandler<String, String> handler = new GeminiModelHandler<>();\n // Should not throw an exception, allowing default initialization via environment variables\n handler.createClient(parameters);\n } ``` ########## sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java: ########## @@ -0,0 +1,61 @@ +/* + * 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 boolean getUseVertexFlexApi(); Review Comment:  The parameter `getUseVertexFlexApi` is defined in `GeminiModelParameters` but is never used anywhere in the codebase. If it is not yet supported or needed, we should remove it to keep the API clean and maintainable. ########## sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java: ########## @@ -0,0 +1,39 @@ +/* + * 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.lang.reflect.Method; + +public class TestReflection { + public static void main(String[] args) throws Exception { + System.out.println("Client.Builder methods:"); + for (Method m : Client.Builder.class.getMethods()) { + if (m.getDeclaringClass() != Object.class) { + System.out.println(m.getName() + " " + m.getParameterCount()); + } + } + System.out.println("Client methods:"); + for (Method m : Client.class.getMethods()) { + if (m.getDeclaringClass() != Object.class) { + System.out.println(m.getName()); + System.out.println(" return: " + m.getReturnType().getName()); + } + } + } +} Review Comment:  The `TestReflection` class appears to be a temporary debugging utility used to inspect the `Client` class via reflection. It is not a unit test and should be removed from the final codebase to keep it clean. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
