.Net: Gemini Task -> ValueTask for BearerTokenProvider (microsoft#6708)

RogerBarreto · web-flow · commit 82192cfeabd1 · 2024-06-17T21:09:21.000Z
### Motivation and Context Resolve microsoft#6679
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiChatGenerationTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiChatGenerationTests.cs
@@ -416,6 +416,57 @@ public async Task ItCreatesPostRequestWithSemanticKernelVersionHeaderAsync()
         Assert.Equal(expectedVersion, header);
     }
 
+    [Fact]
+    public async Task ItCanUseValueTasksSequentiallyForBearerTokenAsync()
+    {
+        // Arrange
+        var bearerTokenGenerator = new BearerTokenGenerator()
+        {
+            BearerKeys = new List<string> { "key1", "key2", "key3" }
+        };
+
+        var responseContent = File.ReadAllText(ChatTestDataFilePath);
+        using var content1 = new HttpResponseMessage { Content = new StringContent(responseContent) };
+        using var content2 = new HttpResponseMessage { Content = new StringContent(responseContent) };
+
+        using MultipleHttpMessageHandlerStub multipleMessageHandlerStub = new()
+        {
+            ResponsesToReturn = [content1, content2]
+        };
+        using var httpClient = new HttpClient(multipleMessageHandlerStub, false);
+
+        var client = new GeminiChatCompletionClient(
+                httpClient: httpClient,
+                modelId: "fake-model",
+                apiVersion: VertexAIVersion.V1,
+                bearerTokenProvider: () => bearerTokenGenerator.GetBearerToken(),
+                location: "fake-location",
+                projectId: "fake-project-id");
+
+        var chatHistory = CreateSampleChatHistory();
+
+        // Act
+        await client.GenerateChatMessageAsync(chatHistory);
+        await client.GenerateChatMessageAsync(chatHistory);
+        var firstRequestHeader = multipleMessageHandlerStub.RequestHeaders[0]?.GetValues("Authorization").SingleOrDefault();
+        var secondRequestHeader = multipleMessageHandlerStub.RequestHeaders[1]?.GetValues("Authorization").SingleOrDefault();
+
+        // Assert
+        Assert.NotNull(firstRequestHeader);
+        Assert.NotNull(secondRequestHeader);
+        Assert.NotEqual(firstRequestHeader, secondRequestHeader);
+        Assert.Equal("Bearer key1", firstRequestHeader);
+        Assert.Equal("Bearer key2", secondRequestHeader);
+    }
+
+    private sealed class BearerTokenGenerator()
+    {
+        private int _index = 0;
+        public required List<string> BearerKeys { get; init; }
+
+        public ValueTask<string> GetBearerToken() => ValueTask.FromResult(this.BearerKeys[this._index++]);
+    }
+
     private static ChatHistory CreateSampleChatHistory()
     {
         var chatHistory = new ChatHistory();
@@ -436,7 +487,7 @@ private GeminiChatCompletionClient CreateChatCompletionClient(
                 httpClient: httpClient ?? this._httpClient,
                 modelId: modelId,
                 apiVersion: VertexAIVersion.V1,
-                bearerTokenProvider: () => Task.FromResult(bearerKey),
+                bearerTokenProvider: () => new ValueTask<string>(bearerKey),
                 location: "fake-location",
                 projectId: "fake-project-id");
         }
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiChatStreamingTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiChatStreamingTests.cs
@@ -368,7 +368,7 @@ private GeminiChatCompletionClient CreateChatCompletionClient(
             return new GeminiChatCompletionClient(
                 httpClient: httpClient ?? this._httpClient,
                 modelId: modelId,
-                bearerTokenProvider: () => Task.FromResult(bearerKey),
+                bearerTokenProvider: () => new ValueTask<string>(bearerKey),
                 apiVersion: VertexAIVersion.V1,
                 location: "fake-location",
                 projectId: "fake-project-id");
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiCountingTokensTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/Clients/GeminiCountingTokensTests.cs
@@ -124,7 +124,7 @@ private GeminiTokenCounterClient CreateTokenCounterClient(
             return new GeminiTokenCounterClient(
                 httpClient: this._httpClient,
                 modelId: modelId,
-                bearerTokenProvider: () => Task.FromResult(bearerKey),
+                bearerTokenProvider: () => ValueTask.FromResult(bearerKey),
                 apiVersion: VertexAIVersion.V1,
                 location: "fake-location",
                 projectId: "fake-project-id");
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Core/VertexAI/VertexAIClientEmbeddingsGenerationTests.cs
@@ -143,7 +143,7 @@ private VertexAIEmbeddingClient CreateEmbeddingsClient(
         var client = new VertexAIEmbeddingClient(
             httpClient: this._httpClient,
             modelId: modelId,
-            bearerTokenProvider: () => Task.FromResult(bearerKey ?? "fake-key"),
+            bearerTokenProvider: () => ValueTask.FromResult(bearerKey ?? "fake-key"),
             apiVersion: VertexAIVersion.V1,
             location: "us-central1",
             projectId: "fake-project-id");
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Extensions/VertexAIMemoryBuilderExtensionsTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Extensions/VertexAIMemoryBuilderExtensionsTests.cs
@@ -39,7 +39,7 @@ public void ShouldBuildMemoryWithVertexAIEmbeddingGeneratorBearerAsFunc()
 
         // Act
         var memory = builder
-            .WithVertexAITextEmbeddingGeneration("fake-model", () => Task.FromResult("fake-bearer-key"), "fake-location", "fake-project")
+            .WithVertexAITextEmbeddingGeneration("fake-model", () => ValueTask.FromResult("fake-bearer-key"), "fake-location", "fake-project")
             .WithMemoryStore(this._mockMemoryStore.Object)
             .Build();
 
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Extensions/VertexAIServiceCollectionExtensionsTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Extensions/VertexAIServiceCollectionExtensionsTests.cs
@@ -38,7 +38,7 @@ public void VertexAIGeminiChatCompletionServiceShouldBeRegisteredInKernelService
         var kernelBuilder = Kernel.CreateBuilder();
 
         // Act
-        kernelBuilder.AddVertexAIGeminiChatCompletion("modelId", () => Task.FromResult("apiKey"), location: "test2", projectId: "projectId");
+        kernelBuilder.AddVertexAIGeminiChatCompletion("modelId", () => ValueTask.FromResult("apiKey"), location: "test2", projectId: "projectId");
         var kernel = kernelBuilder.Build();
 
         // Assert
@@ -70,7 +70,7 @@ public void VertexAIGeminiChatCompletionServiceShouldBeRegisteredInServiceCollec
         var services = new ServiceCollection();
 
         // Act
-        services.AddVertexAIGeminiChatCompletion("modelId", () => Task.FromResult("apiKey"), location: "test2", projectId: "projectId");
+        services.AddVertexAIGeminiChatCompletion("modelId", () => ValueTask.FromResult("apiKey"), location: "test2", projectId: "projectId");
         var serviceProvider = services.BuildServiceProvider();
 
         // Assert
@@ -102,7 +102,7 @@ public void VertexAIEmbeddingGenerationServiceShouldBeRegisteredInKernelServices
         var kernelBuilder = Kernel.CreateBuilder();
 
         // Act
-        kernelBuilder.AddVertexAIEmbeddingGeneration("modelId", () => Task.FromResult("apiKey"), location: "test2", projectId: "projectId");
+        kernelBuilder.AddVertexAIEmbeddingGeneration("modelId", () => ValueTask.FromResult("apiKey"), location: "test2", projectId: "projectId");
         var kernel = kernelBuilder.Build();
 
         // Assert
@@ -134,7 +134,7 @@ public void VertexAIEmbeddingGenerationServiceShouldBeRegisteredInServiceCollect
         var services = new ServiceCollection();
 
         // Act
-        services.AddVertexAIEmbeddingGeneration("modelId", () => Task.FromResult("apiKey"), location: "test2", projectId: "projectId");
+        services.AddVertexAIEmbeddingGeneration("modelId", () => ValueTask.FromResult("apiKey"), location: "test2", projectId: "projectId");
         var serviceProvider = services.BuildServiceProvider();
 
         // Assert
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAIGeminiChatCompletionServiceTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAIGeminiChatCompletionServiceTests.cs
@@ -25,7 +25,7 @@ public void AttributesShouldContainModelIdBearerAsFunc()
     {
         // Arrange & Act
         string model = "fake-model";
-        var service = new VertexAIGeminiChatCompletionService(model, () => Task.FromResult("key"), "location", "project");
+        var service = new VertexAIGeminiChatCompletionService(model, () => new ValueTask<string>("key"), "location", "project");
 
         // Assert
         Assert.Equal(model, service.Attributes[AIServiceExtensions.ModelIdKey]);
diff --git a/dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAITextEmbeddingGenerationServiceTests.cs b/dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAITextEmbeddingGenerationServiceTests.cs
@@ -25,7 +25,7 @@ public void AttributesShouldContainModelIdBearerAsFunc()
     {
         // Arrange & Act
         string model = "fake-model";
-        var service = new VertexAITextEmbeddingGenerationService(model, () => Task.FromResult("key"), "location", "project");
+        var service = new VertexAITextEmbeddingGenerationService(model, () => ValueTask.FromResult("key"), "location", "project");
 
         // Assert
         Assert.Equal(model, service.Attributes[AIServiceExtensions.ModelIdKey]);
diff --git a/dotnet/src/Connectors/Connectors.Google/Connectors.Google.csproj b/dotnet/src/Connectors/Connectors.Google/Connectors.Google.csproj
@@ -10,8 +10,8 @@
   </PropertyGroup>
 
   <!-- IMPORT NUGET PACKAGE SHARED PROPERTIES -->
-  <Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props"/>
-  <Import Project="$(RepoRoot)/dotnet/src/InternalUtilities/src/InternalUtilities.props"/>
+  <Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props" />
+  <Import Project="$(RepoRoot)/dotnet/src/InternalUtilities/src/InternalUtilities.props" />
 
   <PropertyGroup>
     <!-- NuGet Package Settings -->
@@ -21,12 +21,12 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\SemanticKernel.Abstractions\SemanticKernel.Abstractions.csproj" />
-    <ProjectReference Include="..\..\SemanticKernel.Core\SemanticKernel.Core.csproj"/>
+    <ProjectReference Include="..\..\SemanticKernel.Core\SemanticKernel.Core.csproj" />
   </ItemGroup>
 
   <ItemGroup>
-    <InternalsVisibleTo Include="SemanticKernel.Connectors.GoogleVertexAI.UnitTests"/>
-    <InternalsVisibleTo Include="DynamicProxyGenAssembly2"/>
+    <InternalsVisibleTo Include="SemanticKernel.Connectors.GoogleVertexAI.UnitTests" />
+    <InternalsVisibleTo Include="DynamicProxyGenAssembly2" />
   </ItemGroup>
 
 </Project>
diff --git a/dotnet/src/Connectors/Connectors.Google/Core/ClientBase.cs b/dotnet/src/Connectors/Connectors.Google/Core/ClientBase.cs
@@ -14,7 +14,7 @@ namespace Microsoft.SemanticKernel.Connectors.Google.Core;
 
 internal abstract class ClientBase
 {
-    private readonly Func<Task<string>>? _bearerTokenProvider;
+    private readonly Func<ValueTask<string>>? _bearerTokenProvider;
 
     protected ILogger Logger { get; }
 
@@ -23,7 +23,7 @@ internal abstract class ClientBase
     protected ClientBase(
         HttpClient httpClient,
         ILogger? logger,
-        Func<Task<string>> bearerTokenProvider)
+        Func<ValueTask<string>> bearerTokenProvider)
         : this(httpClient, logger)
     {
         Verify.NotNull(bearerTokenProvider);
diff --git a/dotnet/src/Connectors/Connectors.Google/Core/Gemini/Clients/GeminiChatCompletionClient.cs b/dotnet/src/Connectors/Connectors.Google/Core/Gemini/Clients/GeminiChatCompletionClient.cs
@@ -125,7 +125,7 @@ public GeminiChatCompletionClient(
     public GeminiChatCompletionClient(
         HttpClient httpClient,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion,
diff --git a/dotnet/src/Connectors/Connectors.Google/Core/Gemini/Clients/GeminiTokenCounterClient.cs b/dotnet/src/Connectors/Connectors.Google/Core/Gemini/Clients/GeminiTokenCounterClient.cs
@@ -57,7 +57,7 @@ public GeminiTokenCounterClient(
     public GeminiTokenCounterClient(
         HttpClient httpClient,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion,
diff --git a/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs b/dotnet/src/Connectors/Connectors.Google/Core/VertexAI/VertexAIEmbeddingClient.cs
@@ -31,7 +31,7 @@ internal sealed class VertexAIEmbeddingClient : ClientBase
     public VertexAIEmbeddingClient(
         HttpClient httpClient,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion,
diff --git a/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIKernelBuilderExtensions.cs b/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIKernelBuilderExtensions.cs
@@ -37,7 +37,7 @@ public static class VertexAIKernelBuilderExtensions
     public static IKernelBuilder AddVertexAIGeminiChatCompletion(
         this IKernelBuilder builder,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
@@ -122,7 +122,7 @@ public static IKernelBuilder AddVertexAIGeminiChatCompletion(
     public static IKernelBuilder AddVertexAIEmbeddingGeneration(
         this IKernelBuilder builder,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
diff --git a/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIMemoryBuilderExtensions.cs b/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIMemoryBuilderExtensions.cs
@@ -33,7 +33,7 @@ public static class VertexAIMemoryBuilderExtensions
     public static MemoryBuilder WithVertexAITextEmbeddingGeneration(
         this MemoryBuilder builder,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
diff --git a/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIServiceCollectionExtensions.cs b/dotnet/src/Connectors/Connectors.Google/Extensions/VertexAIServiceCollectionExtensions.cs
@@ -35,7 +35,7 @@ public static class VertexAIServiceCollectionExtensions
     public static IServiceCollection AddVertexAIGeminiChatCompletion(
         this IServiceCollection services,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
@@ -116,7 +116,7 @@ public static IServiceCollection AddVertexAIGeminiChatCompletion(
     public static IServiceCollection AddVertexAIEmbeddingGeneration(
         this IServiceCollection services,
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
diff --git a/dotnet/src/Connectors/Connectors.Google/Services/VertexAIGeminiChatCompletionService.cs b/dotnet/src/Connectors/Connectors.Google/Services/VertexAIGeminiChatCompletionService.cs
@@ -39,7 +39,7 @@ public VertexAIGeminiChatCompletionService(
         VertexAIVersion apiVersion = VertexAIVersion.V1,
         HttpClient? httpClient = null,
         ILoggerFactory? loggerFactory = null)
-        : this(modelId, () => Task.FromResult(bearerKey), location, projectId, apiVersion, httpClient, loggerFactory)
+        : this(modelId, () => new ValueTask<string>(bearerKey), location, projectId, apiVersion, httpClient, loggerFactory)
     {
         Verify.NotNullOrWhiteSpace(bearerKey);
     }
@@ -61,7 +61,7 @@ public VertexAIGeminiChatCompletionService(
     /// </remarks>
     public VertexAIGeminiChatCompletionService(
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,
diff --git a/dotnet/src/Connectors/Connectors.Google/Services/VertexAITextEmbeddingGenerationService.cs b/dotnet/src/Connectors/Connectors.Google/Services/VertexAITextEmbeddingGenerationService.cs
@@ -39,7 +39,7 @@ public VertexAITextEmbeddingGenerationService(
         VertexAIVersion apiVersion = VertexAIVersion.V1,
         HttpClient? httpClient = null,
         ILoggerFactory? loggerFactory = null)
-        : this(modelId, () => Task.FromResult(bearerKey), location, projectId, apiVersion, httpClient, loggerFactory)
+        : this(modelId, () => new ValueTask<string>(bearerKey), location, projectId, apiVersion, httpClient, loggerFactory)
     {
         Verify.NotNullOrWhiteSpace(bearerKey);
     }
@@ -61,7 +61,7 @@ public VertexAITextEmbeddingGenerationService(
     /// </remarks>
     public VertexAITextEmbeddingGenerationService(
         string modelId,
-        Func<Task<string>> bearerTokenProvider,
+        Func<ValueTask<string>> bearerTokenProvider,
         string location,
         string projectId,
         VertexAIVersion apiVersion = VertexAIVersion.V1,

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ public void AttributesShouldContainModelIdBearerAsFunc()`
`25`	`25`	`{`
`26`	`26`	`// Arrange & Act`
`27`	`27`	`string model = "fake-model";`
`28`		`- var service = new VertexAIGeminiChatCompletionService(model, () => Task.FromResult("key"), "location", "project");`
	`28`	`+ var service = new VertexAIGeminiChatCompletionService(model, () => new ValueTask<string>("key"), "location", "project");`
`29`	`29`
`30`	`30`	`// Assert`
`31`	`31`	`Assert.Equal(model, service.Attributes[AIServiceExtensions.ModelIdKey]);`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ namespace Microsoft.SemanticKernel.Connectors.Google.Core;`
`14`	`14`
`15`	`15`	`internal abstract class ClientBase`
`16`	`16`	`{`
`17`		`- private readonly Func<Task<string>>? _bearerTokenProvider;`
	`17`	`+ private readonly Func<ValueTask<string>>? _bearerTokenProvider;`
`18`	`18`
`19`	`19`	`protected ILogger Logger { get; }`
`20`	`20`
`@@ -23,7 +23,7 @@ internal abstract class ClientBase`
`23`	`23`	`protected ClientBase(`
`24`	`24`	`HttpClient httpClient,`
`25`	`25`	`ILogger? logger,`
`26`		`- Func<Task<string>> bearerTokenProvider)`
	`26`	`+ Func<ValueTask<string>> bearerTokenProvider)`
`27`	`27`	`: this(httpClient, logger)`
`28`	`28`	`{`
`29`	`29`	`Verify.NotNull(bearerTokenProvider);`