Skip to content

Commit 17f8368

Browse files
shethaaditAdit Shethrogerbarreto
authored
.Net: Fix Labels Implementation in Vertex AI to Use Dictionary<string, string> (#12636)
### Description This pull request addresses issue #12633 by updating the Labels property in the `GeminiPromptExecutionSettings` and `GeminiRequest` classes to use a `Dictionary<string, string>` instead of a string. This change aligns the implementation with the Vertex AI API requirements, which expect labels to be sent as key-value pairs in JSON format. Fixes #12633 --------- Co-authored-by: Adit Sheth <[email protected]> Co-authored-by: Roger Barreto <[email protected]>
1 parent 2602d33 commit 17f8368

File tree

6 files changed

+94
-38
lines changed

6 files changed

+94
-38
lines changed

dotnet/src/Connectors/Connectors.Google.UnitTests/Core/Gemini/GeminiRequestTests.cs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -528,14 +528,14 @@ public void LabelsFromPromptReturnsAsExpected()
528528
var prompt = "prompt-example";
529529
var executionSettings = new GeminiPromptExecutionSettings
530530
{
531-
Labels = "Key1:Value1"
531+
Labels = new Dictionary<string, string> { { "key1", "value1" }, { "key2", "value2" } }
532532
};
533533

534534
// Act
535535
var request = GeminiRequest.FromPromptAndExecutionSettings(prompt, executionSettings);
536536

537537
// Assert
538-
Assert.NotNull(request.Configuration);
538+
Assert.NotNull(request.Labels);
539539
Assert.Equal(executionSettings.Labels, request.Labels);
540540
}
541541

@@ -569,7 +569,7 @@ public void LabelsFromChatHistoryReturnsAsExpected()
569569
chatHistory.AddUserMessage("user-message2");
570570
var executionSettings = new GeminiPromptExecutionSettings
571571
{
572-
Labels = "Key1:Value1"
572+
Labels = new Dictionary<string, string> { { "key1", "value1" }, { "key2", "value2" } }
573573
};
574574

575575
// Act

dotnet/src/Connectors/Connectors.Google.UnitTests/Services/GoogleAIGeminiChatCompletionServiceTests.cs

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

33
using System;
4+
using System.Collections.Generic;
45
using System.IO;
56
using System.Net.Http;
67
using System.Text;
@@ -69,15 +70,13 @@ public async Task RequestCachedContentWorksCorrectlyAsync(string? cachedContent)
6970
}
7071
}
7172

72-
[Theory]
73-
[InlineData(null)]
74-
[InlineData("key:value")]
75-
[InlineData("")]
76-
public async Task RequestLabelsWorksCorrectlyAsync(string? labels)
73+
[Fact]
74+
public async Task RequestLabelsWorksCorrectlyAsync()
7775
{
7876
// Arrange
7977
string model = "fake-model";
8078
var sut = new GoogleAIGeminiChatCompletionService(model, "key", httpClient: this._httpClient);
79+
var labels = new Dictionary<string, string> { { "key1", "value1" }, { "key2", "value2" } };
8180

8281
// Act
8382
var result = await sut.GetChatMessageContentAsync("my prompt", new GeminiPromptExecutionSettings { Labels = labels });
@@ -87,19 +86,28 @@ public async Task RequestLabelsWorksCorrectlyAsync(string? labels)
8786
Assert.NotNull(this._messageHandlerStub.RequestContent);
8887

8988
var requestBody = UTF8Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
90-
if (labels is not null)
91-
{
92-
Assert.Contains($"\"labels\":\"{labels}\"", requestBody);
93-
}
94-
else
95-
{
96-
// Then no quality is provided, it should not be included in the request body
97-
Assert.DoesNotContain("labels", requestBody);
98-
}
89+
Assert.Contains("\"labels\":{\"key1\":\"value1\",\"key2\":\"value2\"}", requestBody);
90+
}
91+
92+
[Fact]
93+
public async Task RequestLabelsNullWorksCorrectlyAsync()
94+
{
95+
// Arrange
96+
string model = "fake-model";
97+
var sut = new GoogleAIGeminiChatCompletionService(model, "key", httpClient: this._httpClient);
98+
99+
// Act
100+
var result = await sut.GetChatMessageContentAsync("my prompt", new GeminiPromptExecutionSettings { Labels = null });
101+
102+
// Assert
103+
Assert.NotNull(result);
104+
Assert.NotNull(this._messageHandlerStub.RequestContent);
105+
106+
var requestBody = UTF8Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
107+
Assert.DoesNotContain("labels", requestBody);
99108
}
100109

