Skip to content

Commit fc71c1a

Browse files
yyyu-googlecopybara-github
yyyu-google
authored and
copybara-github
committed
feat: support count tokens feature in Models module
PiperOrigin-RevId: 833619706
1 parent 0c3adca commit fc71c1a

12 files changed

+1243
-0
lines changed

‎Google.GenAI.E2E.Tests/CountTokens/CountTokensTest.cs‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
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+
using System;
18+
using System.Collections.Generic;
19+
using System.Linq;
20+
using System.Threading.Tasks;
21+
22+
using Google.GenAI;
23+
using Google.GenAI.Types;
24+
25+
using Microsoft.VisualStudio.TestTools.UnitTesting;
26+
27+
using TestServerSdk;
28+
29+
[TestClass]
30+
public class CountTokensTest {
31+
private static TestServerProcess? _server;
32+
private Client vertexClient;
33+
private Client geminiClient;
34+
private string modelName;
35+
public TestContext TestContext { get; set; }
36+
37+
[ClassInitialize]
38+
public static void ClassInit(TestContext _) {
39+
_server = TestServer.StartTestServer();
40+
}
41+
42+
[ClassCleanup]
43+
public static void ClassCleanup() {
44+
TestServer.StopTestServer(_server);
45+
}
46+
47+
[TestInitialize]
48+
public void TestInit() {
49+
// Test server specific setup.
50+
if (_server == null) {
51+
throw new InvalidOperationException("Test server is not initialized.");
52+
}
53+
var geminiClientHttpOptions = new HttpOptions {
54+
Headers = new Dictionary<string, string> { { "Test-Name",
55+
$"{GetType().Name}.{TestContext.TestName}" } },
56+
BaseUrl = "http://localhost:1453"
57+
};
58+
var vertexClientHttpOptions = new HttpOptions {
59+
Headers = new Dictionary<string, string> { { "Test-Name",
60+
$"{GetType().Name}.{TestContext.TestName}" } },
61+
BaseUrl = "http://localhost:1454"
62+
};
63+
64+
// Common setup for both clients.
65+
string project = System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_PROJECT");
66+
string location =
67+
System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_LOCATION") ?? "us-central1";
68+
string apiKey = System.Environment.GetEnvironmentVariable("GOOGLE_API_KEY");
69+
vertexClient = new Client(project: project, location: location, vertexAI: true,
70+
credential: TestServer.GetCredentialForTestMode(),
71+
httpOptions: vertexClientHttpOptions);
72+
geminiClient =
73+
new Client(apiKey: apiKey, vertexAI: false, httpOptions: geminiClientHttpOptions);
74+
75+
// Specific setup for this test class
76+
modelName = "gemini-2.5-flash";
77+
}
78+
79+
[TestMethod]
80+
public async Task CountTokensVertexTest() {
81+
var contents = new List<Content>
82+
{
83+
new Content
84+
{
85+
Role = "user",
86+
Parts = new List<Part>
87+
{
88+
new Part
89+
{
90+
Text = "What is the capital of France?"
91+
}
92+
}
93+
}
94+
};
95+
var vertexResponse = await vertexClient.Models.CountTokensAsync(
96+
model: modelName, contents: contents, config: null);
97+
98+
Assert.IsNotNull(vertexResponse.TotalTokens);
99+
}
100+
101+
[TestMethod]
102+
public async Task CountTokensGeminiTest() {
103+
var contents = new List<Content>
104+
{
105+
new Content
106+
{
107+
Role = "user",
108+
Parts = new List<Part>
109+
{
110+
new Part
111+
{
112+
Text = "What is the capital of France?"
113+
}
114+
}
115+
}
116+
};
117+
var geminiResponse = await geminiClient.Models.CountTokensAsync(
118+
model: modelName, contents: contents, config: null);
119+
120+
Assert.IsNotNull(geminiResponse.TotalTokens);
121+
}
122+
123+
[TestMethod]
124+
public async Task CountTokensContentVertexTest() {
125+
var contents = new Content
126+
{
127+
Role = "user",
128+
Parts = new List<Part>
129+
{
130+
new Part
131+
{
132+
Text = "What is the capital of France?"
133+
}
134+
}
135+
};
136+
var vertexResponse = await vertexClient.Models.CountTokensAsync(
137+
model: modelName, contents: contents);
138+
139+
Assert.IsNotNull(vertexResponse.TotalTokens);
140+
}
141+
142+
[TestMethod]
143+
public async Task CountTokensContentGeminiTest() {
144+
var contents = new Content
145+
{
146+
Role = "user",
147+
Parts = new List<Part>
148+
{
149+
new Part
150+
{
151+
Text = "What is the capital of France?"
152+
}
153+
}
154+
};
155+
var geminiResponse = await geminiClient.Models.CountTokensAsync(
156+
model: modelName, contents: contents);
157+
158+
Assert.IsNotNull(geminiResponse.TotalTokens);
159+
}
160+
161+
[TestMethod]
162+
public async Task CountTokensStringVertexTest() {
163+
var vertexResponse = await vertexClient.Models.CountTokensAsync(
164+
model: modelName, contents: "What is the capital of France?");
165+
166+
Assert.IsNotNull(vertexResponse.TotalTokens);
167+
}
168+
169+
[TestMethod]
170+
public async Task CountTokensStringGeminiTest() {
171+
var geminiResponse = await geminiClient.Models.CountTokensAsync(
172+
model: modelName, contents: "What is the capital of France?");
173+
174+
Assert.IsNotNull(geminiResponse.TotalTokens);
175+
}
176+
177+
[TestMethod]
178+
public async Task CountTokensConfigVertexTest() {
179+
var config = new CountTokensConfig
180+
{
181+
SystemInstruction = new Content
182+
{
183+
Parts = new List<Part>
184+
{
185+
new Part
186+
{
187+
Text = "You are a helpful assistant."
188+
}
189+
}
190+
}
191+
};
192+
var vertexResponse = await vertexClient.Models.CountTokensAsync(
193+
model: modelName, contents: "What is the capital of France?", config: config);
194+
195+
Assert.IsNotNull(vertexResponse.TotalTokens);
196+
}
197+
198+
[TestMethod]
199+
public async Task CountTokensConfigGeminiTest() {
200+
var config = new CountTokensConfig
201+
{
202+
SystemInstruction = new Content
203+
{
204+
Parts = new List<Part>
205+
{
206+
new Part
207+
{
208+
Text = "You are a helpful assistant."
209+
}
210+
}
211+
}
212+
};
213+
var ex = await Assert.ThrowsExceptionAsync<NotSupportedException>(async () => {
214+
await geminiClient.Models.CountTokensAsync(
215+
model: modelName, contents: "What is the capital of France?", config: config);
216+
});
217+
218+
StringAssert.Contains(ex.Message, "systemInstruction");
219+
StringAssert.Contains(ex.Message, "not supported in Gemini API");
220+
}
221+
}

