Skip to content

Commit 5f70751

Browse files
hkt74copybara-github
hkt74
authored and
copybara-github
committed
feat: Added Operations.get which is a generic method which will handle all Operation types.
PiperOrigin-RevId: 821849754
1 parent c205d01 commit 5f70751

File tree

10 files changed

+168
-91
lines changed

10 files changed

+168
-91
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
@@ -92,8 +92,7 @@ public static void main(String[] args) {
9292
while (!generateVideosOperation.done().filter(Boolean::booleanValue).isPresent()) {
9393
try {
9494
Thread.sleep(10000); // Sleep for 10 seconds.
95-
generateVideosOperation =
96-
client.operations.getVideosOperation(generateVideosOperation, null);
95+
generateVideosOperation = client.operations.getVideosOperation(generateVideosOperation, null);
9796
System.out.println("Waiting for operation to complete...");
9897
} catch (InterruptedException e) {
9998
System.out.println("Thread was interrupted while sleeping.");

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ public static void main(String[] args) {
9898
try {
9999
Thread.sleep(10000); // Sleep for 10 seconds.
100100
try {
101-
operation = client.async.operations.getVideosOperation(operation, null).get();
101+
operation = client.async.operations.get(operation, null).get();
102102
} catch (ExecutionException e) {
103103
throw new RuntimeException(e);
104104
}
@@ -115,7 +115,7 @@ public static void main(String[] args) {
115115

116116
Video generatedVideo =
117117
operation.response().get().generatedVideos().get().get(0).video().get();
118-
// Do something with the video.
118+
System.out.println("Video URL: " + generatedVideo.uri().get());
119119
})
120120
.join();
121121
}

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

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@
1818

1919
package com.google.genai;
2020

21+
import com.fasterxml.jackson.databind.JsonNode;
2122
import com.google.genai.Common.BuiltRequest;
2223
import com.google.genai.types.FetchPredictOperationConfig;
2324
import com.google.genai.types.GenerateVideosOperation;
2425
import com.google.genai.types.GetOperationConfig;
26+
import com.google.genai.types.Operation;
2527
import java.util.concurrent.CompletableFuture;
2628

2729
/** Async module of {@link Operations} */
@@ -35,7 +37,7 @@ public AsyncOperations(ApiClient apiClient) {
3537
this.operations = new Operations(apiClient);
3638
}
3739

38-
CompletableFuture<GenerateVideosOperation> privateGetVideosOperation(
40+
CompletableFuture<JsonNode> privateGetVideosOperation(
3941
String operationName, GetOperationConfig config) {
4042
BuiltRequest builtRequest =
4143
operations.buildRequestForPrivateGetVideosOperation(operationName, config);
@@ -49,7 +51,7 @@ CompletableFuture<GenerateVideosOperation> privateGetVideosOperation(
4951
});
5052
}
5153

52-
CompletableFuture<GenerateVideosOperation> privateFetchPredictVideosOperation(
54+
CompletableFuture<JsonNode> privateFetchPredictVideosOperation(
5355
String operationName, String resourceName, FetchPredictOperationConfig config) {
5456
BuiltRequest builtRequest =
5557
operations.buildRequestForPrivateFetchPredictVideosOperation(
@@ -73,19 +75,29 @@ CompletableFuture<GenerateVideosOperation> privateFetchPredictVideosOperation(
7375
*/
7476
public CompletableFuture<GenerateVideosOperation> getVideosOperation(
7577
GenerateVideosOperation operation, GetOperationConfig config) {
78+
return (CompletableFuture<GenerateVideosOperation>) get(operation, config);
79+
}
80+
81+
/**
82+
* Gets the status of an Operation.
83+
*
84+
* @param operation An Operation.
85+
* @param config The configuration for getting the operation.
86+
* @return An Operation with the updated status of the operation.
87+
*/
88+
public <T, U extends Operation<T, U>> CompletableFuture<U> get(
89+
U operation, GetOperationConfig config) {
7690
if (!operation.name().isPresent()) {
7791
throw new Error("Operation name is required.");
7892
}
7993

8094
if (this.apiClient.vertexAI()) {
8195
String resourceName = operation.name().get().split("/operations/")[0];
82-
83-
FetchPredictOperationConfig fetchConfig = FetchPredictOperationConfig.builder().build();
84-
85-
return this.privateFetchPredictVideosOperation(
86-
operation.name().get(), resourceName, fetchConfig);
96+
return this.privateFetchPredictVideosOperation(operation.name().get(), resourceName, null)
97+
.thenApplyAsync(response -> operation.fromApiResponse(response, true));
8798
} else {
88-
return this.privateGetVideosOperation(operation.name().get(), config);
99+
return this.privateGetVideosOperation(operation.name().get(), config)
100+
.thenApplyAsync(response -> operation.fromApiResponse(response, false));
89101
}
90102
}
91103
}

‎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: 23 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import com.google.genai.types.GetOperationConfig;
3131
import com.google.genai.types.GetOperationParameters;
3232
import com.google.genai.types.HttpOptions;
33+
import com.google.genai.types.Operation;
3334
import java.io.IOException;
3435
import java.util.Optional;
3536
import okhttp3.ResponseBody;
@@ -378,7 +379,7 @@ BuiltRequest buildRequestForPrivateGetVideosOperation(
378379
}
379380

380381
/** A shared processResponse function for both sync and async methods. */
381-
GenerateVideosOperation processResponseForPrivateGetVideosOperation(
382+
JsonNode processResponseForPrivateGetVideosOperation(
382383
ApiResponse response, GetOperationConfig config) {
383384
ResponseBody responseBody = response.getBody();
384385
String responseString;
@@ -388,21 +389,10 @@ GenerateVideosOperation processResponseForPrivateGetVideosOperation(
388389
throw new GenAiIOException("Failed to read HTTP response.", e);
389390
}
390391

391-
JsonNode responseNode = JsonSerializable.stringToJsonNode(responseString);
392-
393-
if (this.apiClient.vertexAI()) {
394-
responseNode = generateVideosOperationFromVertex(responseNode, null);
395-
}
396-
397-
if (!this.apiClient.vertexAI()) {
398-
responseNode = generateVideosOperationFromMldev(responseNode, null);
399-
}
400-
401-
return JsonSerializable.fromJsonNode(responseNode, GenerateVideosOperation.class);
392+
return JsonSerializable.stringToJsonNode(responseString);
402393
}
403394

404-
GenerateVideosOperation privateGetVideosOperation(
405-
String operationName, GetOperationConfig config) {
395+
JsonNode privateGetVideosOperation(String operationName, GetOperationConfig config) {
406396
BuiltRequest builtRequest = buildRequestForPrivateGetVideosOperation(operationName, config);
407397

408398
try (ApiResponse response =
@@ -457,7 +447,7 @@ BuiltRequest buildRequestForPrivateFetchPredictVideosOperation(
457447
}
458448

459449
/** A shared processResponse function for both sync and async methods. */
460-
GenerateVideosOperation processResponseForPrivateFetchPredictVideosOperation(
450+
JsonNode processResponseForPrivateFetchPredictVideosOperation(
461451
ApiResponse response, FetchPredictOperationConfig config) {
462452
ResponseBody responseBody = response.getBody();
463453
String responseString;
@@ -467,21 +457,10 @@ GenerateVideosOperation processResponseForPrivateFetchPredictVideosOperation(
467457
throw new GenAiIOException("Failed to read HTTP response.", e);
468458
}
469459

470-
JsonNode responseNode = JsonSerializable.stringToJsonNode(responseString);
471-
472-
if (this.apiClient.vertexAI()) {
473-
responseNode = generateVideosOperationFromVertex(responseNode, null);
474-
}
475-
476-
if (!this.apiClient.vertexAI()) {
477-
throw new UnsupportedOperationException(
478-
"This method is only supported in the Vertex AI client.");
479-
}
480-
481-
return JsonSerializable.fromJsonNode(responseNode, GenerateVideosOperation.class);
460+
return JsonSerializable.stringToJsonNode(responseString);
482461
}
483462

484-
GenerateVideosOperation privateFetchPredictVideosOperation(
463+
JsonNode privateFetchPredictVideosOperation(
485464
String operationName, String resourceName, FetchPredictOperationConfig config) {
486465
BuiltRequest builtRequest =
487466
buildRequestForPrivateFetchPredictVideosOperation(operationName, resourceName, config);
@@ -502,20 +481,29 @@ GenerateVideosOperation privateFetchPredictVideosOperation(
502481
*/
503482
public GenerateVideosOperation getVideosOperation(
504483
GenerateVideosOperation operation, GetOperationConfig config) {
484+
return (GenerateVideosOperation) get(operation, config);
485+
}
505486

487+
/**
488+
* Gets the status of an Operation.
489+
*
490+
* @param operation An Operation.
491+
* @param config The configuration for getting the operation.
492+
* @return An Operation with the updated status of the operation.
493+
*/
494+
public <T, U extends Operation<T, U>> U get(U operation, GetOperationConfig config) {
506495
if (!operation.name().isPresent()) {
507-
throw new Error("Operation name is required.");
496+
throw new IllegalArgumentException("Operation name is required.");
508497
}
509498

510499
if (this.apiClient.vertexAI()) {
511500
String resourceName = operation.name().get().split("/operations/")[0];
512-
513-
FetchPredictOperationConfig fetchConfig = FetchPredictOperationConfig.builder().build();
514-
515-
return this.privateFetchPredictVideosOperation(
516-
operation.name().get(), resourceName, fetchConfig);
501+
JsonNode response =
502+
this.privateFetchPredictVideosOperation(operation.name().get(), resourceName, null);
503+
return operation.fromApiResponse(response, true);
517504
} else {
518-
return this.privateGetVideosOperation(operation.name().get(), config);
505+
JsonNode response = this.privateGetVideosOperation(operation.name().get(), config);
506+
return operation.fromApiResponse(response, false);
519507
}
520508
}
521509
}

‎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,9 @@ ObjectNode fetchPredictOperationParametersToMldev(JsonNode fromObject, ObjectNod
4954
}
5055

5156
@ExcludeFromGeneratedCoverageReport
52-
ObjectNode fetchPredictOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
57+
@InternalApi
58+
public ObjectNode fetchPredictOperationParametersToVertex(
59+
JsonNode fromObject, ObjectNode parentObject) {
5360
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
5461
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
5562
Common.setValueByPath(
@@ -69,7 +76,8 @@ ObjectNode fetchPredictOperationParametersToVertex(JsonNode fromObject, ObjectNo
6976
}
7077

7178
@ExcludeFromGeneratedCoverageReport
72-
ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode parentObject) {
79+
@InternalApi
80+
public ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode parentObject) {
7381
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
7482
if (Common.getValueByPath(fromObject, new String[] {"name"}) != null) {
7583
Common.setValueByPath(
@@ -115,7 +123,9 @@ ObjectNode generateVideosOperationFromMldev(JsonNode fromObject, ObjectNode pare
115123
}
116124

117125
@ExcludeFromGeneratedCoverageReport
118-
ObjectNode generateVideosOperationFromVertex(JsonNode fromObject, ObjectNode parentObject) {
126+
@InternalApi
127+
public ObjectNode generateVideosOperationFromVertex(
128+
JsonNode fromObject, ObjectNode parentObject) {
119129
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
120130
if (Common.getValueByPath(fromObject, new String[] {"name"}) != null) {
121131
Common.setValueByPath(
@@ -159,7 +169,8 @@ ObjectNode generateVideosOperationFromVertex(JsonNode fromObject, ObjectNode par
159169
}
160170

161171
@ExcludeFromGeneratedCoverageReport
162-
ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
172+
@InternalApi
173+
public ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
163174
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
164175
if (Common.getValueByPath(fromObject, new String[] {"generatedSamples"}) != null) {
165176
ArrayNode keyArray =
@@ -191,7 +202,8 @@ ObjectNode generateVideosResponseFromMldev(JsonNode fromObject, ObjectNode paren
191202
}
192203

193204
@ExcludeFromGeneratedCoverageReport
194-
ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode parentObject) {
205+
@InternalApi
206+
public ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode parentObject) {
195207
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
196208
if (Common.getValueByPath(fromObject, new String[] {"videos"}) != null) {
197209
ArrayNode keyArray = (ArrayNode) Common.getValueByPath(fromObject, new String[] {"videos"});
@@ -222,7 +234,8 @@ ObjectNode generateVideosResponseFromVertex(JsonNode fromObject, ObjectNode pare
222234
}
223235

224236
@ExcludeFromGeneratedCoverageReport
225-
ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
237+
@InternalApi
238+
public ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
226239
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
227240
if (Common.getValueByPath(fromObject, new String[] {"video"}) != null) {
228241
Common.setValueByPath(
@@ -238,7 +251,8 @@ ObjectNode generatedVideoFromMldev(JsonNode fromObject, ObjectNode parentObject)
238251
}
239252

240253
@ExcludeFromGeneratedCoverageReport
241-
ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
254+
@InternalApi
255+
public ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
242256
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
243257
if (Common.getValueByPath(fromObject, new String[] {"_self"}) != null) {
244258
Common.setValueByPath(
@@ -254,7 +268,8 @@ ObjectNode generatedVideoFromVertex(JsonNode fromObject, ObjectNode parentObject
254268
}
255269

256270
@ExcludeFromGeneratedCoverageReport
257-
ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentObject) {
271+
@InternalApi
272+
public ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentObject) {
258273
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
259274
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
260275
Common.setValueByPath(
@@ -267,7 +282,8 @@ ObjectNode getOperationParametersToMldev(JsonNode fromObject, ObjectNode parentO
267282
}
268283

269284
@ExcludeFromGeneratedCoverageReport
270-
ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
285+
@InternalApi
286+
public ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parentObject) {
271287
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
272288
if (Common.getValueByPath(fromObject, new String[] {"operationName"}) != null) {
273289
Common.setValueByPath(
@@ -280,7 +296,8 @@ ObjectNode getOperationParametersToVertex(JsonNode fromObject, ObjectNode parent
280296
}
281297

282298
@ExcludeFromGeneratedCoverageReport
283-
ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
299+
@InternalApi
300+
public ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
284301
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
285302
if (Common.getValueByPath(fromObject, new String[] {"uri"}) != null) {
286303
Common.setValueByPath(
@@ -305,7 +322,8 @@ ObjectNode videoFromMldev(JsonNode fromObject, ObjectNode parentObject) {
305322
}
306323

307324
@ExcludeFromGeneratedCoverageReport
308-
ObjectNode videoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
325+
@InternalApi
326+
public ObjectNode videoFromVertex(JsonNode fromObject, ObjectNode parentObject) {
309327
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
310328
if (Common.getValueByPath(fromObject, new String[] {"gcsUri"}) != null) {
311329
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

0 commit comments

Comments
 (0)