Skip to content

Commit 63beb49

Browse files
speedstorm1copybara-github
authored andcommitted
feat: Support Batches module for C# SDK (Create, CreateEmbeddings, Get, Cancel, List, Delete)
PiperOrigin-RevId: 835424117
1 parent dce6173 commit 63beb49

File tree

48 files changed

+6874
-7
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+6874
-7
lines changed
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+
using System;
17+
using System.Collections.Generic;
18+
using System.Linq;
19+
using System.Threading.Tasks;
20+
using Google.GenAI;
21+
using Google.GenAI.Types;
22+
using Microsoft.VisualStudio.TestTools.UnitTesting;
23+
using TestServerSdk;
24+
25+
26+
[TestClass]
27+
public class CancelBatchTest
28+
{
29+
private static TestServerProcess? _server;
30+
private Client vertexClient;
31+
private Client geminiClient;
32+
private string modelName;
33+
public TestContext TestContext { get; set; }
34+
35+
[ClassInitialize]
36+
public static void ClassInit(TestContext _)
37+
{
38+
_server = TestServer.StartTestServer();
39+
}
40+
41+
[ClassCleanup]
42+
public static void ClassCleanup()
43+
{
44+
TestServer.StopTestServer(_server);
45+
}
46+
47+
[TestInitialize]
48+
public void TestInit()
49+
{
50+
// Test server specific setup.
51+
if (_server == null)
52+
{
53+
throw new InvalidOperationException("Test server is not initialized.");
54+
}
55+
var geminiClientHttpOptions = new HttpOptions {
56+
Headers = new Dictionary<string, string> { { "Test-Name",
57+
$"{GetType().Name}.{TestContext.TestName}" } },
58+
BaseUrl = "http://localhost:1453"
59+
};
60+
var vertexClientHttpOptions = new HttpOptions {
61+
Headers = new Dictionary<string, string> { { "Test-Name",
62+
$"{GetType().Name}.{TestContext.TestName}" } },
63+
BaseUrl = "http://localhost:1454"
64+
};
65+
66+
// Common setup for both clients.
67+
string project = System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_PROJECT");
68+
string location =
69+
System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_LOCATION") ?? "us-central1";
70+
string apiKey = System.Environment.GetEnvironmentVariable("GOOGLE_API_KEY");
71+
vertexClient = new Client(project: project, location: location, vertexAI: true,
72+
credential: TestServer.GetCredentialForTestMode(),
73+
httpOptions: vertexClientHttpOptions);
74+
geminiClient =
75+
new Client(apiKey: apiKey, vertexAI: false, httpOptions: geminiClientHttpOptions);
76+
77+
// Specific setup for this test class
78+
modelName = "gemini-2.5-flash";
79+
}
80+
81+
[TestMethod]
82+
public async Task CancelBatchJobVertexTest()
83+
{
84+
var src = new BatchJobSource
85+
{
86+
GcsUri = new List<string> { "gs://unified-genai-tests/batches/input/generate_content_requests.jsonl" },
87+
Format = "jsonl"
88+
};
89+
var config = new CreateBatchJobConfig
90+
{
91+
DisplayName = "test_batch_gcs_cancel",
92+
Dest = new BatchJobDestination
93+
{
94+
GcsUri = "gs://unified-genai-tests/batches/output",
95+
Format = "jsonl"
96+
}
97+
};
98+
var response = await vertexClient.Batches.CreateAsync(modelName, src, config);
99+
await vertexClient.Batches.CancelAsync(response.Name);
100+
}
101+
102+
[TestMethod]
103+
public async Task CancelBatchJobGeminiTest()
104+
{
105+
var src = new BatchJobSource
106+
{
107+
FileName = "files/oi5wupwi0y7e"
108+
};
109+
var config = new CreateBatchJobConfig
110+
{
111+
DisplayName = "test_batch_file_cancel"
112+
};
113+
var createResponse = await geminiClient.Batches.CreateAsync(modelName, src, config);
114+
await geminiClient.Batches.CancelAsync(createResponse.Name);
115+
}
116+
}
Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
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+
using System;
17+
using System.Collections.Generic;
18+
using System.Linq;
19+
using System.Threading.Tasks;
20+
using Google.GenAI;
21+
using Google.GenAI.Types;
22+
using Microsoft.VisualStudio.TestTools.UnitTesting;
23+
using TestServerSdk;
24+
25+
[TestClass]
26+
public class CreateBatchTest
27+
{
28+
private static TestServerProcess? _server;
29+
private Client vertexClient;
30+
private Client geminiClient;
31+
private string modelName;
32+
public TestContext TestContext { get; set; }
33+
34+
[ClassInitialize]
35+
public static void ClassInit(TestContext _)
36+
{
37+
_server = TestServer.StartTestServer();
38+
}
39+
40+
[ClassCleanup]
41+
public static void ClassCleanup()
42+
{
43+
TestServer.StopTestServer(_server);
44+
}
45+
46+
[TestInitialize]
47+
public void TestInit()
48+
{
49+
// Test server specific setup.
50+
if (_server == null)
51+
{
52+
throw new InvalidOperationException("Test server is not initialized.");
53+
}
54+
var geminiClientHttpOptions = new HttpOptions {
55+
Headers = new Dictionary<string, string> { { "Test-Name",
56+
$"{GetType().Name}.{TestContext.TestName}" } },
57+
BaseUrl = "http://localhost:1453"
58+
};
59+
var vertexClientHttpOptions = new HttpOptions {
60+
Headers = new Dictionary<string, string> { { "Test-Name",
61+
$"{GetType().Name}.{TestContext.TestName}" } },
62+
BaseUrl = "http://localhost:1454"
63+
};
64+
65+
// Common setup for both clients.
66+
string project = System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_PROJECT");
67+
string location =
68+
System.Environment.GetEnvironmentVariable("GOOGLE_CLOUD_LOCATION") ?? "us-central1";
69+
string apiKey = System.Environment.GetEnvironmentVariable("GOOGLE_API_KEY");
70+
vertexClient = new Client(project: project, location: location, vertexAI: true,
71+
credential: TestServer.GetCredentialForTestMode(),
72+
httpOptions: vertexClientHttpOptions);
73+
geminiClient =
74+
new Client(apiKey: apiKey, vertexAI: false, httpOptions: geminiClientHttpOptions);
75+
76+
// Specific setup for this test class
77+
modelName = "gemini-2.5-flash";
78+
}
79+
80+
[TestMethod]
81+
public async Task CreateBatchWithInlinedRequestsGeminiTest()
82+
{
83+
var safetySettings = new List<SafetySetting>
84+
{
85+
new SafetySetting { Category = HarmCategory.HARM_CATEGORY_HATE_SPEECH, Threshold = HarmBlockThreshold.BLOCK_ONLY_HIGH },
86+
new SafetySetting { Category = HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, Threshold = HarmBlockThreshold.BLOCK_LOW_AND_ABOVE },
87+
};
88+
var inlineRequest = new InlinedRequest
89+
{
90+
Contents = new List<Content>
91+
{
92+
new Content
93+
{
94+
Parts = new List<Part> { new Part { Text = "Hello!" } },
95+
Role = "user"
96+
}
97+
},
98+
Metadata = new Dictionary<string, string> { { "key", "request-1" } },
99+
Config = new GenerateContentConfig
100+
{
101+
SafetySettings = safetySettings
102+
}
103+
};
104+
var src = new BatchJobSource
105+
{
106+
InlinedRequests = new List<InlinedRequest> { inlineRequest }
107+
};
108+
var config = new CreateBatchJobConfig
109+
{
110+
DisplayName = "test_batch_inlined"
111+
};
112+
113+
var response = await geminiClient.Batches.CreateAsync(modelName, src, config);
114+
Assert.IsNotNull(response);
115+
Assert.IsTrue(response.Name.StartsWith("batches/"));
116+
}
117+
118+
[TestMethod]
119+
public async Task CreateBatchWithInlinedRequestsVertexTest()
120+
{
121+
var src = new BatchJobSource
122+
{
123+
InlinedRequests = new List<InlinedRequest> { new InlinedRequest { Contents = new List<Content> { new Content { Parts = new List<Part> { new Part { Text = "Hello!" } } } } } }
124+
};
125+
await Assert.ThrowsExceptionAsync<NotSupportedException>(async () =>
126+
{
127+
await vertexClient.Batches.CreateAsync(modelName, src, null);
128+
});
129+
}
130+
131+
[TestMethod]
132+
public async Task CreateBatchWithGcsVertexTest()
133+
{
134+
var src = new BatchJobSource
135+
{
136+
GcsUri = new List<string> { "gs://unified-genai-tests/batches/input/generate_content_requests.jsonl" },
137+
Format = "jsonl"
138+
};
139+
var config = new CreateBatchJobConfig
140+
{
141+
DisplayName = "test_batch_gcs",
142+
Dest = new BatchJobDestination
143+
{
144+
GcsUri = "gs://unified-genai-tests/batches/output",
145+
Format = "jsonl"
146+
}
147+
};
148+
var response = await vertexClient.Batches.CreateAsync(modelName, src, config);
149+
Assert.IsNotNull(response);
150+
Assert.IsTrue(response.Name.StartsWith("projects/"));
151+
}
152+
153+
[TestMethod]
154+
public async Task CreateBatchWithGcsGeminiTest()
155+
{
156+
var src = new BatchJobSource
157+
{
158+
GcsUri = new List<string> { "gs://unified-genai-tests/batches/input/generate_content_requests.jsonl" },
159+
Format = "jsonl"
160+
};
161+
await Assert.ThrowsExceptionAsync<ArgumentException>(async () =>
162+
{
163+
await geminiClient.Batches.CreateAsync(modelName, src, null);
164+
});
165+
}
166+
167+
[TestMethod]
168+
public async Task CreateBatchWithBigQueryVertexTest()
169+
{
170+
if (TestServer.IsReplayMode)
171+
{
172+
Assert.Inconclusive("Vertex BigQuery batch tests run in record mode only.");
173+
}
174+
var src = new BatchJobSource
175+
{
176+
BigqueryUri = "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input",
177+
Format = "bigquery"
178+
};
179+
var config = new CreateBatchJobConfig
180+
{
181+
DisplayName = "test_batch_bigquery",
182+
Dest = new BatchJobDestination
183+
{
184+
BigqueryUri = "bq://REDACTED.unified_genai_tests_batches.generate_content_output",
185+
Format = "bigquery"
186+
}
187+
};
188+
var response = await vertexClient.Batches.CreateAsync(modelName, src, config);
189+
Assert.IsNotNull(response);
190+
Assert.IsTrue(response.Name.StartsWith("projects/"));
191+
}
192+
193+
[TestMethod]
194+
public async Task CreateBatchWithBigQueryGeminiTest()
195+
{
196+
var src = new BatchJobSource
197+
{
198+
BigqueryUri = "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input",
199+
Format = "bigquery"
200+
};
201+
await Assert.ThrowsExceptionAsync<ArgumentException>(async () =>
202+
{
203+
await geminiClient.Batches.CreateAsync(modelName, src, null);
204+
});
205+
}
206+
207+
[TestMethod]
208+
public async Task CreateBatchWithFileVertexTest()
209+
{
210+
var src = new BatchJobSource
211+
{
212+
FileName = "files/s0pa54alni6w"
213+
};
214+
await Assert.ThrowsExceptionAsync<NotSupportedException>(async () =>
215+
{
216+
await vertexClient.Batches.CreateAsync(modelName, src, null);
217+
});
218+
}
219+
220+
[TestMethod]
221+
public async Task CreateBatchWithFileGeminiTest()
222+
{
223+
var src = new BatchJobSource
224+
{
225+
FileName = "files/s0pa54alni6w"
226+
};
227+
var config = new CreateBatchJobConfig
228+
{
229+
DisplayName = "test_batch_file"
230+
};
231+
var response = await geminiClient.Batches.CreateAsync(modelName, src, config);
232+
Assert.IsNotNull(response);
233+
Assert.IsTrue(response.Name.StartsWith("batches/"));
234+
}
235+
}

0 commit comments

Comments
 (0)