fix(VertexAI): use :batchEmbedContents for gemini-embedding models

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<string> 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<string> 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<string>()
+        {
+            "Write a story about a magic backpack.",
+            "Print color of backpack."
+        };
+
+        // Act
+        var embeddings = await client.GenerateEmbeddingsAsync(dataToEmbed);
+
+        // Assert
+        VertexAIEmbedContentResponse testDataResponse = JsonSerializer.Deserialize<VertexAIEmbedContentResponse>(
+            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()
     {

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
+      ]
+    }
+  ]
+}

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<EmbedContentRequestItem> Requests { get; set; } = null!;
+
+    public static VertexAIEmbedContentRequest FromData(IEnumerable<string> 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<RequestPart> Parts { get; set; } = null!;
+    }
+
+    internal sealed class RequestPart
+    {
+        [JsonPropertyName("text")]
+        public string Text { get; set; } = null!;
+    }
+}

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<ResponseEmbedding> Embeddings { get; set; } = null!;
+
+    internal sealed class ResponseEmbedding
+    {
+        [JsonPropertyName("values")]
+        [JsonRequired]
+        public ReadOnlyMemory<float> Values { get; set; }
+    }
+}

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;
 
     /// <summary>
     /// 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);
+
     /// <summary>
     /// Generates embeddings for the given data asynchronously.
     /// </summary>
@@ -72,18 +78,26 @@ public async Task<IList<ReadOnlyMemory<float>>> 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<string> data, EmbeddingGenerationOptions? options = null)
         => VertexAIEmbeddingRequest.FromData(data, options?.Dimensions ?? this._dimensions);
 
+    private static List<ReadOnlyMemory<float>> ProcessEmbedContentResponse(string body)
+        => DeserializeResponse<VertexAIEmbedContentResponse>(body).Embeddings.Select(embedding => embedding.Values).ToList();
+
     private static List<ReadOnlyMemory<float>> DeserializeAndProcessEmbeddingsResponse(string body)
         => ProcessEmbeddingsResponse(DeserializeResponse<VertexAIEmbeddingResponse>(body));