Skip to content

Commit 8c1e66e

Browse files
IvanLHcopybara-github
IvanLH
authored and
copybara-github
committed
feat: Added Operations.get which is a generic method which will handle all Operation types.
PiperOrigin-RevId: 795297000
1 parent 5aae52b commit 8c1e66e

File tree

10 files changed

+175
-69
lines changed

10 files changed

+175
-69
lines changed

‎examples/src/main/java/com/google/genai/examples/GenerateVideos.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,7 @@ public static void main(String[] args) {
8686
while (!generateVideosOperation.done().filter(Boolean::booleanValue).isPresent()) {
8787
try {
8888
Thread.sleep(10000); // Sleep for 10 seconds.
89-
generateVideosOperation =
90-
client.operations.getVideosOperation(generateVideosOperation, null);
89+
generateVideosOperation = client.operations.get(generateVideosOperation, null);
9190
System.out.println("Waiting for operation to complete...");
9291
} catch (InterruptedException e) {
9392
System.out.println("Thread was interrupted while sleeping.");

‎examples/src/main/java/com/google/genai/examples/GenerateVideosAsync.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ public static void main(String[] args) {
9292
try {
9393
Thread.sleep(10000); // Sleep for 10 seconds.
9494
try {
95-
operation = client.async.operations.getVideosOperation(operation, null).get();
95+
operation = client.async.operations.get(operation, null).get();
9696
} catch (ExecutionException e) {
9797
throw new RuntimeException(e);
9898
}

‎src/main/java/com/google/genai/AsyncOperations.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import com.google.genai.types.GenerateVideosOperation;
2222
import com.google.genai.types.GetOperationConfig;
23+
import com.google.genai.types.Operation;
2324
import java.util.concurrent.CompletableFuture;
2425

2526
/** Async module of {@link Operations} */
@@ -41,4 +42,16 @@ public CompletableFuture<GenerateVideosOperation> getVideosOperation(
4142
GenerateVideosOperation operation, GetOperationConfig config) {
4243
return CompletableFuture.supplyAsync(() -> operations.getVideosOperation(operation, config));
4344
}
45+
46+
/**
47+
* Gets the status of an Operation.
48+
*
49+
* @param operation An Operation.
50+
* @param config The configuration for getting the operation.
51+
* @return An Operation with the updated status of the operation.
52+
*/
53+
public <T, U extends Operation<T, U>> CompletableFuture<U> get(
54+
U operation, GetOperationConfig config) {
55+
return CompletableFuture.supplyAsync(() -> operations.get(operation, config));
56+
}
4457
}

‎src/main/java/com/google/genai/JsonSerializable.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
/** A class that can be serialized to JSON and deserialized from JSON. */
3939
public abstract class JsonSerializable {
4040

41-
static final ObjectMapper objectMapper = new ObjectMapper();
41+
@InternalApi protected static final ObjectMapper objectMapper = new ObjectMapper();
4242

4343
/** Custom Jackson serializer for {@link java.time.Duration} to output "Xs" format. */
4444
static class CustomDurationSerializer extends JsonSerializer<java.time.Duration> {
@@ -134,7 +134,8 @@ protected static <T extends JsonSerializable> T fromJsonString(
134134
}
135135

136136
/** Deserializes a JsonNode to an object of the given type. */
137-
static <T extends JsonSerializable> T fromJsonNode(JsonNode jsonNode, Class<T> clazz) {
137+
@InternalApi
138+
protected static <T extends JsonSerializable> T fromJsonNode(JsonNode jsonNode, Class<T> clazz) {
138139
try {
139140
return objectMapper.treeToValue(jsonNode, clazz);
140141
} catch (JsonProcessingException e) {

‎src/main/java/com/google/genai/LiveConverters.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.fasterxml.jackson.databind.node.ArrayNode;
2424
import com.fasterxml.jackson.databind.node.ObjectNode;
2525

26+
/** Internal SDK converter functions. */
2627
final class LiveConverters {
2728
private final ApiClient apiClient;
2829

‎src/main/java/com/google/genai/Operations.java‎

Lines changed: 38 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.google.genai.types.GetOperationConfig;
3030
import com.google.genai.types.GetOperationParameters;
3131
import com.google.genai.types.HttpOptions;
32+
import com.google.genai.types.Operation;
3233
import java.io.IOException;
3334
import java.util.Optional;
3435
import okhttp3.ResponseBody;
@@ -355,8 +356,7 @@ ObjectNode generateVideosOperationFromVertex(JsonNode fromObject, ObjectNode par
355356
return toObject;
356357
}
357358

358-
GenerateVideosOperation privateGetVideosOperation(
359-
String operationName, GetOperationConfig config) {
359+
JsonNode privateGetVideosOperation(String operationName, GetOperationConfig config) {
360360

361361
GetOperationParameters.Builder parameterBuilder = GetOperationParameters.builder();
362362

@@ -408,17 +408,11 @@ GenerateVideosOperation privateGetVideosOperation(
408408
throw new GenAiIOException("Failed to read HTTP response.", e);
409409
}
410410

411-
JsonNode responseNode = JsonSerializable.stringToJsonNode(responseString);
412-
if (this.apiClient.vertexAI()) {
413-
responseNode = generateVideosOperationFromVertex(responseNode, null);
414-
} else {
415-
responseNode = generateVideosOperationFromMldev(responseNode, null);
416-
}
417-
return JsonSerializable.fromJsonNode(responseNode, GenerateVideosOperation.class);
411+
return JsonSerializable.stringToJsonNode(responseString);
418412
}
419413
}
420414

421-
GenerateVideosOperation privateFetchPredictVideosOperation(
415+
JsonNode privateFetchPredictVideosOperation(
422416
String operationName, String resourceName, FetchPredictOperationConfig config) {
423417

424418
FetchPredictOperationParameters.Builder parameterBuilder =
@@ -471,14 +465,7 @@ GenerateVideosOperation privateFetchPredictVideosOperation(
471465
throw new GenAiIOException("Failed to read HTTP response.", e);
472466
}
473467

474-
JsonNode responseNode = JsonSerializable.stringToJsonNode(responseString);
475-
if (this.apiClient.vertexAI()) {
476-
responseNode = generateVideosOperationFromVertex(responseNode, null);
477-
} else {
478-
throw new UnsupportedOperationException(
479-
"This method is only supported in the Vertex AI client.");
480-
}
481-
return JsonSerializable.fromJsonNode(responseNode, GenerateVideosOperation.class);
468+
return JsonSerializable.stringToJsonNode(responseString);
482469
}
483470
}
484471

@@ -491,20 +478,49 @@ GenerateVideosOperation privateFetchPredictVideosOperation(
491478
*/
492479
public GenerateVideosOperation getVideosOperation(
493480
GenerateVideosOperation operation, GetOperationConfig config) {
481+
if (!operation.name().isPresent()) {
482+
throw new IllegalArgumentException("Operation name is required.");
483+
}
494484

485+
if (this.apiClient.vertexAI()) {
486+
String resourceName = operation.name().get().split("/operations/")[0];
487+
488+
FetchPredictOperationConfig fetchConfig = FetchPredictOperationConfig.builder().build();
489+
490+
JsonNode response =
491+
this.privateFetchPredictVideosOperation(
492+
operation.name().get(), resourceName, fetchConfig);
493+
return operation.fromApiResponse(response, true);
494+
} else {
495+
JsonNode response = this.privateGetVideosOperation(operation.name().get(), config);
496+
return operation.fromApiResponse(response, false);
497+
}
498+
}
499+
500+
/**
501+
* Gets the status of an Operation.
502+
*
503+
* @param operation An Operation.
504+
* @param config The configuration for getting the operation.
505+
* @return An Operation with the updated status of the operation.
506+
*/
507+
public <T, U extends Operation<T, U>> U get(U operation, GetOperationConfig config) {
495508
if (!operation.name().isPresent()) {
496-
throw new Error("Operation name is required.");
509+
throw new IllegalArgumentException("Operation name is required.");
497510
}
498511

499512
if (this.apiClient.vertexAI()) {
500513
String resourceName = operation.name().get().split("/operations/")[0];
501514

502515
FetchPredictOperationConfig fetchConfig = FetchPredictOperationConfig.builder().build();
503516

504-
return this.privateFetchPredictVideosOperation(
505-
operation.name().get(), resourceName, fetchConfig);
517+
JsonNode response =
518+
this.privateFetchPredictVideosOperation(
519+
operation.name().get(), resourceName, fetchConfig);
520+
return operation.fromApiResponse(response, true);
506521
} else {
507-
return this.privateGetVideosOperation(operation.name().get(), config);
522+
JsonNode response = this.privateGetVideosOperation(operation.name().get(), config);
523+
return operation.fromApiResponse(response, false);
508524
}
509525
}
510526
}

‎src/main/java/com/google/genai/OperationsConverters.java‎

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,21 @@
2222
import com.fasterxml.jackson.databind.ObjectMapper;
2323
import com.fasterxml.jackson.databind.node.ArrayNode;
2424
import com.fasterxml.jackson.databind.node.ObjectNode;
25+
import com.google.api.core.InternalApi;
2526

26-
final class OperationsConverters {
27+
/** Internal SDK converter functions. */
28+
@InternalApi
29+
public final class OperationsConverters {
2730
private final ApiClient apiClient;
2831

2932
public OperationsConverters(ApiClient apiClient) {
3033
this.apiClient = apiClient;
3134
}
3235

3336
@ExcludeFromGeneratedCoverageReport
34-
ObjectNode fetchPredictOperationParametersToMldev(JsonNode fromObject, ObjectNode parentObject) {
37+
@InternalApi
38+
public ObjectNode fetchPredictOperationParametersToMldev(
39+
JsonNode fromObject, ObjectNode parentObject) {
3540
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
3641
if (!Common.isZero(Common.getValueByPath(fromObject, new String[] {"operationName"}))) {
3742
throw new IllegalArgumentException("operationName parameter is not supported in Gemini API.");
@@ -49,7 +54,8 @@ ObjectNode fetchPredictOperationParametersToMldev(JsonNode fromObject, ObjectNod
4954
}
5055

5156
@ExcludeFromGeneratedCoverageReport
52-
ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentObject) {
57+
@InternalApi
58+
public ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentObject) {
5359
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
5460
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
5561
Common.setValueByPath(
@@ -69,7 +75,9 @@ ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentO
6975
}
7076

7177
@ExcludeFromGeneratedCoverageReport
72-
ObjectNode fetchPredictOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
78+
@InternalApi
79+
public ObjectNode fetchPredictOperationParametersToVertex(
80+
JsonNode fromObject, ObjectNode parentObject) {
7381
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
7482
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
7583
Common.setValueByPath(
@@ -96,7 +104,8 @@ ObjectNode fetchPredictOperationParametersToVertex(JsonNode fromObject, ObjectNo
96104
}
97105

98106
@ExcludeFromGeneratedCoverageReport
99-
ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
107+
@InternalApi
108+
public ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
100109
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
101110
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
102111
Common.setValueByPath(
@@ -116,7 +125,8 @@ ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parent
116125
}
117126

118127
@ExcludeFromGeneratedCoverageReport
119-
ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
128+
@InternalApi
129+
public ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
120130
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
121131
if (Common.getValueByPath(fromObject, new String[] {"video", "uri"}) != null) {
122132
Common.setValueByPath(
@@ -144,7 +154,8 @@ ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
144154
}
145155

146156
@ExcludeFromGeneratedCoverageReport
147-
ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
157+
@InternalApi
158+
public ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
148159
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
149160
if (Common.getValueByPath(fromObject, new String[] {"_self"}) != null) {
150161
Common.setValueByPath(
@@ -160,7 +171,8 @@ ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject)
160171
}
161172

162173
@ExcludeFromGeneratedCoverageReport
163-
ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
174+
@InternalApi
175+
public ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
164176
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
165177
if (Common.getValueByPath(fromObject, new String[] {"generatedSamples"}) != null) {
166178
ArrayNode keyArray =
@@ -192,7 +204,8 @@ ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode paren
192204
}
193205

194206
@ExcludeFromGeneratedCoverageReport
195-
ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode parentObject) {
207+
@InternalApi
208+
public ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode parentObject) {
196209
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
197210
if (Common.getValueByPath(fromObject, new String[] {"name"}) != null) {
198211
Common.setValueByPath(
@@ -238,7 +251,8 @@ ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode pare
238251
}
239252

240253
@ExcludeFromGeneratedCoverageReport
241-
ObjectNode videoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
254+
@InternalApi
255+
public ObjectNode videoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
242256
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
243257
if (Common.getValueByPath(fromObject, new String[] {"gcsUri"}) != null) {
244258
Common.setValueByPath(
@@ -266,7 +280,8 @@ ObjectNode videoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
266280
}
267281

268282
@ExcludeFromGeneratedCoverageReport
269-
ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
283+
@InternalApi
284+
public ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
270285
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
271286
if (Common.getValueByPath(fromObject, new String[] {"_self"}) != null) {
272287
Common.setValueByPath(
@@ -282,7 +297,8 @@ ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject
282297
}
283298

284299
@ExcludeFromGeneratedCoverageReport
285-
ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode parentObject) {
300+
@InternalApi
301+
public ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode parentObject) {
286302
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
287303
if (Common.getValueByPath(fromObject, new String[] {"videos"}) != null) {
288304
ArrayNode keyArray = (ArrayNode) Common.getValueByPath(fromObject, new String[] {"videos"});
@@ -313,7 +329,9 @@ ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode pare
313329
}
314330

315331
@ExcludeFromGeneratedCoverageReport
316-
ObjectNode generateVideosOperationFromVertex(JsonNode fromObject, ObjectNode parentObject) {
332+
@InternalApi
333+
public ObjectNode generateVideosOperationFromVertex(
334+
JsonNode fromObject, ObjectNode parentObject) {
317335
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
318336
if (Common.getValueByPath(fromObject, new String[] {"name"}) != null) {
319337
Common.setValueByPath(

‎src/main/java/com/google/genai/TokensConverters.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.fasterxml.jackson.databind.node.ArrayNode;
2424
import com.fasterxml.jackson.databind.node.ObjectNode;
2525

26+
/** Internal SDK converter functions. */
2627
final class TokensConverters {
2728
private final ApiClient apiClient;
2829

‎src/main/java/com/google/genai/types/GenerateVideosOperation.java‎

Lines changed: 17 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -20,48 +20,36 @@
2020

2121
import com.fasterxml.jackson.annotation.JsonCreator;
2222
import com.fasterxml.jackson.annotation.JsonProperty;
23+
import com.fasterxml.jackson.databind.JsonNode;
2324
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
2425
import com.google.auto.value.AutoValue;
2526
import com.google.genai.JsonSerializable;
27+
import com.google.genai.OperationsConverters;
2628
import java.util.Map;
2729
import java.util.Optional;
2830

2931
/** A video generation operation. */
3032
@AutoValue
3133
@JsonDeserialize(builder = GenerateVideosOperation.Builder.class)
32-
public abstract class GenerateVideosOperation extends JsonSerializable {
33-
/**
34-
* The server-assigned name, which is only unique within the same service that originally returns
35-
* it. If you use the default HTTP mapping, the `name` should be a resource name ending with
36-
* `operations/{unique_id}`.
37-
*/
38-
@JsonProperty("name")
39-
public abstract Optional<String> name();
40-
41-
/**
42-
* Service-specific metadata associated with the operation. It typically contains progress
43-
* information and common metadata such as create time. Some services might not provide such
44-
* metadata. Any method that returns a long-running operation should document the metadata type,
45-
* if any.
46-
*/
47-
@JsonProperty("metadata")
48-
public abstract Optional<Map<String, Object>> metadata();
49-
50-
/**
51-
* If the value is `false`, it means the operation is still in progress. If `true`, the operation
52-
* is completed, and either `error` or `response` is available.
53-
*/
54-
@JsonProperty("done")
55-
public abstract Optional<Boolean> done();
56-
57-
/** The error result of the operation in case of failure or cancellation. */
58-
@JsonProperty("error")
59-
public abstract Optional<Map<String, Object>> error();
60-
34+
public abstract class GenerateVideosOperation
35+
extends Operation<GenerateVideosResponse, GenerateVideosOperation> {
6136
/** The generated videos. */
6237
@JsonProperty("response")
38+
@Override
6339
public abstract Optional<GenerateVideosResponse> response();
6440

41+
@Override
42+
public GenerateVideosOperation fromApiResponse(JsonNode apiResponse, boolean isVertexAi) {
43+
OperationsConverters converter = new OperationsConverters(null);
44+
JsonNode response = null;
45+
if (isVertexAi) {
46+
response = converter.generateVideosOperationFromVertex(apiResponse, null);
47+
} else {
48+
response = converter.generateVideosOperationFromMldev(apiResponse, null);
49+
}
50+
return JsonSerializable.fromJsonNode(response, GenerateVideosOperation.class);
51+
}
52+
6553
/** Instantiates a builder for GenerateVideosOperation. */
6654
@ExcludeFromGeneratedCoverageReport
6755
public static Builder builder() {

0 commit comments

Comments
 (0)