From 4db07a542396ef71237b01513904a36546ce2abd Mon Sep 17 00:00:00 2001
From: Daniel Cazzulino <daniel@cazzulino.com>
Date: Wed, 16 Oct 2024 14:23:50 -0300
Subject: [PATCH] Use low level JSON API to manipulate the wrapper node

Rather than relying on the type system, since a source-generated serializer options would not be able to deal with it.
---
 .../ChatClientStructuredOutputExtensions.cs         | 13 +++++++------
 .../ChatCompletion/ChatCompletion{T}.cs             | 10 ++++------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatClientStructuredOutputExtensions.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatClientStructuredOutputExtensions.cs
index 251dacdd673..ef7dc3d3443 100644
--- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatClientStructuredOutputExtensions.cs
+++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatClientStructuredOutputExtensions.cs
@@ -116,7 +116,7 @@ public static async Task<ChatCompletion<T>> CompleteAsync<T>(
 
         serializerOptions.MakeReadOnly();
 
-        var exporterOptions = new JsonSchemaExporterOptions
+        var schemaNode = (JsonObject)serializerOptions.GetJsonSchemaAsNode(typeof(T), new()
         {
             TreatNullObliviousAsNonNullable = true,
             TransformSchemaNode = static (context, node) =>
@@ -132,9 +132,8 @@ public static async Task<ChatCompletion<T>> CompleteAsync<T>(
 
                 return node;
             },
-        };
+        });
 
-        var schemaNode = (JsonObject)serializerOptions.GetJsonSchemaAsNode(typeof(T), exporterOptions);
         var isObject = schemaNode.TryGetPropertyValue("type", out var schemaType) &&
                 schemaType?.GetValueKind() == JsonValueKind.String &&
                 schemaType.GetValue<string>() is { } type &&
@@ -145,7 +144,11 @@ public static async Task<ChatCompletion<T>> CompleteAsync<T>(
         // We wrap regardless of native structured output, since it also applies to Azure Inference
         if (!isObject)
         {
-            schemaNode = (JsonObject)serializerOptions.GetJsonSchemaAsNode(typeof(Payload<T>), exporterOptions);
+            schemaNode = new JsonObject
+            {
+                { "type", "object" },
+                { "properties", new JsonObject { { "data", schemaNode } } },
+            };
             wrapped = true;
         }
 
@@ -205,6 +208,4 @@ public static async Task<ChatCompletion<T>> CompleteAsync<T>(
             }
         }
     }
-
-    private sealed record Payload<TValue>(TValue Data);
 }
diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatCompletion{T}.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatCompletion{T}.cs
index 7dc3df10263..2fecf70afbf 100644
--- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatCompletion{T}.cs
+++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatCompletion{T}.cs
@@ -83,7 +83,7 @@ public bool TryGetResult([NotNullWhen(true)] out T? result)
 #pragma warning restore CA1031 // Do not catch general exception types
     }
 
-    private static TValue? DeserializeFirstTopLevelObject<TValue>(string json, JsonTypeInfo<TValue> typeInfo)
+    private static T? DeserializeFirstTopLevelObject(string json, JsonTypeInfo<T> typeInfo)
     {
         // We need to deserialize only the first top-level object as a workaround for a common LLM backend
         // issue. GPT 3.5 Turbo commonly returns multiple top-level objects after doing a function call.
@@ -132,10 +132,10 @@ public bool TryGetResult([NotNullWhen(true)] out T? result)
 
         if (wrapped)
         {
-            var result = DeserializeFirstTopLevelObject(json!, (JsonTypeInfo<Payload<T>>)_serializerOptions.GetTypeInfo(typeof(Payload<T>)));
-            if (result != null)
+            var doc = JsonDocument.Parse(json!);
+            if (doc.RootElement.TryGetProperty("data", out var data))
             {
-                deserialized = result.Data;
+                deserialized = DeserializeFirstTopLevelObject(data.GetRawText(), (JsonTypeInfo<T>)_serializerOptions.GetTypeInfo(typeof(T)));
             }
         }
         else
@@ -155,8 +155,6 @@ public bool TryGetResult([NotNullWhen(true)] out T? result)
         return deserialized;
     }
 
-    private sealed record Payload<TValue>(TValue Data);
-
     private enum FailureReason
     {
         ResultDidNotContainJson,