101110
[Theory]
102-
[InlineData(null, false)]
103111
[InlineData(0, true)]
104112
[InlineData(500, true)]
105113
[InlineData(2048, true)]

dotnet/src/Connectors/Connectors.Google.UnitTests/Services/VertexAIGeminiChatCompletionServiceTests.cs

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

33
using System;
4+
using System.Collections.Generic;
45
using System.IO;
56
using System.Net.Http;
67
using System.Text;
@@ -80,15 +81,13 @@ public async Task RequestCachedContentWorksCorrectlyAsync(string? cachedContent)
8081
}
8182
}
8283

83-
[Theory]
84-
[InlineData(null)]
85-
[InlineData("key:value")]
86-
[InlineData("")]
87-
public async Task RequestLabelsWorksCorrectlyAsync(string? labels)
84+
[Fact]
85+
public async Task RequestLabelsWorksCorrectlyAsync()
8886
{
8987
// Arrange
9088
string model = "fake-model";
91-
var sut = new GoogleAIGeminiChatCompletionService(model, "key", httpClient: this._httpClient);
89+
var sut = new VertexAIGeminiChatCompletionService(model, () => new ValueTask<string>("key"), "location", "project", httpClient: this._httpClient);
90+
var labels = new Dictionary<string, string> { { "key1", "value1" }, { "key2", "value2" } };
9291

9392
// Act
9493
var result = await sut.GetChatMessageContentAsync("my prompt", new GeminiPromptExecutionSettings { Labels = labels });
@@ -98,15 +97,25 @@ public async Task RequestLabelsWorksCorrectlyAsync(string? labels)
9897
Assert.NotNull(this._messageHandlerStub.RequestContent);
9998

10099
var requestBody = UTF8Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
101-
if (labels is not null)
102-
{
103-
Assert.Contains($"\"labels\":\"{labels}\"", requestBody);
104-
}
105-
else
106-
{
107-
// Then no quality is provided, it should not be included in the request body
108-
Assert.DoesNotContain("labels", requestBody);
109-
}
100+
Assert.Contains("\"labels\":{\"key1\":\"value1\",\"key2\":\"value2\"}", requestBody);
101+
}
102+
103+
[Fact]
104+
public async Task RequestLabelsNullWorksCorrectlyAsync()
105+
{
106+
// Arrange
107+
string model = "fake-model";
108+
var sut = new VertexAIGeminiChatCompletionService(model, () => new ValueTask<string>("key"), "location", "project", httpClient: this._httpClient);
109+
110+
// Act
111+
var result = await sut.GetChatMessageContentAsync("my prompt", new GeminiPromptExecutionSettings { Labels = null });
112+
113+
// Assert
114+
Assert.NotNull(result);
115+
Assert.NotNull(this._messageHandlerStub.RequestContent);
116+
117+
var requestBody = UTF8Encoding.UTF8.GetString(this._messageHandlerStub.RequestContent);
118+
Assert.DoesNotContain("labels", requestBody);
110119
}
111120

112121
[Theory]

dotnet/src/Connectors/Connectors.Google/Core/Gemini/Models/GeminiRequest.cs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ internal sealed class GeminiRequest
4848

4949
[JsonPropertyName("labels")]
5050
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
51-
public string? Labels { get; set; }
51+
public IDictionary<string, string>? Labels { get; set; }
5252

