Skip to content

Commit 34dedf9

Browse files
abbottdev-krollabbottdev-kroll
committed
feat: Add support for dimensions in Vertex AI embedding services and tests
1 parent 429dd1c commit 34dedf9

9 files changed

+383
-47
lines changed

‎dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAIEmbeddingGeneratorTests.cs‎

Lines changed: 140 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,166 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

3+
using System;
4+
using System.Collections.Generic;
5+
using System.Net.Http;
6+
using System.Text;
37
using System.Threading.Tasks;
48
using Microsoft.Extensions.AI;
59
using Microsoft.SemanticKernel.Connectors.Google;
610
using Xunit;
711

812
namespace SemanticKernel.Connectors.Google.UnitTests.Services;
913

10-
public sealed class VertexAIEmbeddingGeneratorTests
14+
public sealed class VertexAIEmbeddingGeneratorTests : IDisposable
1115
{
16+
private const string Model = "fake-model";
17+
private const string BearerKey = "fake-key";
18+
private const int Dimensions = 512;
19+
private readonly HttpMessageHandlerStub _messageHandlerStub;
20+
private readonly HttpClient _httpClient;
21+
22+
public VertexAIEmbeddingGeneratorTests()
23+
{
24+
this._messageHandlerStub = new HttpMessageHandlerStub
25+
{
26+
ResponseToReturn = new HttpResponseMessage(System.Net.HttpStatusCode.OK)
27+
{
28+
Content = new StringContent(
29+
"""
30+
{
31+
"predictions": [
32+
{
33+
"embeddings": {
34+
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
35+
}
36+
}
37+
]
38+
}
39+
""",
40+
Encoding.UTF8,
41+
"application/json"
42+
)
43+
}
44+
};
45+
46+
this._httpClient = new HttpClient(this._messageHandlerStub, disposeHandler: false);
47+
}
48+
1249
[Fact]
1350
public void AttributesShouldContainModelIdBearerAsString()
1451
{
1552
// Arrange & Act
16-
string model = "fake-model";
17-
using var service = new VertexAIEmbeddingGenerator(model, "key", "location", "project");
53+
using var service = new VertexAIEmbeddingGenerator(Model, BearerKey, "location", "project");
1854

1955
// Assert
20-
Assert.Equal(model, service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelId);
56+
Assert.Equal(Model, service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelId);
2157
}
2258

2359
[Fact]
2460
public void AttributesShouldContainModelIdBearerAsFunc()
2561
{
2662
// Arrange & Act
27-
string model = "fake-model";
28-
using var service = new VertexAIEmbeddingGenerator(model, () => ValueTask.FromResult("key"), "location", "project");
63+
using var service = new VertexAIEmbeddingGenerator(Model, () => ValueTask.FromResult(BearerKey), "location", "project");
2964

3065
// Assert
31-
Assert.Equal(model, service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelId);
66+
Assert.Equal(Model, service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelId);
67+
}
68+
69+
[Fact]
70+
public void AttributesShouldNotContainDimensionsWhenNotProvided()
71+
{
72+
// Arrange & Act
73+
using var service = new VertexAIEmbeddingGenerator(Model, BearerKey, "location", "project");
74+
75+
// Assert
76+
Assert.Null(service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelDimensions);
77+
}
78+
79+
[Fact]
80+
public void AttributesShouldContainDimensionsWhenProvided()
81+
{
82+
// Arrange & Act
83+
using var service = new VertexAIEmbeddingGenerator(Model, BearerKey, "location", "project", dimensions: Dimensions);
84+
85+
// Assert
86+
Assert.Equal(Dimensions, service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelDimensions);
87+
}
88+
89+
[Fact]
90+
public void GetDimensionsReturnsCorrectValue()
91+
{
92+
// Arrange
93+
using var service = new VertexAIEmbeddingGenerator(Model, BearerKey, "location", "project", dimensions: Dimensions);
94+
95+
// Act
96+
var result = service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelDimensions;
97+
98+
// Assert
99+
Assert.Equal(Dimensions, result);
100+
}
101+
102+
[Fact]
103+
public void GetDimensionsReturnsNullWhenNotProvided()
104+
{
105+
// Arrange
106+
using var service = new VertexAIEmbeddingGenerator(Model, BearerKey, "location", "project");
107+
108+
// Act
109+
var result = service.GetService<EmbeddingGeneratorMetadata>()!.DefaultModelDimensions;
110+
111+
// Assert
112+
Assert.Null(result);
113+
}
114+
115+
[Fact]
116+
public async Task ShouldNotIncludeDimensionsInRequestWhenNotProvidedAsync()
117+
{
118+
// Arrange
119+
using var service = new VertexAIEmbeddingGenerator(
120+
modelId: Model,
121+
bearerKey: BearerKey,
122+
location: "location",
123+
projectId: "project",
124+
dimensions: null,
125+
httpClient: this._httpClient);
126+
var dataToEmbed = new List<string> { "Text to embed" };
127+
128+
// Act
129+
await service.GenerateAsync(dataToEmbed);
130+
131+
// Assert
132+
Assert.NotNull(this._messageHandlerStub.RequestContent);
133+
var requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
134+
Assert.DoesNotContain("outputDimensionality", requestBody);
135+
}
136+
137+
[Theory]
138+
[InlineData(Dimensions)]
139+
[InlineData(Dimensions * 2)]
140+
public async Task ShouldIncludeDimensionsInRequestWhenProvidedAsync(int? dimensions)
141+
{
142+
// Arrange
143+
using var service = new VertexAIEmbeddingGenerator(
144+
modelId: Model,
145+
bearerKey: BearerKey,
146+
location: "location",
147+
projectId: "project",
148+
dimensions: dimensions,
149+
httpClient: this._httpClient);
150+
var dataToEmbed = new List<string> { "Text to embed" };
151+
152+
// Act
153+
await service.GenerateAsync(dataToEmbed);
154+
155+
// Assert
156+
Assert.NotNull(this._messageHandlerStub.RequestContent);
157+
var requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
158+
Assert.Contains($"\"outputDimensionality\":{dimensions}", requestBody);
159+
}
160+
161+
public void Dispose()
162+
{
163+
this._messageHandlerStub.Dispose();
164+
this._httpClient.Dispose();
32165
}
33166
}

