From 7ecb0f80356101059cc88acf10ec028f11a0ce96 Mon Sep 17 00:00:00 2001
From: Rodion Mostovoi <36400912+rodion-m@users.noreply.github.com>
Date: Sat, 27 Jul 2024 21:22:03 +0500
Subject: [PATCH] Handle extra text in JSON responses and improve error
 handling

Enhanced the handling of JSON responses to account for extraneous text often returned by weaker LLMs. Consolidated error checking and improved error messages for deserialization failures, ensuring any issues are clearly communicated. Added more detailed checks for the presence of JSON brackets in the response.
---
 ...iClientExtensions.GetStructuredResponse.cs | 39 ++++++++++++-------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/modules/OpenAI.ChatGpt.Modules.StructuredResponse/OpenAiClientExtensions.GetStructuredResponse.cs b/src/modules/OpenAI.ChatGpt.Modules.StructuredResponse/OpenAiClientExtensions.GetStructuredResponse.cs
index 8097b4b..14d5e30 100644
--- a/src/modules/OpenAI.ChatGpt.Modules.StructuredResponse/OpenAiClientExtensions.GetStructuredResponse.cs
+++ b/src/modules/OpenAI.ChatGpt.Modules.StructuredResponse/OpenAiClientExtensions.GetStructuredResponse.cs
@@ -154,8 +154,11 @@ private static TObject DeserializeOrThrow<TObject>(JsonSerializerOptions? jsonDe
             }
         }
 
-        if(!response.StartsWith('{') || !response.EndsWith('}'))
+        if (!response.StartsWith('{') || !response.EndsWith('}'))
         {
+            // Sometimes, weak LLMs return responses with extra text,
+            // like `Output: {"key": "value"}`.
+            // Try to handle such cases.
             var (openBracketIndex, closeBracketIndex) = FindFirstAndLastBracket(response);
             response = response[openBracketIndex..(closeBracketIndex + 1)];
         }
@@ -169,16 +172,17 @@ private static TObject DeserializeOrThrow<TObject>(JsonSerializerOptions? jsonDe
         try
         {
             deserialized = JsonSerializer.Deserialize<TObject>(response, jsonDeserializerOptions);
-            if (deserialized is null)
-            {
-                throw new InvalidJsonException(
-                    $"Failed to deserialize response to {typeof(TObject)}. Response: {response}.", response);
-            }
         }
-        catch (JsonException jsonException)
+        catch (JsonException exception)
+        {
+            throw new InvalidJsonException(
+                $"Failed to deserialize response to {typeof(TObject)}. Response: {response}.", response, exception);
+        }
+
+        if (deserialized is null)
         {
             throw new InvalidJsonException(
-                $"Failed to deserialize response to {typeof(TObject)}. Response: {response}.", response, jsonException);
+                $"Failed to deserialize response to {typeof(TObject)}. Response: {response}.", response);
         }
 
         return deserialized;
@@ -186,13 +190,22 @@ private static TObject DeserializeOrThrow<TObject>(JsonSerializerOptions? jsonDe
         static (int openBracketIndex, int closeBracketIndex) FindFirstAndLastBracket(string response)
         {
             ArgumentNullException.ThrowIfNull(response);
-            var openBracketIndex = response.IndexOf('{');
-            var closeBracketIndex = response.LastIndexOf('}');
-            if (openBracketIndex < 0 || closeBracketIndex < 0)
+            int openBracketIndex = response.IndexOf('{');
+            if (openBracketIndex < 0)
             {
-                throw new InvalidJsonException(
-                    $"Failed to deserialize response to {typeof(TObject)}. Response: {response}.", response);
+                string message = $"Failed to deserialize response to {typeof(TObject)}" +
+                                 $", opening bracket not found. Response: {response}.";
+                throw new InvalidJsonException(message, response);
             }
+
+            int closeBracketIndex = response.LastIndexOf('}');
+            if (closeBracketIndex < 0)
+            {
+                string message = $"Failed to deserialize response to {typeof(TObject)}" +
+                                 $", closing bracket not found. Response: {response}.";
+                throw new InvalidJsonException(message, response);
+            }
+
             return (openBracketIndex, closeBracketIndex);
         }
     }