5353
public void AddFunction(GeminiFunction function)
5454
{
@@ -450,7 +450,12 @@ private static void AddSafetySettings(GeminiPromptExecutionSettings executionSet
450450
private static void AddAdditionalBodyFields(GeminiPromptExecutionSettings executionSettings, GeminiRequest request)
451451
{
452452
request.CachedContent = executionSettings.CachedContent;
453-
request.Labels = executionSettings.Labels;
453+
454+
if (executionSettings.Labels is not null)
455+
{
456+
request.Labels = executionSettings.Labels;
457+
}
458+
454459
if (executionSettings.ThinkingConfig is not null)
455460
{
456461
request.Configuration ??= new ConfigurationElement();

dotnet/src/Connectors/Connectors.Google/GeminiPromptExecutionSettings.cs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public sealed class GeminiPromptExecutionSettings : PromptExecutionSettings
2828
private string? _responseMimeType;
2929
private object? _responseSchema;
3030
private string? _cachedContent;
31-
private string? _labels;
31+
private IDictionary<string, string>? _labels;
3232
private IList<GeminiSafetySetting>? _safetySettings;
3333
private GeminiToolCallBehavior? _toolCallBehavior;
3434
private GeminiThinkingConfig? _thinkingConfig;
@@ -147,9 +147,12 @@ public IList<GeminiSafetySetting>? SafetySettings
147147
/// Gets or sets the labels.
148148
/// </summary>
149149
/// <value>
150-
/// Metadata that can be added to the API call in the format of key-value pairs.
150+
/// The labels with user-defined metadata for the request. It is used for billing and reporting only.
151+
/// label keys and values can be no longer than 63 characters (Unicode codepoints) and can only contain lowercase letters, numeric characters, underscores, and dashes. International characters are allowed. label values are optional. label keys must start with a letter.
151152
/// </value>
152-
public string? Labels
153+
[JsonPropertyName("labels")]
154+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
155+
public IDictionary<string, string>? Labels
153156
{
154157
get => this._labels;
155158
set

dotnet/src/IntegrationTests/Connectors/Google/Gemini/GeminiChatCompletionTests.cs

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

33
using System;
4+
using System.Collections.Generic;
45
using System.IO;
56
using System.Linq;
67
using System.Net.Http;
@@ -600,4 +601,34 @@ public async Task GoogleAIChatReturnsResponseWorksWithThinkingBudgetAsync()
600601
Assert.NotNull(streamResponses[0].Content);
601602
Assert.NotNull(responses[0].Content);
602603
}
604+
605+
[RetryTheory(Skip = "This test is for manual verification.")]
606+
[InlineData(ServiceType.VertexAI)] // GoogleAI does not support labels yet
607+
public async Task GoogleAIChatReturnsResponseWorksWithLabelsAsync(ServiceType serviceType)
608+
{
609+
// Arrange
610+
ChatHistory chatHistory = [];
611+
chatHistory.AddUserMessage("Hello, I'm Brandon, how are you?");
612+
chatHistory.AddAssistantMessage("I'm doing well, thanks for asking.");
613+
chatHistory.AddUserMessage("Call me by my name and expand this abbreviation: LLM");
614+
615+
var sut = this.GetChatService(serviceType);
616+
617+
var settings = new GeminiPromptExecutionSettings
618+
{
619+
Labels = new Dictionary<string, string>()
620+
{
621+
["label1"] = "value1",
622+
["label2"] = "value2"
623+
}
624+
};
625+
626+
// Act
627+
var streamResponses = await sut.GetStreamingChatMessageContentsAsync(chatHistory, settings).ToListAsync();
628+
var responses = await sut.GetChatMessageContentsAsync(chatHistory, settings);
629+
630+
// Assert
631+
Assert.NotNull(streamResponses[0].Content);
632+
Assert.NotNull(responses[0].Content);
633+
}
603634
}

0 commit comments

Comments
 (0)