Skip to content

Commit 713d4f8

Browse files
patrickswedishpatrickswedish
committed
.NET: Support Vertex AI embedContent for Gemini embedding models
Select :embedContent endpoint and wire contract for Gemini Embedding 2 models, while preserving :predict for legacy models.
1 parent c028a0c commit 713d4f8

5 files changed

Lines changed: 539 additions & 3 deletions

File tree

‎dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIEmbedContentRequestTests.cs‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System.Text.Json;
4+
using Microsoft.Extensions.AI;
5+
using Microsoft.SemanticKernel.Connectors.Google.Core;
6+
using Xunit;
7+
8+
namespace SemanticKernel.Connectors.Google.UnitTests.Core.VertexAI;
9+
10+
public sealed class VertexAIEmbedContentRequestTests
11+
{
12+
private const string DimensionalityJsonPropertyName = "\"outputDimensionality\"";
13+
private const int Dimensions = 512;
14+
15+
[Fact]
16+
public void FromTextReturnsValidRequestWithContent()
17+
{
18+
// Arrange
19+
const string Text = "sample text to embed";
20+
21+
// Act
22+
var request = VertexAIEmbedContentRequest.FromText(Text);
23+
24+
// Assert
25+
Assert.NotNull(request.Content);
26+
Assert.NotNull(request.Content.Parts);
27+
Assert.Single(request.Content.Parts);
28+
Assert.Equal(Text, request.Content.Parts[0].Text);
29+
}
30+
31+
[Fact]
32+
public void FromTextSetsDimensionsToNullWhenNotProvided()
33+
{
34+
// Act
35+
var request = VertexAIEmbedContentRequest.FromText("sample text");
36+
37+
// Assert
38+
Assert.Null(request.OutputDimensionality);
39+
}
40+
41+
[Fact]
42+
public void FromTextJsonDoesNotIncludeDimensionsWhenNull()
43+
{
44+
// Act
45+
var request = VertexAIEmbedContentRequest.FromText("sample text");
46+
string json = JsonSerializer.Serialize(request);
47+
48+
// Assert
49+
Assert.DoesNotContain(DimensionalityJsonPropertyName, json);
50+
}
51+
52+
[Fact]
53+
public void FromTextSetsDimensionsWhenProvided()
54+
{
55+
// Act
56+
var request = VertexAIEmbedContentRequest.FromText("sample text", Dimensions);
57+
58+
// Assert
59+
Assert.Equal(Dimensions, request.OutputDimensionality);
60+
}
61+
62+
[Fact]
63+
public void FromTextJsonIncludesDimensionsWhenProvided()
64+
{
65+
// Act
66+
var request = VertexAIEmbedContentRequest.FromText("sample text", Dimensions);
67+
string json = JsonSerializer.Serialize(request);
68+
69+
// Assert
70+
Assert.Contains($"{DimensionalityJsonPropertyName}:{Dimensions}", json);
71+
}
72+
73+
[Theory]
74+
[InlineData("TaskType")]
75+
[InlineData("Task_Type")]
76+
[InlineData("taskType")]
77+
[InlineData("task_Type")]
78+
[InlineData("tasktype")]
79+
[InlineData("task_type")]
80+
public void FromTextShouldIncludeTaskTypeWhenProvided(string additionalPropertyKeyName)
81+
{
82+
// Arrange
83+
const string TaskType = "RETRIEVAL_DOCUMENT";
84+
var options = new EmbeddingGenerationOptions
85+
{
86+
AdditionalProperties = new AdditionalPropertiesDictionary
87+
{
88+
[additionalPropertyKeyName] = TaskType
89+
}
90+
};
91+
92+
// Act
93+
var request = VertexAIEmbedContentRequest.FromText("sample text", Dimensions, options);
94+
string json = JsonSerializer.Serialize(request);
95+
96+
// Assert
97+
Assert.Equal(TaskType, request.TaskType);
98+
Assert.Contains("\"taskType\":\"RETRIEVAL_DOCUMENT\"", json);
99+
}
100+
101+
[Fact]
102+
public void FromTextShouldIncludeTitleWhenProvided()
103+
{
104+
// Arrange
105+
const string Title = "Document Title";
106+
var options = new EmbeddingGenerationOptions
107+
{
108+
AdditionalProperties = new AdditionalPropertiesDictionary
109+
{
110+
["title"] = Title
111+
}
112+
};
113+
114+
// Act
115+
var request = VertexAIEmbedContentRequest.FromText("sample text", Dimensions, options);
116+
string json = JsonSerializer.Serialize(request);
117+
118+
// Assert
119+
Assert.Equal(Title, request.Title);
120+
Assert.Contains("\"title\":\"Document Title\"", json);
121+
}
122+
}