‎Google.GenAI.E2E.Tests/Recordings/CountTokensTest.CountTokensConfigVertexTest.json‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
{
2+
"recordID": "CountTokensTest.CountTokensConfigVertexTest",
3+
"interactions": [
4+
{
5+
"request": {
6+
"method": "POST",
7+
"url": "/v1beta1/projects/REDACTED/locations/us-central1/publishers/google/models/gemini-2.5-flash:countTokens",
8+
"request": "POST /v1beta1/projects/REDACTED/locations/us-central1/publishers/google/models/gemini-2.5-flash:countTokens HTTP/1.1",
9+
"headers": {
10+
"Content-Length": "154",
11+
"Content-Type": "application/json; charset=utf-8",
12+
"Test-Name": "CountTokensTest.CountTokensConfigVertexTest"
13+
},
14+
"bodySegments": [
15+
{
16+
"contents": [
17+
{
18+
"parts": [
19+
{
20+
"text": "What is the capital of France?"
21+
}
22+
],
23+
"role": "user"
24+
}
25+
],
26+
"systemInstruction": {
27+
"parts": [
28+
{
29+
"text": "You are a helpful assistant."
30+
}
31+
]
32+
}
33+
}
34+
],
35+
"previousRequest": "b4d6e60a9b97e7b98c63df9308728c5c88c0b40c398046772c63447b94608b4d",
36+
"serverAddress": "us-central1-aiplatform.googleapis.com",
37+
"port": 443,
38+
"protocol": "https"
39+
},
40+
"shaSum": "fd721119f6ac02e8f92a3bef3156934ed0f425917c6ea5e8af48cd847269c354",
41+
"response": {
42+
"statusCode": 200,
43+
"headers": {
44+
"Content-Type": "application/json; charset=UTF-8",
45+
"Date": "Tue, 18 Nov 2025 04:52:19 GMT",
46+
"Server": "scaffolding on HTTPServer2",
47+
"Vary": "Origin, X-Origin, Referer",
48+
"X-Content-Type-Options": "nosniff",
49+
"X-Frame-Options": "SAMEORIGIN",
50+
"X-Xss-Protection": "0"
51+
},
52+
"bodySegments": [
53+
{
54+
"promptTokensDetails": [
55+
{
56+
"modality": "TEXT",
57+
"tokenCount": 13
58+
}
59+
],
60+
"totalBillableCharacters": 49,
61+
"totalTokens": 13
62+
}
63+
]
64+
}
65+
}
66+
]
67+
}

