Add Count Tokens Samples

Author: holtskinner

genai/count_tokens/count_tokens_compute_with_txt.py

@@ -0,0 +1,35 @@
+# 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
+#
+#    https://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.
+
+
+def compute_tokens() -> int:
+    # [START googlegenaisdk_count_tokens_compute_with_txt]
+    from google import genai
+
+    client = genai.Client()
+    response = client.models.compute_tokens(
+        model="gemini-2.0-flash-001",
+        contents="What's the longest word in the English language?",
+    )
+    print(response)
+
+    # Example output:
+    # tokens_info=[TokensInfo(role='user', token_ids=[1841, 235303, 235256, 573, 32514, 2204, 575, 573, 4645, 5255, 235336], tokens=[b'What', b"'", b's', b' the', b' longest', b' word', b' in', b' the', b' English', b' language', b'?'])]
+
+    # [END googlegenaisdk_count_tokens_compute_with_txt]
+    return response.tokens_info
+
+
+if __name__ == "__main__":
+    compute_tokens()

genai/count_tokens/count_tokens_resp_with_txt.py

@@ -0,0 +1,40 @@
+# 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
+#
+#    https://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.
+
+
+def generate_content() -> int:
+    # [START googlegenaisdk_count_tokens_resp_with_txt]
+    from google import genai
+
+    client = genai.Client()
+
+    prompt = "Why is the sky blue?"
+
+    # Send text to Gemini
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001", contents=prompt
+    )
+
+    # Prompt and response tokens count
+    print(response.usage_metadata)
+
+    # Example output:
+    # cached_content_token_count=None candidates_token_count=311 prompt_token_count=6 total_token_count=317
+
+    # [END googlegenaisdk_count_tokens_resp_with_txt]
+    return response.usage_metadata
+
+
+if __name__ == "__main__":
+    generate_content()

genai/count_tokens/count_tokens_with_txt.py

@@ -0,0 +1,35 @@
+# 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
+#
+#    https://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.
+
+
+def count_tokens() -> int:
+    # [START googlegenaisdk_count_tokens_with_txt]
+    from google import genai
+
+    client = genai.Client()
+    response = client.models.count_tokens(
+        model="gemini-2.0-flash-001",
+        contents="What's the highest mountain in Africa?",
+    )
+    print(response)
+
+    # Example output:
+    # total_tokens=10 cached_content_token_count=None
+
+    # [END googlegenaisdk_count_tokens_with_txt]
+    return response.total_tokens
+
+
+if __name__ == "__main__":
+    count_tokens()

genai/count_tokens/count_tokens_with_txt_img_vid.py

@@ -0,0 +1,45 @@
+# 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
+#
+#    https://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.
+
+
+def count_tokens() -> int:
+    # [START googlegenaisdk_count_tokens_with_txt_img_vid]
+    from google import genai
+    from google.genai.types import Part
+
+    client = genai.Client()
+
+    contents = [
+        Part.from_uri(
+            file_uri="gs://cloud-samples-data/generative-ai/video/pixel8.mp4",
+            mime_type="video/mp4",
+        ),
+        "Provide a description of the video.",
+    ]
+
+    response = client.models.count_tokens(
+        model="gemini-2.0-flash-001",
+        contents=contents,
+    )
+    print(response)
+
+    # Example output:
+    # total_tokens=10 cached_content_token_count=None
+
+    # [END googlegenaisdk_count_tokens_with_txt_img_vid]
+    return response.total_tokens
+
+
+if __name__ == "__main__":
+    count_tokens()

genai/count_tokens/counttoken_with_txt.py

@@ -12,15 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import advanced_example
-import simple_example
+#
+# Using Google Cloud Vertex AI to test the code samples.
+#
+
+import os
+
+import count_tokens_compute_with_txt
+import count_tokens_resp_with_txt
+import count_tokens_with_txt
+import count_tokens_with_txt_img_vid
+
+os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
+os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
+# The project name is included in the CICD pipeline
+# os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
+
+
+def test_count_tokens_compute_with_txt() -> None:
+    assert count_tokens_compute_with_txt.compute_tokens()
+
+
+def test_count_tokens_resp_with_txt() -> None:
+    assert count_tokens_resp_with_txt.generate_content()
 
 
-def test_simple_example() -> None:
-    response = simple_example.simple_example()
-    assert response
+def test_count_tokens_with_txt() -> None:
+    assert count_tokens_with_txt.count_tokens()
 
 
-def test_advanced_example() -> None:
-    response = advanced_example.advanced_example()
-    assert response
+def test_count_tokens_with_txt_img_vid() -> None:
+    assert count_tokens_with_txt_img_vid.count_tokens()