From 349875b242fc1ab612fc6c45d50e150b11ce9e69 Mon Sep 17 00:00:00 2001
From: Martin Sosic <sosic.martin@gmail.com>
Date: Tue, 27 Jun 2023 18:24:12 +0200
Subject: [PATCH] Upped number of retries for JSON from 1 to 2. Added pretty
 debugging of chatGPT prompts.

---
 .../src/Wasp/AI/GenerateNewProject/Common.hs  | 42 ++++++++++---------
 waspc/src/Wasp/AI/OpenAI/ChatGPT.hs           | 13 +++++-
 waspc/waspc.cabal                             |  1 +
 3 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/waspc/src/Wasp/AI/GenerateNewProject/Common.hs b/waspc/src/Wasp/AI/GenerateNewProject/Common.hs
index f8f311d4f..1a1ac8352 100644
--- a/waspc/src/Wasp/AI/GenerateNewProject/Common.hs
+++ b/waspc/src/Wasp/AI/GenerateNewProject/Common.hs
@@ -31,30 +31,34 @@ type File = (FilePath, Text)
 data AuthProvider = UsernameAndPassword
 
 queryChatGPTForJSON :: FromJSON a => ChatGPTParams -> [ChatMessage] -> CodeAgent a
-queryChatGPTForJSON chatGPTParams = doQueryForJSON True
+queryChatGPTForJSON chatGPTParams = doQueryForJSON 0
   where
-    doQueryForJSON isFirstTry chatMsgs = do
+    doQueryForJSON :: (FromJSON a) => Int -> [ChatMessage] -> CodeAgent a
+    doQueryForJSON numPrevFailures chatMsgs = do
       response <- queryChatGPT chatGPTParams chatMsgs
       case Aeson.eitherDecode . textToLazyBS . naiveTrimJSON $ response of
         Right result -> return result
         Left errMsg ->
-          if isFirstTry
-            then
-              doQueryForJSON False $
-                chatMsgs
-                  ++ [ GPT.ChatMessage {GPT.role = GPT.Assistant, GPT.content = response},
-                       GPT.ChatMessage
-                         { GPT.role = GPT.User,
-                           GPT.content =
-                             "You did not respond with valid JSON. Please fix it and respond with only"
-                               <> " valid JSON, no other text or explanations. Error I got parsing JSON"
-                               <> " from your last message: "
-                               <> T.pack errMsg
-                         }
-                     ]
-            else do
-              writeToLog "Failed to parse ChatGPT response as JSON."
-              error $ "Failed to parse ChatGPT response as JSON: " <> errMsg
+          let numFailures = numPrevFailures + 1
+           in if numFailures <= maxNumFailuresBeforeGivingUp
+                then
+                  doQueryForJSON (numPrevFailures + 1) $
+                    chatMsgs
+                      ++ [ GPT.ChatMessage {GPT.role = GPT.Assistant, GPT.content = response},
+                           GPT.ChatMessage
+                             { GPT.role = GPT.User,
+                               GPT.content =
+                                 "You did not respond with valid JSON. Please fix it and respond with only"
+                                   <> " valid JSON, no other text or explanations. Error I got parsing JSON"
+                                   <> " from your last message: "
+                                   <> T.pack errMsg
+                             }
+                         ]
+                else do
+                  writeToLog "Failed to parse ChatGPT response as JSON."
+                  error $ "Failed to parse ChatGPT response as JSON: " <> errMsg
+
+    maxNumFailuresBeforeGivingUp = 2
 
 -- TODO: Test more for the optimal temperature (possibly higher).
 defaultChatGPTParams :: ChatGPTParams
diff --git a/waspc/src/Wasp/AI/OpenAI/ChatGPT.hs b/waspc/src/Wasp/AI/OpenAI/ChatGPT.hs
index 38aedfd32..c5158cbd2 100644
--- a/waspc/src/Wasp/AI/OpenAI/ChatGPT.hs
+++ b/waspc/src/Wasp/AI/OpenAI/ChatGPT.hs
@@ -1,5 +1,4 @@
 {-# LANGUAGE DeriveGeneric #-}
-{-# LANGUAGE ViewPatterns #-}
 
 module Wasp.AI.OpenAI.ChatGPT
   ( queryChatGPT,
@@ -14,10 +13,12 @@ module Wasp.AI.OpenAI.ChatGPT
 where
 
 import Control.Arrow ()
+import Control.Monad (when)
 import Data.Aeson ((.=))
 import qualified Data.Aeson as Aeson
 import Data.ByteString.UTF8 as BSU
 import Data.Text (Text)
+import Debug.Pretty.Simple (pTrace)
 import GHC.Generics (Generic)
 import qualified Network.HTTP.Conduit as HTTP.C
 import qualified Network.HTTP.Simple as HTTP
@@ -48,6 +49,16 @@ queryChatGPT apiKey params requestMessages = do
 
   let (chatResponse :: ChatResponse) = HTTP.getResponseBody response
 
+  when False $
+    pTrace
+      ( "\n\n\n\n==================================\n\n"
+          <> show requestMessages
+          <> "\n\n============\n\n"
+          <> show (usage chatResponse)
+          <> "\n\n==================================\n\n\n\n"
+      )
+      $ return ()
+
   return $ content $ message $ head $ choices chatResponse
   where
     secondsToMicroSeconds :: Int -> Int
diff --git a/waspc/waspc.cabal b/waspc/waspc.cabal
index 14dd4dfa6..d726cc063 100644
--- a/waspc/waspc.cabal
+++ b/waspc/waspc.cabal
@@ -119,6 +119,7 @@ library
     , parsec                ^>= 3.1.14
     , path                  ^>= 0.9.2
     , path-io               ^>= 1.6.3
+    , pretty-simple         ^>= 4.1.2
     , regex-tdfa            ^>= 1.3.1
     , strong-path           ^>= 1.1.4
     , unliftio              ^>= 0.2.20