Skip to content

Commit 169a681

Browse files
speedstorm1copybara-github
authored andcommitted
feat: Populate per-modality prompt token count in embedding responses for gemini-embedding-2
PiperOrigin-RevId: 940552765
1 parent 631f4e0 commit 169a681

2 files changed

Lines changed: 65 additions & 0 deletions

File tree

src/main/java/com/google/genai/Models.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,13 @@ ObjectNode contentEmbeddingStatisticsFromVertex(
469469
Common.getValueByPath(fromObject, new String[] {"token_count"}));
470470
}
471471

472+
if (Common.getValueByPath(fromObject, new String[] {"tokensDetails"}) != null) {
473+
Common.setValueByPath(
474+
toObject,
475+
new String[] {"tokensDetails"},
476+
Common.getValueByPath(fromObject, new String[] {"tokensDetails"}));
477+
}
478+
472479
return toObject;
473480
}
474481

@@ -1349,6 +1356,9 @@ ObjectNode embedContentResponseFromVertex(
13491356
if (usageMetadata != null && usageMetadata.get("promptTokenCount") != null) {
13501357
stats.set("token_count", usageMetadata.get("promptTokenCount"));
13511358
}
1359+
if (usageMetadata != null && usageMetadata.get("promptTokensDetails") != null) {
1360+
stats.set("tokensDetails", usageMetadata.get("promptTokensDetails"));
1361+
}
13521362
if (truncated != null) {
13531363
stats.set("truncated", truncated);
13541364
}

src/main/java/com/google/genai/types/ContentEmbeddingStatistics.java

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,16 @@
1818

1919
package com.google.genai.types;
2020

21+
import static com.google.common.collect.ImmutableList.toImmutableList;
22+
2123
import com.fasterxml.jackson.annotation.JsonCreator;
2224
import com.fasterxml.jackson.annotation.JsonProperty;
2325
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
2426
import com.google.auto.value.AutoValue;
2527
import com.google.errorprone.annotations.CanIgnoreReturnValue;
2628
import com.google.genai.JsonSerializable;
29+
import java.util.Arrays;
30+
import java.util.List;
2731
import java.util.Optional;
2832

2933
/** Statistics of the input text associated with the result of content embedding. */
@@ -41,6 +45,13 @@ public abstract class ContentEmbeddingStatistics extends JsonSerializable {
4145
@JsonProperty("tokenCount")
4246
public abstract Optional<Float> tokenCount();
4347

48+
/**
49+
* Gemini Enterprise Agent Platform only. List of modalities and their token count for the input
50+
* content.
51+
*/
52+
@JsonProperty("tokensDetails")
53+
public abstract Optional<List<ModalityTokenCount>> tokensDetails();
54+
4455
/** Instantiates a builder for ContentEmbeddingStatistics. */
4556
@ExcludeFromGeneratedCoverageReport
4657
public static Builder builder() {
@@ -96,6 +107,50 @@ public Builder clearTokenCount() {
96107
return tokenCount(Optional.empty());
97108
}
98109

110+
/**
111+
* Setter for tokensDetails.
112+
*
113+
* <p>tokensDetails: Gemini Enterprise Agent Platform only. List of modalities and their token
114+
* count for the input content.
115+
*/
116+
@JsonProperty("tokensDetails")
117+
public abstract Builder tokensDetails(List<ModalityTokenCount> tokensDetails);
118+
119+
/**
120+
* Setter for tokensDetails.
121+
*
122+
* <p>tokensDetails: Gemini Enterprise Agent Platform only. List of modalities and their token
123+
* count for the input content.
124+
*/
125+
@CanIgnoreReturnValue
126+
public Builder tokensDetails(ModalityTokenCount... tokensDetails) {
127+
return tokensDetails(Arrays.asList(tokensDetails));
128+
}
129+
130+
/**
131+
* Setter for tokensDetails builder.
132+
*
133+
* <p>tokensDetails: Gemini Enterprise Agent Platform only. List of modalities and their token
134+
* count for the input content.
135+
*/
136+
@CanIgnoreReturnValue
137+
public Builder tokensDetails(ModalityTokenCount.Builder... tokensDetailsBuilders) {
138+
return tokensDetails(
139+
Arrays.asList(tokensDetailsBuilders).stream()
140+
.map(ModalityTokenCount.Builder::build)
141+
.collect(toImmutableList()));
142+
}
143+
144+
@ExcludeFromGeneratedCoverageReport
145+
abstract Builder tokensDetails(Optional<List<ModalityTokenCount>> tokensDetails);
146+
147+
/** Clears the value of tokensDetails field. */
148+
@ExcludeFromGeneratedCoverageReport
149+
@CanIgnoreReturnValue
150+
public Builder clearTokensDetails() {
151+
return tokensDetails(Optional.empty());
152+
}
153+
99154
public abstract ContentEmbeddingStatistics build();
100155
}
101156

0 commit comments

Comments
 (0)