feat: add new GenAI SDK count tokens samples (#10145)

* remove unused imports

* feat: add new GenAI count tokens samples

* feat: add new tests for count tokens samples

* refactor: add new asserts to make sure CountTokensReponseWithText return data

* refactor: update compute tokens to change logic that prints tokens info

* fix controlled generation test case

---------

Co-authored-by: Juan Dominguez <[email protected]>

Author: jdomingr

genai/snippets/src/main/java/genai/counttokens/CountTokensComputeWithText.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed 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 genai.counttokens;
+
+// [START googlegenaisdk_counttoken_compute_with_txt]
+
+import com.google.genai.Client;
+import com.google.genai.types.ComputeTokensResponse;
+import com.google.genai.types.HttpOptions;
+import com.google.genai.types.TokensInfo;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Optional;
+
+public class CountTokensComputeWithText {
+
+  public static void main(String[] args) {
+    // TODO(developer): Replace these variables before running the sample.
+    String modelId = "gemini-2.5-flash";
+    computeTokens(modelId);
+  }
+
+  // Computes tokens with text input
+  public static Optional<List<TokensInfo>> computeTokens(String modelId) {
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests.
+    try (Client client =
+        Client.builder()
+            .location("global")
+            .vertexAI(true)
+            .httpOptions(HttpOptions.builder().apiVersion("v1").build())
+            .build()) {
+
+      ComputeTokensResponse response = client.models.computeTokens(
+              modelId, "What's the longest word in the English language?", null);
+
+      // Print TokensInfo
+      response.tokensInfo().ifPresent(tokensInfoList -> {
+        for (TokensInfo info : tokensInfoList) {
+          info.role().ifPresent(role -> System.out.println("role: " + role));
+          info.tokenIds().ifPresent(tokenIds -> System.out.println("tokenIds: " + tokenIds));
+          // print tokens input as strings since they are in a form of byte array
+          System.out.println("tokens: ");
+          info.tokens().ifPresent(tokens ->
+              tokens.forEach(token ->
+                  System.out.println(new String(token, StandardCharsets.UTF_8))
+              )
+          );
+        }
+      });
+      // Example response.tokensInfo()
+      // role: user
+      // tokenIds: [1841, 235303, 235256, 573, 32514, 2204, 575, 573, 4645, 5255, 235336]
+      // tokens:
+      // What
+      // '
+      // s
+      // the
+      return response.tokensInfo();
+    }
+  }
+}
+// [END googlegenaisdk_counttoken_compute_with_txt]

genai/snippets/src/main/java/genai/counttokens/CountTokensResponseWithText.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed 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 genai.counttokens;
+
+// [START googlegenaisdk_counttoken_resp_with_txt]
+
+import com.google.genai.Client;
+import com.google.genai.types.GenerateContentResponse;
+import com.google.genai.types.GenerateContentResponseUsageMetadata;
+import com.google.genai.types.HttpOptions;
+import java.util.Optional;
+
+public class CountTokensResponseWithText {
+
+  public static void main(String[] args) {
+    // TODO(developer): Replace these variables before running the sample.
+    String modelId = "gemini-2.5-flash";
+    countTokens(modelId);
+  }
+
+  // Generates content response usage metadata that contains prompt and response token counts
+  public static Optional<GenerateContentResponseUsageMetadata> countTokens(String modelId) {
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests.
+    try (Client client =
+        Client.builder()
+            .location("global")
+            .vertexAI(true)
+            .httpOptions(HttpOptions.builder().apiVersion("v1").build())
+            .build()) {
+
+      GenerateContentResponse response =
+              client.models.generateContent(modelId, "Why is the sky blue?", null);
+
+      response.usageMetadata().ifPresent(System.out::println);
+      // Example response:
+      // GenerateContentResponseUsageMetadata{cacheTokensDetails=Optional.empty,
+      // cachedContentTokenCount=Optional.empty, candidatesTokenCount=Optional[569],
+      // candidatesTokensDetails=Optional[[ModalityTokenCount{modality=Optional[TEXT],
+      // tokenCount=Optional[569]}]], promptTokenCount=Optional[6],
+      // promptTokensDetails=Optional[[ModalityTokenCount{modality=Optional[TEXT],
+      // tokenCount=Optional[6]}]], thoughtsTokenCount=Optional[1132],
+      // toolUsePromptTokenCount=Optional.empty, toolUsePromptTokensDetails=Optional.empty,
+      // totalTokenCount=Optional[1707], trafficType=Optional[ON_DEMAND]}
+      return response.usageMetadata();
+    }
+  }
+}
+// [END googlegenaisdk_counttoken_resp_with_txt]

genai/snippets/src/main/java/genai/counttokens/CountTokensWithText.java

@@ -19,11 +19,8 @@
 // [START googlegenaisdk_counttoken_with_txt]
 
 import com.google.genai.Client;
-import com.google.genai.types.Content;
 import com.google.genai.types.CountTokensResponse;
 import com.google.genai.types.HttpOptions;
-import com.google.genai.types.Part;
-import java.util.List;
 import java.util.Optional;
 
 public class CountTokensWithText {

genai/snippets/src/test/java/genai/counttokens/CountTokensIT.java

@@ -19,8 +19,12 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import com.google.genai.types.GenerateContentResponseUsageMetadata;
+import com.google.genai.types.TokensInfo;
 import java.io.ByteArrayOutputStream;
 import java.io.PrintStream;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import org.junit.After;
 import org.junit.Before;
@@ -32,7 +36,6 @@
 @RunWith(JUnit4.class)
 public class CountTokensIT {
 
-
   private static final String GEMINI_FLASH = "gemini-2.5-flash";
   private ByteArrayOutputStream bout;
   private PrintStream out;
@@ -74,4 +77,37 @@ public void testCountTokensWithTextAndVideo() {
     assertThat(response).isPresent();
     assertThat(response.get()).isGreaterThan(6);
   }
+
+  @Test
+  public void testCountTokensComputeWithText() {
+
+    List<TokensInfo> response =
+            CountTokensComputeWithText.computeTokens(GEMINI_FLASH).orElse(new ArrayList<>());
+
+    assertThat(response).isNotEmpty();
+    TokensInfo tokensInfo = response.get(0);
+
+    assertThat(tokensInfo.role()).isPresent();
+
+    assertThat(tokensInfo.tokenIds()).isPresent();
+    assertThat(tokensInfo.tokenIds().get()).isNotEmpty();
+
+    assertThat(tokensInfo.tokens()).isPresent();
+    assertThat(tokensInfo.tokens().get()).isNotEmpty();
+
+  }
+
+  @Test
+  public void testCountTokensResponseWithText() {
+
+    Optional<GenerateContentResponseUsageMetadata> response =
+            CountTokensResponseWithText.countTokens(GEMINI_FLASH);
+
+    assertThat(response).isPresent();
+    assertThat(response.get().totalTokenCount()).isPresent();
+    assertThat(response.get().totalTokenCount().get()).isGreaterThan(0);
+    assertThat(response.get().promptTokenCount()).isPresent();
+    assertThat(response.get().promptTokenCount().get()).isGreaterThan(0);
+
+  }
 }