Skip to content

Commit 14e5356

Browse files
Diwak4rDiwak4r
committed
fix(Connectors.Google): use embedContent API for gemini-embedding models
1 parent c028a0c commit 14e5356

5 files changed

Lines changed: 181 additions & 4 deletions

File tree

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,77 @@ public void ItAcceptsValidHostnameSegments(string validLocation)
206206
Assert.Null(exception);
207207
}
208208

209+
[Fact]
210+
public async Task ShouldUseBatchEmbedContentsEndpointForGeminiEmbeddingModelAsync()
211+
{
212+
// Arrange
213+
string modelId = "gemini-embedding-2";
214+
var client = this.CreateEmbeddingsClient(modelId: modelId);
215+
this._messageHandlerStub.ResponseToReturn.Content = new StringContent(
216+
File.ReadAllText("./TestData/vertex_embed_content_response.json"));
217+
IList<string> data = ["sample data"];
218+
219+
// Act
220+
await client.GenerateEmbeddingsAsync(data);
221+
222+
// Assert
223+
Assert.NotNull(this._messageHandlerStub.RequestUri);
224+
Assert.EndsWith(":batchEmbedContents", this._messageHandlerStub.RequestUri.ToString(), StringComparison.Ordinal);
225+
Assert.NotNull(this._messageHandlerStub.RequestContent);
226+
string requestBody = System.Text.Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
227+
using var requestJson = JsonDocument.Parse(requestBody);
228+
Assert.Equal(JsonValueKind.Array, requestJson.RootElement.GetProperty("requests").ValueKind);
229+
var firstRequest = requestJson.RootElement.GetProperty("requests")[0];
230+
Assert.Equal("sample data", firstRequest.GetProperty("content").GetProperty("parts")[0].GetProperty("text").GetString());
231+
Assert.False(requestJson.RootElement.TryGetProperty("instances", out _));
232+
}
233+
234+
[Fact]
235+
public async Task ShouldUsePredictEndpointForLegacyEmbeddingModelAsync()
236+
{
237+
// Arrange
238+
string modelId = "text-embedding-004";
239+
var client = this.CreateEmbeddingsClient(modelId: modelId);
240+
IList<string> data = ["sample data"];
241+
242+
// Act
243+
await client.GenerateEmbeddingsAsync(data);
244+
245+
// Assert
246+
Assert.NotNull(this._messageHandlerStub.RequestUri);
247+
Assert.EndsWith(":predict", this._messageHandlerStub.RequestUri.ToString(), StringComparison.Ordinal);
248+
Assert.NotNull(this._messageHandlerStub.RequestContent);
249+
string requestBody = System.Text.Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
250+
using var requestJson = JsonDocument.Parse(requestBody);
251+
Assert.Equal(JsonValueKind.Array, requestJson.RootElement.GetProperty("instances").ValueKind);
252+
}
253+
254+
[Fact]
255+
public async Task ShouldReturnValidEmbeddingsResponseForGeminiEmbeddingModelAsync()
256+
{
257+
// Arrange
258+
string modelId = "gemini-embedding-2";
259+
var client = this.CreateEmbeddingsClient(modelId: modelId);
260+
this._messageHandlerStub.ResponseToReturn.Content = new StringContent(
261+
File.ReadAllText("./TestData/vertex_embed_content_response.json"));
262+
var dataToEmbed = new List<string>()
263+
{
264+
"Write a story about a magic backpack.",
265+
"Print color of backpack."
266+
};
267+
268+
// Act
269+
var embeddings = await client.GenerateEmbeddingsAsync(dataToEmbed);
270+
271+
// Assert
272+
VertexAIEmbedContentResponse testDataResponse = JsonSerializer.Deserialize<VertexAIEmbedContentResponse>(
273+
await File.ReadAllTextAsync("./TestData/vertex_embed_content_response.json"))!;
274+
Assert.NotNull(embeddings);
275+
Assert.Collection(embeddings,
276+
values => Assert.Equal(testDataResponse.Embeddings[0].Values, values),
277+
values => Assert.Equal(testDataResponse.Embeddings[1].Values, values));
278+
}
279+
209280
[Fact]
210281
public async Task ShouldUseGlobalEndpointWhenLocationIsGlobalAsync()
211282
{

‎dotnet/src/Connectors/Connectors.Google.UnitTests/TestData/vertex_embed_content_response.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
{
2+
"embeddings": [
3+
{
4+
"values": [
5+
0.1,
6+
0.2,
7+
0.3
8+
]
9+
},
10+
{
11+
"values": [
12+
0.4,
13+
0.5,
14+
0.6
15+
]
16+
}
17+
]
18+
}

‎dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentRequest.cs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System.Collections.Generic;
4+
using System.Linq;
5+
using System.Text.Json.Serialization;
6+
7+
namespace Microsoft.SemanticKernel.Connectors.Google.Core;
8+
9+
internal sealed class VertexAIEmbedContentRequest
10+
{
11+
[JsonPropertyName("requests")]
12+
public IList<EmbedContentRequestItem> Requests { get; set; } = null!;
13+
14+
public static VertexAIEmbedContentRequest FromData(IEnumerable<string> data, int? dimensions = null) => new()
15+
{
16+
Requests = data.Select(text => new EmbedContentRequestItem
17+
{
18+
Content = new RequestContent
19+
{
20+
Parts =
21+
[
22+
new RequestPart
23+
{
24+
Text = text
25+
}
26+
]
27+
},
28+
OutputDimensionality = dimensions
29+
}).ToList()
30+
};
31+
32+
internal sealed class EmbedContentRequestItem
33+
{
34+
[JsonPropertyName("content")]
35+
public RequestContent Content { get; set; } = null!;
36+
37+
[JsonPropertyName("outputDimensionality")]
38+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
39+
public int? OutputDimensionality { get; set; }
40+
}
41+
42+
internal sealed class RequestContent
43+
{
44+
[JsonPropertyName("parts")]
45+
public IList<RequestPart> Parts { get; set; } = null!;
46+
}
47+
48+
internal sealed class RequestPart
49+
{
50+
[JsonPropertyName("text")]
51+
public string Text { get; set; } = null!;
52+
}
53+
}

‎dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentResponse.cs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Text.Json.Serialization;
6+
7+
namespace Microsoft.SemanticKernel.Connectors.Google.Core;
8+
9+
internal sealed class VertexAIEmbedContentResponse
10+
{
11+
[JsonPropertyName("embeddings")]
12+
[JsonRequired]
13+
public IList<ResponseEmbedding> Embeddings { get; set; } = null!;
14+
15+
internal sealed class ResponseEmbedding
16+
{
17+
[JsonPropertyName("values")]
18+
[JsonRequired]
19+
public ReadOnlyMemory<float> Values { get; set; }
20+
}
21+
}

‎dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ internal sealed class VertexAIEmbeddingClient : ClientBase
1919
private readonly string _embeddingModelId;
2020
private readonly Uri _embeddingEndpoint;
2121
private readonly int? _dimensions;
22+
private readonly bool _useEmbedContentMethod;
2223

2324
/// <summary>
2425
/// Represents a client for interacting with the embeddings models by Vertex AI.
@@ -54,10 +55,15 @@ public VertexAIEmbeddingClient(
5455
string baseUri = GetVertexAIBaseUri(location);
5556

5657
this._embeddingModelId = modelId;
57-
this._embeddingEndpoint = new Uri($"{baseUri}/{versionSubLink}/projects/{projectId}/locations/{location}/publishers/google/models/{this._embeddingModelId}:predict");
58+
this._useEmbedContentMethod = UsesEmbedContentMethod(modelId);
59+
string embeddingMethod = this._useEmbedContentMethod ? "batchEmbedContents" : "predict";
60+
this._embeddingEndpoint = new Uri($"{baseUri}/{versionSubLink}/projects/{projectId}/locations/{location}/publishers/google/models/{this._embeddingModelId}:{embeddingMethod}");
5861
this._dimensions = dimensions;
5962
}
6063

64+
private static bool UsesEmbedContentMethod(string modelId)
65+
=> modelId.StartsWith("gemini-embedding", StringComparison.Ordinal);
66+
6167
/// <summary>
6268
/// Generates embeddings for the given data asynchronously.
6369
/// </summary>
@@ -72,18 +78,26 @@ public async Task<IList<ReadOnlyMemory<float>>> GenerateEmbeddingsAsync(
7278
{
7379
Verify.NotNullOrEmpty(data);
7480

75-
var geminiRequest = this.GetEmbeddingRequest(data, options);
76-
using var httpRequestMessage = await this.CreateHttpRequestAsync(geminiRequest, this._embeddingEndpoint).ConfigureAwait(false);
81+
object request = this._useEmbedContentMethod
82+
? VertexAIEmbedContentRequest.FromData(data, options?.Dimensions ?? this._dimensions)
83+
: this.GetEmbeddingRequest(data, options);
84+
85+
using var httpRequestMessage = await this.CreateHttpRequestAsync(request, this._embeddingEndpoint).ConfigureAwait(false);
7786

7887
string body = await this.SendRequestAndGetStringBodyAsync(httpRequestMessage, cancellationToken)
7988
.ConfigureAwait(false);
8089

81-
return DeserializeAndProcessEmbeddingsResponse(body);
90+
return this._useEmbedContentMethod
91+
? ProcessEmbedContentResponse(body)
92+
: DeserializeAndProcessEmbeddingsResponse(body);
8293
}
8394

8495
private VertexAIEmbeddingRequest GetEmbeddingRequest(IEnumerable<string> data, EmbeddingGenerationOptions? options = null)
8596
=> VertexAIEmbeddingRequest.FromData(data, options?.Dimensions ?? this._dimensions);
8697

98+
private static List<ReadOnlyMemory<float>> ProcessEmbedContentResponse(string body)
99+
=> DeserializeResponse<VertexAIEmbedContentResponse>(body).Embeddings.Select(embedding => embedding.Values).ToList();
100+
87101
private static List<ReadOnlyMemory<float>> DeserializeAndProcessEmbeddingsResponse(string body)
88102
=> ProcessEmbeddingsResponse(DeserializeResponse<VertexAIEmbeddingResponse>(body));
89103

0 commit comments

Comments
 (0)