Skip to content

Commit dbf0726

Browse files
authored
.Net: Fix Google Gemini Enums Schema definition (#11617)
### Motivation and Context - Fix #11395 AIJsonUtilities Schema generation does not generate the `type` for `enum` properties, this workaround became necessary to allow enums to be used with VertexAI/GoogleAI API's. FYI @eiriktsarpalis
1 parent 04dcf86 commit dbf0726

File tree

3 files changed

+98
-11
lines changed

3 files changed

+98
-11
lines changed

dotnet/samples/Concepts/ChatCompletion/Google_GeminiStructuredOutputs.cs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,9 +191,26 @@ private struct Movie
191191

192192
public bool IsAvailableOnStreaming { get; set; }
193193

194+
public MovieGenre? Genre { get; set; }
195+
194196
public List<string> Tags { get; set; }
195197
}
196198

199+
private enum MovieGenre
200+
{
201+
Action,
202+
Adventure,
203+
Comedy,
204+
Drama,
205+
Fantasy,
206+
Horror,
207+
Mystery,
208+
Romance,
209+
SciFi,
210+
Thriller,
211+
Western
212+
}
213+
197214
private sealed class EmailResult
198215
{
199216
public List<Email> Emails { get; set; }
@@ -256,6 +273,7 @@ private void OutputResult(MovieResult movieResult)
256273
Director: {movie.Director}
257274
Release year: {movie.ReleaseYear}
258275
Rating: {movie.Rating}
276+
Genre: {movie.Genre}
259277
Is available on streaming: {movie.IsAvailableOnStreaming}
260278
Tags: {string.Join(",", movie.Tags ?? [])}
261279
""");

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -584,6 +584,60 @@ public void ResponseSchemaConvertsNullableTypesToOpenApiFormat()
584584
Assert.True(ageProperty.GetProperty("nullable").GetBoolean());
585585
}
586586

587+
[Fact]
588+
public void ResponseSchemaAddsTypeToEnumProperties()
589+
{
590+
// Arrange
591+
var prompt = "prompt-example";
592+
var schemaWithEnum = """
593+
{
594+
"properties" : {
595+
"Movies": {
596+
"type" : "array",
597+
"items" : {
598+
"type" : "object",
599+
"properties" : {
600+
"status": {
601+
"enum": ["active", "inactive", null],
602+
"description": "user status"
603+
},
604+
"role": {
605+
"enum": ["admin", "user"],
606+
"description": "user role"
607+
}
608+
}
609+
}
610+
}
611+
}
612+
}
613+
""";
614+
615+
var executionSettings = new GeminiPromptExecutionSettings
616+
{
617+
ResponseMimeType = "application/json",
618+
ResponseSchema = JsonSerializer.Deserialize<JsonElement>(schemaWithEnum)
619+
};
620+
621+
// Act
622+
var request = GeminiRequest.FromPromptAndExecutionSettings(prompt, executionSettings);
623+
624+
// Assert
625+
Assert.NotNull(request.Configuration?.ResponseSchema);
626+
var properties = request.Configuration.ResponseSchema.Value
627+
.GetProperty("properties")
628+
.GetProperty("Movies")
629+
.GetProperty("items")
630+
.GetProperty("properties");
631+
632+
var statusProperty = properties.GetProperty("status");
633+
Assert.Equal("string", statusProperty.GetProperty("type").GetString());
634+
Assert.Equal(3, statusProperty.GetProperty("enum").GetArrayLength());
635+
636+
var roleProperty = properties.GetProperty("role");
637+
Assert.Equal("string", roleProperty.GetProperty("type").GetString());
638+
Assert.Equal(2, roleProperty.GetProperty("enum").GetArrayLength());
639+
}
640+
587641
private sealed class DummyContent(object? innerContent, string? modelId = null, IReadOnlyDictionary<string, object?>? metadata = null) :
588642
KernelContent(innerContent, modelId, metadata);
589643

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

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,7 @@ private static void AddConfiguration(GeminiPromptExecutionSettings executionSett
326326
_ => CreateSchema(responseSchemaSettings.GetType(), GetDefaultOptions())
327327
};
328328

329-
jsonElement = AdjustOpenApi3Nullables(jsonElement);
329+
jsonElement = TransformToOpenApi3Schema(jsonElement);
330330
return jsonElement;
331331
}
332332

@@ -343,37 +343,52 @@ private static void AddConfiguration(GeminiPromptExecutionSettings executionSett
343343
/// - Replaces the type array with a single type value
344344
/// - Adds "nullable": true as a property
345345
/// </remarks>
346-
private static JsonElement AdjustOpenApi3Nullables(JsonElement jsonElement)
346+
private static JsonElement TransformToOpenApi3Schema(JsonElement jsonElement)
347347
{
348348
JsonNode? node = JsonNode.Parse(jsonElement.GetRawText());
349349
if (node is JsonObject rootObject)
350350
{
351-
AdjustOpenApi3Object(rootObject);
351+
TransformOpenApi3Object(rootObject);
352352
}
353353

354354
return JsonSerializer.SerializeToElement(node, GetDefaultOptions());
355355

356-
static void AdjustOpenApi3Object(JsonObject obj)
356+
static void TransformOpenApi3Object(JsonObject obj)
357357
{
358358
if (obj.TryGetPropertyValue("properties", out JsonNode? propsNode) && propsNode is JsonObject properties)
359359
{
360360
foreach (var property in properties)
361361
{
362362
if (property.Value is JsonObject propertyObj)
363363
{
364-
if (propertyObj.TryGetPropertyValue("type", out JsonNode? typeNode) && typeNode is JsonArray typeArray)
364+
// Handle enum properties - add "type": "string" if missing
365+
if (propertyObj.TryGetPropertyValue("enum", out JsonNode? enumNode) && !propertyObj.ContainsKey("type"))
365366
{
366-
var types = typeArray.Select(t => t?.GetValue<string>()).Where(t => t != null).ToList();
367-
if (types.Contains("null"))
367+
propertyObj["type"] = JsonValue.Create("string");
368+
}
369+
else if (propertyObj.TryGetPropertyValue("type", out JsonNode? typeNode))
370+
{
371+
if (typeNode is JsonArray typeArray)
372+
{
373+
var types = typeArray.Select(t => t?.GetValue<string>()).Where(t => t != null).ToList();
374+
if (types.Contains("null"))
375+
{
376+
var mainType = types.First(t => t != "null");
377+
propertyObj["type"] = JsonValue.Create(mainType);
378+
propertyObj["nullable"] = JsonValue.Create(true);
379+
}
380+
}
381+
else if (typeNode is JsonValue typeValue && typeValue.GetValue<string>() == "array")
368382
{
369-
var mainType = types.First(t => t != "null");
370-
propertyObj["type"] = JsonValue.Create(mainType);
371-
propertyObj["nullable"] = JsonValue.Create(true);
383+
if (propertyObj.TryGetPropertyValue("items", out JsonNode? itemsNode) && itemsNode is JsonObject itemsObj)
384+
{
385+
TransformOpenApi3Object(itemsObj);
386+
}
372387
}
373388
}
374389

375390
// Recursively process nested objects
376-
AdjustOpenApi3Object(propertyObj);
391+
TransformOpenApi3Object(propertyObj);
377392
}
378393
}
379394
}

0 commit comments

Comments
 (0)