Skip to content

Commit 12ddfce

Browse files
MarkDaoustcopybara-github
MarkDaoust
authored and
copybara-github
committed
feat: [Python] Implement async embedding batches for MLDev.
feat: [Python] Add BatchJob.done property. - Vertex's interface is generic enough that it already works with embedding batches. - MLDev's interface uses a separate method for embedding batches. - The SDK uses needs different argument names to indicate which method to call. - The different argument/result field names are also important to avoid the use of Union Types with inlinedRequests and inlinedResponses. PiperOrigin-RevId: 804658916
1 parent a76ec37 commit 12ddfce

File tree

6 files changed

+380
-5
lines changed

6 files changed

+380
-5
lines changed

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

Lines changed: 67 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2040,6 +2040,68 @@ ObjectNode inlinedResponseFromMldev(JsonNode fromObject, ObjectNode parentObject
20402040
return toObject;
20412041
}
20422042

2043+
@ExcludeFromGeneratedCoverageReport
2044+
ObjectNode contentEmbeddingFromMldev(JsonNode fromObject, ObjectNode parentObject) {
2045+
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
2046+
if (Common.getValueByPath(fromObject, new String[] {"values"}) != null) {
2047+
Common.setValueByPath(
2048+
toObject,
2049+
new String[] {"values"},
2050+
Common.getValueByPath(fromObject, new String[] {"values"}));
2051+
}
2052+
2053+
return toObject;
2054+
}
2055+
2056+
@ExcludeFromGeneratedCoverageReport
2057+
ObjectNode singleEmbedContentResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
2058+
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
2059+
if (Common.getValueByPath(fromObject, new String[] {"embedding"}) != null) {
2060+
Common.setValueByPath(
2061+
toObject,
2062+
new String[] {"embedding"},
2063+
contentEmbeddingFromMldev(
2064+
JsonSerializable.toJsonNode(
2065+
Common.getValueByPath(fromObject, new String[] {"embedding"})),
2066+
toObject));
2067+
}
2068+
2069+
if (Common.getValueByPath(fromObject, new String[] {"tokenCount"}) != null) {
2070+
Common.setValueByPath(
2071+
toObject,
2072+
new String[] {"tokenCount"},
2073+
Common.getValueByPath(fromObject, new String[] {"tokenCount"}));
2074+
}
2075+
2076+
return toObject;
2077+
}
2078+
2079+
@ExcludeFromGeneratedCoverageReport
2080+
ObjectNode inlinedEmbedContentResponseFromMldev(JsonNode fromObject, ObjectNode parentObject) {
2081+
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
2082+
if (Common.getValueByPath(fromObject, new String[] {"response"}) != null) {
2083+
Common.setValueByPath(
2084+
toObject,
2085+
new String[] {"response"},
2086+
singleEmbedContentResponseFromMldev(
2087+
JsonSerializable.toJsonNode(
2088+
Common.getValueByPath(fromObject, new String[] {"response"})),
2089+
toObject));
2090+
}
2091+
2092+
if (Common.getValueByPath(fromObject, new String[] {"error"}) != null) {
2093+
Common.setValueByPath(
2094+
toObject,
2095+
new String[] {"error"},
2096+
jobErrorFromMldev(
2097+
JsonSerializable.toJsonNode(
2098+
Common.getValueByPath(fromObject, new String[] {"error"})),
2099+
toObject));
2100+
}
2101+
2102+
return toObject;
2103+
}
2104+
20432105
@ExcludeFromGeneratedCoverageReport
20442106
ObjectNode batchJobDestinationFromMldev(JsonNode fromObject, ObjectNode parentObject) {
20452107
ObjectNode toObject = JsonSerializable.objectMapper.createObjectNode();
@@ -2128,7 +2190,8 @@ ObjectNode batchJobFromMldev(JsonNode fromObject, ObjectNode parentObject) {
21282190
new String[] {"dest"},
21292191
batchJobDestinationFromMldev(
21302192
JsonSerializable.toJsonNode(
2131-
Common.getValueByPath(fromObject, new String[] {"metadata", "output"})),
2193+
Transformers.tRecvBatchJobDestination(
2194+
Common.getValueByPath(fromObject, new String[] {"metadata", "output"}))),
21322195
toObject));
21332196
}
21342197