‎dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAITextEmbeddingGenerationServiceTests.cs‎

Lines changed: 140 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,168 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

33
using System;
4+
using System.Collections.Generic;
5+
using System.Net.Http;
6+
using System.Text;
47
using System.Threading.Tasks;
58
using Microsoft.SemanticKernel.Connectors.Google;
9+
using Microsoft.SemanticKernel.Embeddings;
610
using Microsoft.SemanticKernel.Services;
711
using Xunit;
812

913
namespace SemanticKernel.Connectors.Google.UnitTests.Services;
1014

1115
[Obsolete("Temporary test for Obsolete ITextEmbeddingGenerationService")]
12-
public sealed class VertexAITextEmbeddingGenerationServiceTests
16+
public sealed class VertexAITextEmbeddingGenerationServiceTests : IDisposable
1317
{
18+
private const string Model = "fake-model";
19+
private const string BearerKey = "fake-key";
20+
private const int Dimensions = 512;
21+
private readonly HttpMessageHandlerStub _messageHandlerStub;
22+
private readonly HttpClient _httpClient;
23+
24+
public VertexAITextEmbeddingGenerationServiceTests()
25+
{
26+
this._messageHandlerStub = new HttpMessageHandlerStub
27+
{
28+
ResponseToReturn = new HttpResponseMessage(System.Net.HttpStatusCode.OK)
29+
{
30+
Content = new StringContent(
31+
"""
32+
{
33+
"predictions": [
34+
{
35+
"embeddings": {
36+
"values": [0.1, 0.2, 0.3, 0.4, 0.5]
37+
}
38+
}
39+
]
40+
}
41+
""",
42+
Encoding.UTF8,
43+
"application/json"
44+
)
45+
}
46+
};
47+
48+
this._httpClient = new HttpClient(this._messageHandlerStub, disposeHandler: false);
49+
}
50+
1451
[Fact]
1552
public void AttributesShouldContainModelIdBearerAsString()
1653
{
1754
// Arrange & Act
18-
string model = "fake-model";
19-
var service = new VertexAITextEmbeddingGenerationService(model, "key", "location", "project");
55+
var service = new VertexAITextEmbeddingGenerationService(Model, BearerKey, "location", "project");
2056

2157
// Assert
22-
Assert.Equal(model, service.Attributes[AIServiceExtensions.ModelIdKey]);
58+
Assert.Equal(Model, service.Attributes[AIServiceExtensions.ModelIdKey]);
2359
}
2460

2561
[Fact]
2662
public void AttributesShouldContainModelIdBearerAsFunc()
2763
{
2864
// Arrange & Act
29-
string model = "fake-model";
30-
var service = new VertexAITextEmbeddingGenerationService(model, () => ValueTask.FromResult("key"), "location", "project");
65+
var service = new VertexAITextEmbeddingGenerationService(Model, () => ValueTask.FromResult(BearerKey), "location", "project");
3166

3267
// Assert
33-
Assert.Equal(model, service.Attributes[AIServiceExtensions.ModelIdKey]);
68+
Assert.Equal(Model, service.Attributes[AIServiceExtensions.ModelIdKey]);
69+
}
70+
71+
[Fact]
72+
public void AttributesShouldNotContainDimensionsWhenNotProvided()
73+
{
74+
// Arrange & Act
75+
var service = new VertexAITextEmbeddingGenerationService(Model, BearerKey, "location", "project");
76+
77+
// Assert
78+
Assert.False(service.Attributes.ContainsKey(EmbeddingGenerationExtensions.DimensionsKey));
79+
}
80+
81+
[Fact]
82+
public void AttributesShouldContainDimensionsWhenProvided()
83+
{
84+
// Arrange & Act
85+
var service = new VertexAITextEmbeddingGenerationService(Model, BearerKey, "location", "project", dimensions: Dimensions);
86+
87+
// Assert
88+
Assert.Equal(Dimensions, service.Attributes[EmbeddingGenerationExtensions.DimensionsKey]);
89+
}
90+
91+
[Fact]
92+
public void GetDimensionsReturnsCorrectValue()
93+
{
94+
// Arrange
95+
var service = new VertexAITextEmbeddingGenerationService(Model, BearerKey, "location", "project", dimensions: Dimensions);
96+
97+
// Act
98+
var result = service.GetDimensions();
99+
100+
// Assert
101+
Assert.Equal(Dimensions, result);
102+
}
103+
104+
[Fact]
105+
public void GetDimensionsReturnsNullWhenNotProvided()
106+
{
107+
// Arrange
108+
var service = new VertexAITextEmbeddingGenerationService(Model, BearerKey, "location", "project");
109+
110+
// Act
111+
var result = service.GetDimensions();
112+
113+
// Assert
114+
Assert.Null(result);
115+
}
116+
117+
[Fact]
118+
public async Task ShouldNotIncludeDimensionsInRequestWhenNotProvidedAsync()
119+
{
120+
// Arrange
121+
var service = new VertexAITextEmbeddingGenerationService(
122+
modelId: Model,
123+
bearerKey: BearerKey,
124+
location: "location",
125+
projectId: "project",
126+
dimensions: null,
127+
httpClient: this._httpClient);
128+
var dataToEmbed = new List<string> { "Text to embed" };
129+
130+
// Act
131+
await service.GenerateEmbeddingsAsync(dataToEmbed);
132+
133+
// Assert
134+
Assert.NotNull(this._messageHandlerStub.RequestContent);
135+
var requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
136+
Assert.DoesNotContain("outputDimensionality", requestBody);
137+
}
138+
139+
[Theory]
140+
[InlineData(Dimensions)]
141+
[InlineData(Dimensions * 2)]
142+
public async Task ShouldIncludeDimensionsInRequestWhenProvidedAsync(int? dimensions)
143+
{
144+
// Arrange
145+
var service = new VertexAITextEmbeddingGenerationService(
146+
modelId: Model,
147+
bearerKey: BearerKey,
148+
location: "location",
149+
projectId: "project",
150+
dimensions: dimensions,
151+
httpClient: this._httpClient);
152+
var dataToEmbed = new List<string> { "Text to embed" };
153+
154+
// Act
155+
await service.GenerateEmbeddingsAsync(dataToEmbed);
156+
157+
// Assert
158+
Assert.NotNull(this._messageHandlerStub.RequestContent);
159+
var requestBody = Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
160+
Assert.Contains($"\"outputDimensionality\":{dimensions}", requestBody);
161+
}
162+
163+
public void Dispose()
164+
{
165+
this._messageHandlerStub.Dispose();
166+
this._httpClient.Dispose();
34167
}
35168
}

0 commit comments

Comments
 (0)