‎Google.GenAI.E2E.Tests/Recordings/CountTokensTest.CountTokensContentGeminiTest.json‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
{
2+
"recordID": "CountTokensTest.CountTokensContentGeminiTest",
3+
"interactions": [
4+
{
5+
"request": {
6+
"method": "POST",
7+
"url": "/v1beta/models/gemini-2.5-flash:countTokens",
8+
"request": "POST /v1beta/models/gemini-2.5-flash:countTokens HTTP/1.1",
9+
"headers": {
10+
"Content-Length": "82",
11+
"Content-Type": "application/json; charset=utf-8",
12+
"Test-Name": "CountTokensTest.CountTokensContentGeminiTest"
13+
},
14+
"bodySegments": [
15+
{
16+
"contents": [
17+
{
18+
"parts": [
19+
{
20+
"text": "What is the capital of France?"
21+
}
22+
],
23+
"role": "user"
24+
}
25+
]
26+
}
27+
],
28+
"previousRequest": "b4d6e60a9b97e7b98c63df9308728c5c88c0b40c398046772c63447b94608b4d",
29+
"serverAddress": "generativelanguage.googleapis.com",
30+
"port": 443,
31+
"protocol": "https"
32+
},
33+
"shaSum": "6f82a5a817d114c72184e8c84277173b4ba2ed15d23acfe6f3f2f58278f81f6f",
34+
"response": {
35+
"statusCode": 200,
36+
"headers": {
37+
"Content-Type": "application/json; charset=UTF-8",
38+
"Date": "Tue, 18 Nov 2025 04:27:12 GMT",
39+
"Server": "scaffolding on HTTPServer2",
40+
"Server-Timing": "gfet4t7; dur=53",
41+
"Vary": "Origin, X-Origin, Referer",
42+
"X-Content-Type-Options": "nosniff",
43+
"X-Frame-Options": "SAMEORIGIN",
44+
"X-Xss-Protection": "0"
45+
},
46+
"bodySegments": [
47+
{
48+
"promptTokensDetails": [
49+
{
50+
"modality": "TEXT",
51+
"tokenCount": 8
52+
}
53+
],
54+
"totalTokens": 8
55+
}
56+
]
57+
}
58+
}
59+
]
60+
}

0 commit comments

Comments
 (0)