@@ -2372,7 +2435,8 @@ ObjectNode batchJobFromVertex(JsonNode fromObject, ObjectNode parentObject) {
23722435
new String[] {"dest"},
23732436
batchJobDestinationFromVertex(
23742437
JsonSerializable.toJsonNode(
2375-
Common.getValueByPath(fromObject, new String[] {"outputConfig"})),
2438+
Transformers.tRecvBatchJobDestination(
2439+
Common.getValueByPath(fromObject, new String[] {"outputConfig"}))),
23762440
toObject));
23772441
}
23782442

@@ -2818,7 +2882,7 @@ public BatchJob create(String model, BatchJobSource src, CreateBatchJobConfig co
28182882
throw new GenAiIOException("Only one of fileName and InlinedRequests can be set.");
28192883
}
28202884
if (!src.fileName().isPresent() && !src.inlinedRequests().isPresent()) {
2821-
throw new GenAiIOException("One of fileName and InlinedRequests must be set.");
2885+
throw new GenAiIOException("one of fileName and InlinedRequests must be set.");
28222886
}
28232887
}
28242888
return this.privateCreate(model, src, config);

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,10 @@ public static Object tBatchJobDestination(Object dest) {
578578
return dest;
579579
}
580580

581+
public static Object tRecvBatchJobDestination(Object dest) {
582+
return dest;
583+
}
584+
581585
/**
582586
* It validates and extracts the batch job name based on the backend (Vertex AI or MLDev).
583587
*

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
// Auto-generated code. Do not edit.
18+
19+
package com.google.genai.types;
20+
21+
import static com.google.common.collect.ImmutableList.toImmutableList;
22+
23+
import com.fasterxml.jackson.annotation.JsonCreator;
24+
import com.fasterxml.jackson.annotation.JsonProperty;
25+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
26+
import com.google.auto.value.AutoValue;
27+
import com.google.genai.JsonSerializable;
28+
import java.util.Arrays;
29+
import java.util.List;
30+
import java.util.Optional;
31+
32+
/** Parameters for the embed_content method. */
33+
@AutoValue
34+
@JsonDeserialize(builder = EmbedContentBatch.Builder.class)
35+
public abstract class EmbedContentBatch extends JsonSerializable {
36+
/** The content to embed. Only the `parts.text` fields will be counted. */
37+
@JsonProperty("contents")
38+
public abstract Optional<List<Content>> contents();
39+
40+
/** Configuration that contains optional parameters. */
41+
@JsonProperty("config")
42+
public abstract Optional<EmbedContentConfig> config();
43+
44+
/** Instantiates a builder for EmbedContentBatch. */
45+
@ExcludeFromGeneratedCoverageReport
46+
public static Builder builder() {
47+
return new AutoValue_EmbedContentBatch.Builder();
48+
}
49+
50+
/** Creates a builder with the same values as this instance. */
51+
public abstract Builder toBuilder();
52+
53+
/** Builder for EmbedContentBatch. */
54+
@AutoValue.Builder
55+
public abstract static class Builder {
56+
/** For internal usage. Please use `EmbedContentBatch.builder()` for instantiation. */
57+
@JsonCreator
58+
private static Builder create() {
59+
return new AutoValue_EmbedContentBatch.Builder();
60+
}
61+
62+
/**
63+
* Setter for contents.
64+
*
65+
* <p>contents: The content to embed. Only the `parts.text` fields will be counted.
66+
*/
67+
@JsonProperty("contents")
68+
public abstract Builder contents(List<Content> contents);
69+
70+
/**
71+
* Setter for contents.
72+
*
73+
* <p>contents: The content to embed. Only the `parts.text` fields will be counted.
74+
*/
75+
public Builder contents(Content... contents) {
76+
return contents(Arrays.asList(contents));
77+
}
78+
79+
/**
80+
* Setter for contents builder.
81+
*
82+
* <p>contents: The content to embed. Only the `parts.text` fields will be counted.
83+
*/
84+
public Builder contents(Content.Builder... contentsBuilders) {
85+
return contents(
86+
Arrays.asList(contentsBuilders).stream()
87+
.map(Content.Builder::build)
88+
.collect(toImmutableList()));
89+
}
90+
91+
/**
92+
* Setter for config.
93+
*
94+
* <p>config: Configuration that contains optional parameters.
95+
*/
96+
@JsonProperty("config")
97+
public abstract Builder config(EmbedContentConfig config);
98+
99+
/**
100+
* Setter for config builder.
101+
*
102+
* <p>config: Configuration that contains optional parameters.
103+
*/
104+
public Builder config(EmbedContentConfig.Builder configBuilder) {
105+
return config(configBuilder.build());
106+
}
107+
108+
public abstract EmbedContentBatch build();
109+
}
110+
111+
/** Deserializes a JSON string to a EmbedContentBatch object. */
112+
@ExcludeFromGeneratedCoverageReport
113+
public static EmbedContentBatch fromJson(String jsonString) {
114+
return JsonSerializable.fromJsonString(jsonString, EmbedContentBatch.class);
115+
}
116+
}

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
// Auto-generated code. Do not edit.
18+
19+
package com.google.genai.types;
20+
21+
import com.fasterxml.jackson.annotation.JsonCreator;
22+
import com.fasterxml.jackson.annotation.JsonProperty;
23+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
24+
import com.google.auto.value.AutoValue;
25+
import com.google.genai.JsonSerializable;
26+
import java.util.Optional;
27+
28+
/** Config for `inlined_embedding_responses` parameter. */
29+
@AutoValue
30+
@JsonDeserialize(builder = InlinedEmbedContentResponse.Builder.class)
31+
public abstract class InlinedEmbedContentResponse extends JsonSerializable {
32+
/** The response to the request. */
33+
@JsonProperty("response")
34+
public abstract Optional<SingleEmbedContentResponse> response();
35+
36+
/** The error encountered while processing the request. */
37+
@JsonProperty("error")
38+
public abstract Optional<JobError> error();
39+
40+
/** Instantiates a builder for InlinedEmbedContentResponse. */
41+
@ExcludeFromGeneratedCoverageReport
42+
public static Builder builder() {
43+
return new AutoValue_InlinedEmbedContentResponse.Builder();
44+
}
45+
46+
/** Creates a builder with the same values as this instance. */
47+
public abstract Builder toBuilder();
48+
49+
/** Builder for InlinedEmbedContentResponse. */
50+
@AutoValue.Builder
51+
public abstract static class Builder {
52+
/** For internal usage. Please use `InlinedEmbedContentResponse.builder()` for instantiation. */
53+
@JsonCreator
54+
private static Builder create() {
55+
return new AutoValue_InlinedEmbedContentResponse.Builder();
56+
}
57+
58+
/**
59+
* Setter for response.
60+
*
61+
* <p>response: The response to the request.
62+
*/
63+
@JsonProperty("response")
64+
public abstract Builder response(SingleEmbedContentResponse response);
65+
66+
/**
67+
* Setter for response builder.
68+
*
69+
* <p>response: The response to the request.
70+
*/
71+
public Builder response(SingleEmbedContentResponse.Builder responseBuilder) {
72+
return response(responseBuilder.build());
73+
}
74+
75+
/**
76+
* Setter for error.
77+
*
78+
* <p>error: The error encountered while processing the request.
79+
*/
80+
@JsonProperty("error")
81+
public abstract Builder error(JobError error);
82+
83+
/**
84+
* Setter for error builder.
85+
*
86+
* <p>error: The error encountered while processing the request.
87+
*/
88+
public Builder error(JobError.Builder errorBuilder) {
89+
return error(errorBuilder.build());
90+
}
91+
92+
public abstract InlinedEmbedContentResponse build();
93+
}
94+
95+
/** Deserializes a JSON string to a InlinedEmbedContentResponse object. */
96+
@ExcludeFromGeneratedCoverageReport
97+
public static InlinedEmbedContentResponse fromJson(String jsonString) {
98+
return JsonSerializable.fromJsonString(jsonString, InlinedEmbedContentResponse.class);
99+
}
100+
}

0 commit comments

Comments
 (0)