diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs index 7928428448ef..3be2311b53bf 100644 --- a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs +++ b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs @@ -206,6 +206,77 @@ public void ItAcceptsValidHostnameSegments(string validLocation) Assert.Null(exception); } + [Fact] + public async Task ShouldUseBatchEmbedContentsEndpointForGeminiEmbeddingModelAsync() + { + // Arrange + string modelId = "gemini-embedding-2"; + var client = this.CreateEmbeddingsClient(modelId: modelId); + this._messageHandlerStub.ResponseToReturn.Content = new StringContent( + File.ReadAllText("./TestData/vertex_embed_content_response.json")); + IList data = ["sample data"]; + + // Act + await client.GenerateEmbeddingsAsync(data); + + // Assert + Assert.NotNull(this._messageHandlerStub.RequestUri); + Assert.EndsWith(":batchEmbedContents", this._messageHandlerStub.RequestUri.ToString(), StringComparison.Ordinal); + Assert.NotNull(this._messageHandlerStub.RequestContent); + string requestBody = System.Text.Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent); + using var requestJson = JsonDocument.Parse(requestBody); + Assert.Equal(JsonValueKind.Array, requestJson.RootElement.GetProperty("requests").ValueKind); + var firstRequest = requestJson.RootElement.GetProperty("requests")[0]; + Assert.Equal("sample data", firstRequest.GetProperty("content").GetProperty("parts")[0].GetProperty("text").GetString()); + Assert.False(requestJson.RootElement.TryGetProperty("instances", out _)); + } + + [Fact] + public async Task ShouldUsePredictEndpointForLegacyEmbeddingModelAsync() + { + // Arrange + string modelId = "text-embedding-004"; + var client = this.CreateEmbeddingsClient(modelId: modelId); + IList data = ["sample data"]; + + // Act + await client.GenerateEmbeddingsAsync(data); + + // Assert + Assert.NotNull(this._messageHandlerStub.RequestUri); + Assert.EndsWith(":predict", this._messageHandlerStub.RequestUri.ToString(), StringComparison.Ordinal); + Assert.NotNull(this._messageHandlerStub.RequestContent); + string requestBody = System.Text.Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent); + using var requestJson = JsonDocument.Parse(requestBody); + Assert.Equal(JsonValueKind.Array, requestJson.RootElement.GetProperty("instances").ValueKind); + } + + [Fact] + public async Task ShouldReturnValidEmbeddingsResponseForGeminiEmbeddingModelAsync() + { + // Arrange + string modelId = "gemini-embedding-2"; + var client = this.CreateEmbeddingsClient(modelId: modelId); + this._messageHandlerStub.ResponseToReturn.Content = new StringContent( + File.ReadAllText("./TestData/vertex_embed_content_response.json")); + var dataToEmbed = new List() + { + "Write a story about a magic backpack.", + "Print color of backpack." + }; + + // Act + var embeddings = await client.GenerateEmbeddingsAsync(dataToEmbed); + + // Assert + VertexAIEmbedContentResponse testDataResponse = JsonSerializer.Deserialize( + await File.ReadAllTextAsync("./TestData/vertex_embed_content_response.json"))!; + Assert.NotNull(embeddings); + Assert.Collection(embeddings, + values => Assert.Equal(testDataResponse.Embeddings[0].Values, values), + values => Assert.Equal(testDataResponse.Embeddings[1].Values, values)); + } + [Fact] public async Task ShouldUseGlobalEndpointWhenLocationIsGlobalAsync() { diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/TestData/vertex_embed_content_response.json b/dotnet/src/Connectors/Connectors.Google.UnitTests/TestData/vertex_embed_content_response.json new file mode 100644 index 000000000000..282984010ead --- /dev/null +++ b/dotnet/src/Connectors/Connectors.Google.UnitTests/TestData/vertex_embed_content_response.json @@ -0,0 +1,18 @@ +{ + "embeddings": [ + { + "values": [ + 0.1, + 0.2, + 0.3 + ] + }, + { + "values": [ + 0.4, + 0.5, + 0.6 + ] + } + ] +} diff --git a/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentRequest.cs b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentRequest.cs new file mode 100644 index 000000000000..8877623fc8cd --- /dev/null +++ b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentRequest.cs @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Linq; +using System.Text.Json.Serialization; + +namespace Microsoft.SemanticKernel.Connectors.Google.Core; + +internal sealed class VertexAIEmbedContentRequest +{ + [JsonPropertyName("requests")] + public IList Requests { get; set; } = null!; + + public static VertexAIEmbedContentRequest FromData(IEnumerable data, int? dimensions = null) => new() + { + Requests = data.Select(text => new EmbedContentRequestItem + { + Content = new RequestContent + { + Parts = + [ + new RequestPart + { + Text = text + } + ] + }, + OutputDimensionality = dimensions + }).ToList() + }; + + internal sealed class EmbedContentRequestItem + { + [JsonPropertyName("content")] + public RequestContent Content { get; set; } = null!; + + [JsonPropertyName("outputDimensionality")] + [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)] + public int? OutputDimensionality { get; set; } + } + + internal sealed class RequestContent + { + [JsonPropertyName("parts")] + public IList Parts { get; set; } = null!; + } + + internal sealed class RequestPart + { + [JsonPropertyName("text")] + public string Text { get; set; } = null!; + } +} diff --git a/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentResponse.cs b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentResponse.cs new file mode 100644 index 000000000000..e57a3b475f5c --- /dev/null +++ b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbedContentResponse.cs @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json.Serialization; + +namespace Microsoft.SemanticKernel.Connectors.Google.Core; + +internal sealed class VertexAIEmbedContentResponse +{ + [JsonPropertyName("embeddings")] + [JsonRequired] + public IList Embeddings { get; set; } = null!; + + internal sealed class ResponseEmbedding + { + [JsonPropertyName("values")] + [JsonRequired] + public ReadOnlyMemory Values { get; set; } + } +} diff --git a/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs index cb59e0087481..2192070f1f0e 100644 --- a/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs +++ b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs @@ -19,6 +19,7 @@ internal sealed class VertexAIEmbeddingClient : ClientBase private readonly string _embeddingModelId; private readonly Uri _embeddingEndpoint; private readonly int? _dimensions; + private readonly bool _useEmbedContentMethod; /// /// Represents a client for interacting with the embeddings models by Vertex AI. @@ -54,10 +55,15 @@ public VertexAIEmbeddingClient( string baseUri = GetVertexAIBaseUri(location); this._embeddingModelId = modelId; - this._embeddingEndpoint = new Uri($"{baseUri}/{versionSubLink}/projects/{projectId}/locations/{location}/publishers/google/models/{this._embeddingModelId}:predict"); + this._useEmbedContentMethod = UsesEmbedContentMethod(modelId); + string embeddingMethod = this._useEmbedContentMethod ? "batchEmbedContents" : "predict"; + this._embeddingEndpoint = new Uri($"{baseUri}/{versionSubLink}/projects/{projectId}/locations/{location}/publishers/google/models/{this._embeddingModelId}:{embeddingMethod}"); this._dimensions = dimensions; } + private static bool UsesEmbedContentMethod(string modelId) + => modelId.StartsWith("gemini-embedding", StringComparison.Ordinal); + /// /// Generates embeddings for the given data asynchronously. /// @@ -72,18 +78,26 @@ public async Task>> GenerateEmbeddingsAsync( { Verify.NotNullOrEmpty(data); - var geminiRequest = this.GetEmbeddingRequest(data, options); - using var httpRequestMessage = await this.CreateHttpRequestAsync(geminiRequest, this._embeddingEndpoint).ConfigureAwait(false); + object request = this._useEmbedContentMethod + ? VertexAIEmbedContentRequest.FromData(data, options?.Dimensions ?? this._dimensions) + : this.GetEmbeddingRequest(data, options); + + using var httpRequestMessage = await this.CreateHttpRequestAsync(request, this._embeddingEndpoint).ConfigureAwait(false); string body = await this.SendRequestAndGetStringBodyAsync(httpRequestMessage, cancellationToken) .ConfigureAwait(false); - return DeserializeAndProcessEmbeddingsResponse(body); + return this._useEmbedContentMethod + ? ProcessEmbedContentResponse(body) + : DeserializeAndProcessEmbeddingsResponse(body); } private VertexAIEmbeddingRequest GetEmbeddingRequest(IEnumerable data, EmbeddingGenerationOptions? options = null) => VertexAIEmbeddingRequest.FromData(data, options?.Dimensions ?? this._dimensions); + private static List> ProcessEmbedContentResponse(string body) + => DeserializeResponse(body).Embeddings.Select(embedding => embedding.Values).ToList(); + private static List> DeserializeAndProcessEmbeddingsResponse(string body) => ProcessEmbeddingsResponse(DeserializeResponse(body));