‎dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIEmbeddingEndpointTests.cs‎

Lines changed: 271 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,271 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Net;
6+
using System.Net.Http;
7+
using System.Text;
8+
using System.Threading;
9+
using System.Threading.Tasks;
10+
using Microsoft.Extensions.AI;
11+
using Microsoft.SemanticKernel;
12+
using Microsoft.SemanticKernel.Connectors.Google;
13+
using Microsoft.SemanticKernel.Connectors.Google.Core;
14+
using Xunit;
15+
16+
namespace SemanticKernel.Connectors.Google.UnitTests.Core.VertexAI;
17+
18+
public sealed class VertexAIEmbeddingEndpointTests : IDisposable
19+
{
20+
private readonly HttpMessageHandlerStub _messageHandlerStub;
21+
private readonly HttpClient _httpClient;
22+
private readonly List<IDisposable> _disposables = [];
23+
24+
public VertexAIEmbeddingEndpointTests()
25+
{
26+
this._messageHandlerStub = new HttpMessageHandlerStub();
27+
this._messageHandlerStub.ResponseToReturn.Content = new StringContent(
28+
"""
29+
{
30+
"embedding": {
31+
"values": [0.1, 0.2, 0.3]
32+
}
33+
}
34+
""",
35+
Encoding.UTF8,
36+
"application/json");
37+
this._httpClient = new HttpClient(this._messageHandlerStub, false);
38+
}
39+
40+
[Theory]
41+
[InlineData("gemini-embedding-2", true)]
42+
[InlineData("gemini-embedding-2-preview", true)]
43+
[InlineData("gemini-embedding-2-0", true)]
44+
[InlineData("GEMINI-EMBEDDING-2", true)]
45+
[InlineData("gemini-embedding-001", false)]
46+
[InlineData("textembedding-gecko", false)]
47+
[InlineData("textembedding-gecko@003", false)]
48+
[InlineData("text-embedding-004", false)]
49+
[InlineData("custom-model", false)]
50+
public void UsesEmbedContentEndpoint_ReturnsExpectedValue(string modelId, bool expected)
51+
{
52+
Assert.Equal(expected, VertexAIEmbeddingClient.UsesEmbedContentEndpoint(modelId));
53+
}
54+
55+
[Fact]
56+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_SendsCorrectEmbedContentWireContractAsync()
57+
{
58+
// Arrange
59+
var client = this.CreateClient("gemini-embedding-2");
60+
const string InputText = "hello world";
61+
62+
// Act
63+
var result = await client.GenerateEmbeddingsAsync([InputText]);
64+
65+
// Assert - URI validation
66+
Assert.NotNull(this._messageHandlerStub.RequestUri);
67+
string uri = this._messageHandlerStub.RequestUri.ToString();
68+
Assert.Contains(":embedContent", uri, StringComparison.Ordinal);
69+
Assert.DoesNotContain(":predict", uri, StringComparison.Ordinal);
70+
71+
// Assert - Request wire payload validation
72+
Assert.NotNull(this._messageHandlerStub.RequestContent);
73+
string requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
74+
Assert.Contains("\"content\"", requestBody, StringComparison.Ordinal);
75+
Assert.Contains("\"parts\"", requestBody, StringComparison.Ordinal);
76+
Assert.Contains("\"text\":\"hello world\"", requestBody, StringComparison.Ordinal);
77+
Assert.DoesNotContain("\"instances\"", requestBody, StringComparison.Ordinal);
78+
Assert.DoesNotContain("\"predictions\"", requestBody, StringComparison.Ordinal);
79+
80+
// Assert - Response parsing validation
81+
Assert.NotNull(result);
82+
Assert.Single(result);
83+
Assert.Equal(new float[] { 0.1f, 0.2f, 0.3f }, result[0].ToArray());
84+
}
85+
86+
[Fact]
87+
public async Task GenerateEmbeddingsAsync_ForLegacyModel_SendsCorrectPredictWireContractAsync()
88+
{
89+
// Arrange
90+
this._messageHandlerStub.ResponseToReturn.Content = new StringContent(
91+
"""
92+
{
93+
"predictions": [
94+
{
95+
"embeddings": {
96+
"values": [0.4, 0.5, 0.6]
97+
}
98+
}
99+
]
100+
}
101+
""",
102+
Encoding.UTF8,
103+
"application/json");
104+
var client = this.CreateClient("text-embedding-004");
105+
const string InputText = "hello legacy";
106+
107+
// Act
108+
var result = await client.GenerateEmbeddingsAsync([InputText]);
109+
110+
// Assert - URI validation
111+
Assert.NotNull(this._messageHandlerStub.RequestUri);
112+
string uri = this._messageHandlerStub.RequestUri.ToString();
113+
Assert.Contains(":predict", uri, StringComparison.Ordinal);
114+
Assert.DoesNotContain(":embedContent", uri, StringComparison.Ordinal);
115+
116+
// Assert - Request wire payload validation
117+
Assert.NotNull(this._messageHandlerStub.RequestContent);
118+
string requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
119+
Assert.Contains("\"instances\"", requestBody, StringComparison.Ordinal);
120+
Assert.Contains("\"content\":\"hello legacy\"", requestBody, StringComparison.Ordinal);
121+
Assert.DoesNotContain("\"parts\"", requestBody, StringComparison.Ordinal);
122+
123+
// Assert - Response parsing validation
124+
Assert.NotNull(result);
125+
Assert.Single(result);
126+
Assert.Equal(new float[] { 0.4f, 0.5f, 0.6f }, result[0].ToArray());
127+
}
128+
129+
[Fact]
130+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_IncludesDimensionsInPayloadWhenProvidedAsync()
131+
{
132+
// Arrange
133+
var client = this.CreateClient("gemini-embedding-2", dimensions: 256);
134+
135+
// Act
136+
await client.GenerateEmbeddingsAsync(["test with dimensions"]);
137+
138+
// Assert
139+
Assert.NotNull(this._messageHandlerStub.RequestContent);
140+
string requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
141+
Assert.Contains("\"outputDimensionality\":256", requestBody, StringComparison.Ordinal);
142+
}
143+
144+
[Fact]
145+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_OmitsDimensionsWhenNullAsync()
146+
{
147+
// Arrange
148+
var client = this.CreateClient("gemini-embedding-2", dimensions: null);
149+
150+
// Act
151+
await client.GenerateEmbeddingsAsync(["test without dimensions"]);
152+
153+
// Assert
154+
Assert.NotNull(this._messageHandlerStub.RequestContent);
155+
string requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
156+
Assert.DoesNotContain("outputDimensionality", requestBody, StringComparison.Ordinal);
157+
}
158+
159+
[Fact]
160+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_IncludesTaskTypeAndTitleFromOptionsAsync()
161+
{
162+
// Arrange
163+
var client = this.CreateClient("gemini-embedding-2");
164+
var options = new EmbeddingGenerationOptions
165+
{
166+
AdditionalProperties = new AdditionalPropertiesDictionary
167+
{
168+
["task_type"] = "RETRIEVAL_DOCUMENT",
169+
["title"] = "Document Title"
170+
}
171+
};
172+
173+
// Act
174+
await client.GenerateEmbeddingsAsync(["test with task_type and title"], options);
175+
176+
// Assert
177+
Assert.NotNull(this._messageHandlerStub.RequestContent);
178+
string requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
179+
Assert.Contains("\"taskType\":\"RETRIEVAL_DOCUMENT\"", requestBody, StringComparison.Ordinal);
180+
Assert.Contains("\"title\":\"Document Title\"", requestBody, StringComparison.Ordinal);
181+
}
182+
183+
[Fact]
184+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_MultipleInputs_SendsSequentialRequestsAndPreservesOrderAsync()
185+
{
186+
// Arrange
187+
var client = this.CreateClient("gemini-embedding-2");
188+
var response1 = this.TrackDisposable(new HttpResponseMessage(HttpStatusCode.OK)
189+
{
190+
Content = new StringContent("""{"embedding": {"values": [1.0, 1.1]}}""", Encoding.UTF8, "application/json")
191+
});
192+
var response2 = this.TrackDisposable(new HttpResponseMessage(HttpStatusCode.OK)
193+
{
194+
Content = new StringContent("""{"embedding": {"values": [2.0, 2.1]}}""", Encoding.UTF8, "application/json")
195+
});
196+
var response3 = this.TrackDisposable(new HttpResponseMessage(HttpStatusCode.OK)
197+
{
198+
Content = new StringContent("""{"embedding": {"values": [3.0, 3.1]}}""", Encoding.UTF8, "application/json")
199+
});
200+
201+
this._messageHandlerStub.ResponseQueue.Enqueue(response1);
202+
this._messageHandlerStub.ResponseQueue.Enqueue(response2);
203+
this._messageHandlerStub.ResponseQueue.Enqueue(response3);
204+
205+
// Act
206+
var results = await client.GenerateEmbeddingsAsync(["text1", "text2", "text3"]);
207+
208+
// Assert
209+
Assert.NotNull(results);
210+
Assert.Equal(3, results.Count);
211+
Assert.Equal(new float[] { 1.0f, 1.1f }, results[0].ToArray());
212+
Assert.Equal(new float[] { 2.0f, 2.1f }, results[1].ToArray());
213+
Assert.Equal(new float[] { 3.0f, 3.1f }, results[2].ToArray());
214+
}
215+
216+
[Fact]
217+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_PropagatesHttpExceptionOnFailureAsync()
218+
{
219+
// Arrange
220+
var client = this.CreateClient("gemini-embedding-2");
221+
this._messageHandlerStub.ResponseToReturn = this.TrackDisposable(new HttpResponseMessage(HttpStatusCode.InternalServerError)
222+
{
223+
Content = new StringContent("""{"error": "Internal server error"}""", Encoding.UTF8, "application/json")
224+
});
225+
226+
// Act & Assert
227+
await Assert.ThrowsAsync<HttpOperationException>(() =>
228+
client.GenerateEmbeddingsAsync(["test failing call"]));
229+
}
230+
231+
[Fact]
232+
public async Task GenerateEmbeddingsAsync_ForGeminiEmbedding2_Cancellation_ThrowsOperationCanceledExceptionAsync()
233+
{
234+
// Arrange
235+
var client = this.CreateClient("gemini-embedding-2");
236+
using var cts = new CancellationTokenSource();
237+
cts.Cancel();
238+
239+
// Act & Assert
240+
await Assert.ThrowsAnyAsync<OperationCanceledException>(() =>
241+
client.GenerateEmbeddingsAsync(["test cancelled"], cancellationToken: cts.Token));
242+
}
243+
244+
public void Dispose()
245+
{
246+
this._httpClient.Dispose();
247+
this._messageHandlerStub.Dispose();
248+
foreach (var disposable in this._disposables)
249+
{
250+
disposable.Dispose();
251+
}
252+
}
253+
254+
private T TrackDisposable<T>(T disposable) where T : IDisposable
255+
{
256+
this._disposables.Add(disposable);
257+
return disposable;
258+
}
259+
260+
private VertexAIEmbeddingClient CreateClient(string modelId, int? dimensions = null)
261+
{
262+
return new VertexAIEmbeddingClient(
263+
httpClient: this._httpClient,
264+
modelId: modelId,
265+
bearerTokenProvider: () => ValueTask.FromResult("fake-key"),
266+
apiVersion: VertexAIVersion.V1,
267+
location: "us-central1",
268+
projectId: "fake-project-id",
269+
dimensions: dimensions);
270+
}
271+
}

0 commit comments

Comments
 (0)