Wasp AI: got prisma format working, fixed recovery on http errors

* Put some skeleton logic in.

* Fixed wrong handling of http errors, got prisma format working.

* fix

waspc/packages/prisma/README.md

@@ -8,7 +8,7 @@ Run `npm install`.
 Run `npm run build` to compile TS to JS. Do this after any changes to TS files, or if you have never run it before.
 
 ### format
-Run `npm run format` and pass prisma schema source (so PSL) via stdin.
+Run `npm start format` and pass prisma schema source (so PSL) via stdin.
 It can even be an incomplete schema -> e.g. just model declarations.
 
 This will run `prisma format` on it and return formatted schema (PSL) + any warnings/errors.

waspc/packages/prisma/package.json

@@ -6,8 +6,8 @@
   "type": "module",
   "scripts": {
     "build": "npx tsc",
-    "start": "npm run format",
-    "format": "node ./dist/format-cli.js",
+    "start": "node ./dist/index.js",
+    "format": "npm start format",
     "test": "jest"
   },
   "dependencies": {

waspc/packages/prisma/src/index.ts

@@ -1,9 +1,17 @@
+import process from 'process';
 import { formatSchema } from './format.js';
 
 async function main() {
-  const schemaPsl = await readWholeStdin();
-  const { formattedSchemaPsl, stderr, exitCode } = await formatSchema(schemaPsl);
-  console.log(JSON.stringify({ formattedSchemaPsl, errors: exitCode === 0 ? undefined : stderr }));
+  const args = process.argv.slice(2);
+  switch (args[0]) {
+    case 'format':
+      const schemaPsl = await readWholeStdin();
+      const { formattedSchemaPsl, stderr, exitCode } = await formatSchema(schemaPsl);
+      console.log(JSON.stringify({ formattedSchemaPsl, errors: exitCode === 0 ? undefined : stderr }));
+      break;
+    default:
+      console.log("Specify a command!");
+  }
 }
 
 async function readWholeStdin(): Promise<string> {
@@ -18,4 +26,4 @@ async function readWholeStdin(): Promise<string> {
   });
 }
 
-main().catch((err) => { console.error(err); process.exit(1); });
+main().catch((err) => { console.error(err); process.exit(1); });

waspc/src/Wasp/AI/CodeAgent.hs

@@ -24,8 +24,9 @@ import Control.Monad.State (MonadState, StateT (runStateT), gets, modify)
 import qualified Data.HashMap.Strict as H
 import Data.Text (Text)
 import qualified Data.Text as T
+import qualified Network.HTTP.Simple as HTTP
 import System.IO (hPutStrLn, stderr)
-import UnliftIO (catch, throwIO)
+import UnliftIO (Handler (Handler), catches, throwIO)
 import Wasp.AI.OpenAI (OpenAIApiKey)
 import Wasp.AI.OpenAI.ChatGPT (ChatGPTParams (..), ChatMessage, ChatResponse)
 import qualified Wasp.AI.OpenAI.ChatGPT as ChatGPT
@@ -33,6 +34,7 @@ import qualified Wasp.Util as Util
 import Wasp.Util.IO.Retry (MonadRetry)
 import qualified Wasp.Util.IO.Retry as R
 import Wasp.Util.Network.HTTP (catchRetryableHttpException)
+import qualified Wasp.Util.Network.HTTP as Utils.HTTP
 
 newtype CodeAgent a = CodeAgent {_unCodeAgent :: ReaderT CodeAgentConfig (StateT CodeAgentState IO) a}
   deriving (Monad, Applicative, Functor, MonadIO, MonadReader CodeAgentConfig, MonadState CodeAgentState)
@@ -50,11 +52,21 @@ instance MonadRetry CodeAgent where
 runCodeAgent :: CodeAgentConfig -> CodeAgent a -> IO a
 runCodeAgent config codeAgent =
   (fst <$> (_unCodeAgent codeAgent `runReaderT` config) `runStateT` initialState)
-    `catch` ( \(e :: SomeException) -> do
-                _writeLog config $
-                  "Code agent failed with the following error: " <> T.pack (shortenWithEllipsisTo 30 $ displayException e)
-                throwIO e
-            )
+    `catches` [ Handler
+                  ( \(e :: HTTP.HttpException) -> do
+                      let errorInfo =
+                            maybe (showShortException e) show $ Utils.HTTP.getHttpExceptionStatusCode e
+                          logMsg = T.pack $ "Code agent failed with the http error: " <> errorInfo
+                      _writeLog config logMsg
+                      throwIO e
+                  ),
+                Handler
+                  ( \(e :: SomeException) -> do
+                      _writeLog config $
+                        "Code agent failed with the following error: " <> T.pack (showShortException e)
+                      throwIO e
+                  )
+              ]
   where
     initialState =
       CodeAgentState
@@ -65,6 +77,9 @@ runCodeAgent config codeAgent =
 
     shortenWithEllipsisTo maxLen text = if length text <= maxLen then text else (take maxLen text) <> "..."
 
+    showShortException :: forall e. Exception e => e -> String
+    showShortException = shortenWithEllipsisTo 30 . displayException
+
 writeToLog :: Text -> CodeAgent ()
 writeToLog msg = asks _writeLog >>= \f -> liftIO $ f msg
 

waspc/src/Wasp/AI/GenerateNewProject.hs

@@ -7,7 +7,6 @@ import Control.Monad (forM, forM_)
 import Data.List (nub)
 import Data.Text (Text)
 import qualified Data.Text as T
-import NeatInterpolation (trimming)
 import StrongPath (File', Path, Rel, System)
 import Text.Printf (printf)
 import Wasp.AI.CodeAgent (CodeAgent, getTotalTokensUsage, writeToLog)
@@ -39,9 +38,7 @@ generateNewProject newProjectDetails waspProjectSkeletonFiles = do
     generateAndWriteProjectSkeletonAndPresetFiles newProjectDetails waspProjectSkeletonFiles
   writeToLog "Generated project skeleton."
 
-  writeToLog "Generating plan..."
   plan <- generatePlan newProjectDetails planRules
-  writeToLog $ "Plan generated!\n" <> summarizePlan plan
 
   writeEntitiesToWaspFile waspFilePath (Plan.entities plan)
   writeToLog "Updated wasp file with entities."
@@ -88,21 +85,3 @@ generateNewProject newProjectDetails waspProjectSkeletonFiles = do
 
   writeToLog "Done!"
   where
-    summarizePlan plan =
-      let numQueries = showT $ length $ Plan.queries plan
-          numActions = showT $ length $ Plan.actions plan
-          numPages = showT $ length $ Plan.pages plan
-          numEntities = showT $ length $ Plan.entities plan
-          queryNames = showT $ Plan.opName <$> Plan.queries plan
-          actionNames = showT $ Plan.opName <$> Plan.actions plan
-          pageNames = showT $ Plan.pageName <$> Plan.pages plan
-          entityNames = showT $ Plan.entityName <$> Plan.entities plan
-       in [trimming|
-            - ${numQueries} queries: ${queryNames}
-            - ${numActions} actions: ${actionNames}
-            - ${numEntities} entities: ${entityNames}
-            - ${numPages} pages: ${pageNames}
-          |]
-
-    showT :: Show a => a -> Text
-    showT = T.pack . show

waspc/src/Wasp/AI/GenerateNewProject/Plan.hs

@@ -10,14 +10,18 @@ module Wasp.AI.GenerateNewProject.Plan
   )
 where
 
+import Control.Monad.IO.Class (liftIO)
 import Data.Aeson (FromJSON)
 import Data.Aeson.Types (ToJSON)
+import Data.Char (isSpace)
 import Data.List (find, intercalate, isPrefixOf)
+import Data.Maybe (isNothing)
+import Data.Text (Text)
 import qualified Data.Text as T
 import GHC.Generics (Generic)
 import NeatInterpolation (trimming)
 import qualified Text.Parsec as Parsec
-import Wasp.AI.CodeAgent (CodeAgent)
+import Wasp.AI.CodeAgent (CodeAgent, writeToLog)
 import Wasp.AI.GenerateNewProject.Common
   ( NewProjectDetails (..),
     defaultChatGPTParams,
@@ -27,6 +31,7 @@ import Wasp.AI.GenerateNewProject.Common
 import Wasp.AI.GenerateNewProject.Common.Prompts (appDescriptionBlock)
 import qualified Wasp.AI.GenerateNewProject.Common.Prompts as Prompts
 import Wasp.AI.OpenAI.ChatGPT (ChatGPTParams (_model), ChatMessage (..), ChatRole (..), Model (GPT_4))
+import qualified Wasp.Psl.Format as Prisma
 import qualified Wasp.Psl.Parser.Model as Psl.Parser
 import qualified Wasp.Util.Aeson as Util.Aeson
 
@@ -35,8 +40,13 @@ type PlanRule = String
 
 generatePlan :: NewProjectDetails -> [PlanRule] -> CodeAgent Plan
 generatePlan newProjectDetails planRules = do
-  queryChatGPTForJSON (defaultChatGPTParams {_model = planGptModel}) chatMessages
-    >>= fixPlanIfNeeded
+  writeToLog "Generating plan... (slowest step, usually takes 30 to 60 seconds)"
+  initialPlan <- queryChatGPTForJSON (defaultChatGPTParams {_model = planGptModel}) chatMessages
+  writeToLog $ "Initial plan generated!\n" <> summarizePlan initialPlan
+  writeToLog "Fixing initial plan..."
+  fixedPlan <- fixPlan initialPlan
+  writeToLog $ "Plan fixed!\n" <> summarizePlan initialPlan
+  return fixedPlan
   where
     chatMessages =
       [ ChatMessage {role = System, content = Prompts.systemPrompt},
@@ -105,28 +115,39 @@ generatePlan newProjectDetails planRules = do
         ${appDescriptionBlockText}
       |]
 
-    fixPlanIfNeeded :: Plan -> CodeAgent Plan
-    fixPlanIfNeeded plan = do
+    fixPlan :: Plan -> CodeAgent Plan
+    fixPlan initialPlan = do
+      (maybePrismaFormatErrorsMsg, formattedEntities) <- liftIO $ prismaFormat $ entities initialPlan
+      let plan' = initialPlan {entities = formattedEntities}
       let issues =
-            checkPlanForEntityIssues plan
-              <> checkPlanForOperationIssues Query plan
-              <> checkPlanForOperationIssues Action plan
-              <> checkPlanForPageIssues plan
-      if null issues
-        then return plan
+            checkPlanForEntityIssues plan'
+              <> checkPlanForOperationIssues Query plan'
+              <> checkPlanForOperationIssues Action plan'
+              <> checkPlanForPageIssues plan'
+      if null issues && isNothing maybePrismaFormatErrorsMsg
+        then return plan'
         else do
-          let issuesText = T.pack $ intercalate "\n" ((" - " <>) <$> issues)
+          let issuesText =
+                if null issues
+                  then ""
+                  else
+                    "I found following potential issues with the plan that you made:\n"
+                      <> T.pack (intercalate "\n" ((" - " <>) <$> issues))
+          let prismaFormatErrorsText =
+                case maybePrismaFormatErrorsMsg of
+                  Nothing -> ""
+                  Just msg -> "Following errors were reported by the 'prisma format' command:\n" <> msg
           queryChatGPTForJSON (defaultChatGPTParamsForFixing {_model = planGptModel}) $
             chatMessages
-              <> [ ChatMessage {role = Assistant, content = Util.Aeson.encodeToText plan},
+              <> [ ChatMessage {role = Assistant, content = Util.Aeson.encodeToText plan'},
                    ChatMessage
                      { role = User,
                        content =
                          [trimming|
-                           I found following potential issues with the plan that you made:
-
                            ${issuesText}
 
+                           ${prismaFormatErrorsText}
+
                            Please improve the plan with regard to these issues and any other potential issues that you find.
 
                            Respond ONLY with a valid JSON that is a plan.
@@ -168,6 +189,34 @@ checkPlanForEntityIssues plan =
 
     parsePslBody = Parsec.parse Psl.Parser.body ""
 
+-- | Calls "prisma format" on given entities, and returns formatted/fixed entities + error message
+-- that captures all schema errors that prisma returns, if any.
+-- Prisma format does not only do formatting, but also fixes some small mistakes and reports errors.
+prismaFormat :: [Entity] -> IO (Maybe Text, [Entity])
+prismaFormat unformattedEntities = do
+  let pslModels = getPslModelTextForEntity <$> unformattedEntities
+  (maybeErrorsMsg, formattedPslModels) <- Prisma.prismaFormatModels pslModels
+  let formattedEntities =
+        zipWith
+          (\e m -> e {entityBodyPsl = T.unpack $ getPslBodyFromPslModelText m})
+          unformattedEntities
+          formattedPslModels
+  return (maybeErrorsMsg, formattedEntities)
+  where
+    getPslModelTextForEntity :: Entity -> Text
+    getPslModelTextForEntity entity =
+      let modelName = T.pack $ entityName entity
+          modelBody = T.pack $ entityBodyPsl entity
+       in [trimming|model ${modelName} {
+                      ${modelBody}
+                    }|]
+
+    -- Example: @getPslBodyFromPslModelText "model Task {\n  id Int\n  desc String\n}" == "  id Int\n  desc String"@.
+    getPslBodyFromPslModelText = removeEnd . removeStart . T.strip
+      where
+        removeStart = T.dropWhile (== '\n') . T.drop 1 . T.dropWhile (/= '{')
+        removeEnd = T.dropWhileEnd isSpace . T.dropEnd 1 . T.dropWhileEnd (/= '}')
+
 checkPlanForOperationIssues :: OperationType -> Plan -> [String]
 checkPlanForOperationIssues opType plan =
   checkNumOperations
@@ -218,6 +267,26 @@ checkPlanForPageIssues plan =
           ]
         else []
 
+summarizePlan :: Plan -> Text
+summarizePlan plan =
+  let numQueries = showT $ length $ queries plan
+      numActions = showT $ length $ actions plan
+      numPages = showT $ length $ pages plan
+      numEntities = showT $ length $ entities plan
+      queryNames = showT $ opName <$> queries plan
+      actionNames = showT $ opName <$> actions plan
+      pageNames = showT $ pageName <$> pages plan
+      entityNames = showT $ entityName <$> entities plan
+   in [trimming|
+        - ${numQueries} queries: ${queryNames}
+        - ${numActions} actions: ${actionNames}
+        - ${numEntities} entities: ${entityNames}
+        - ${numPages} pages: ${pageNames}
+      |]
+  where
+    showT :: Show a => a -> Text
+    showT = T.pack . show
+
 -- TODO: Alternative idea is to give quite more autonomy and equip it with tools (functions) it
 --   needs to build correct context, and then let it drive itself completely on its own. So we give
 --   it app description, and then let it just go for it -> it would be making freeform plans for

waspc/src/Wasp/AI/OpenAI/ChatGPT.hs

@@ -15,11 +15,11 @@ module Wasp.AI.OpenAI.ChatGPT
   )
 where
 
-import Control.Arrow ()
 import Control.Monad (when)
 import Data.Aeson (FromJSON, ToJSON, (.=))
 import qualified Data.Aeson as Aeson
 import Data.ByteString.UTF8 as BSU
+import Data.Functor ((<&>))
 import Data.Text (Text)
 import Debug.Pretty.Simple (pTrace)
 import GHC.Generics (Generic)
@@ -27,6 +27,7 @@ import qualified Network.HTTP.Conduit as HTTP.C
 import qualified Network.HTTP.Simple as HTTP
 import Wasp.AI.OpenAI (OpenAIApiKey)
 import qualified Wasp.Util as Util
+import qualified Wasp.Util.Network.HTTP as Utils.HTTP
 
 -- | Might throw an HttpException.
 queryChatGPT :: OpenAIApiKey -> ChatGPTParams -> [ChatMessage] -> IO ChatResponse
@@ -45,9 +46,9 @@ queryChatGPT apiKey params requestMessages = do
             HTTP.setRequestBodyJSON reqBodyJson $
               HTTP.parseRequest_ "POST https://api.openai.com/v1/chat/completions"
 
-  response <- HTTP.httpJSON request
-
-  let (chatResponse :: ChatResponse) = HTTP.getResponseBody response
+  (chatResponse :: ChatResponse) <-
+    Utils.HTTP.httpJSONThatThrowsIfNot2xx request
+      <&> either (error . ("Failed to parse ChatGPT response body as JSON: " <>)) Prelude.id
 
   when True $
     pTrace

waspc/src/Wasp/Psl/Format.hs

@@ -1,16 +1,63 @@
-module Wasp.Psl.Format () where
+module Wasp.Psl.Format
+  ( prismaFormat,
+    PrismaFormatResult,
+    PslModelText,
+    PslErrorsMsg,
+    prismaFormatModels,
+  )
+where
 
+import Data.Aeson ((.:), (.:?))
+import qualified Data.Aeson as Aeson
+import Data.Function ((&))
 import Data.Text (Text)
-import Wasp.Psl.Ast.Model (Model)
+import qualified Data.Text as T
+import System.Exit (ExitCode (..))
+import qualified System.Process as P
+import qualified Wasp.Package as WP
+import Wasp.Util.Aeson (decodeFromString)
 
 -- | For given prisma schema source, returns formatted schema + any warnings/errors,
 -- by calling "prisma format" in the background.
+-- It might fail while attempting to do all that, in which case it will return error message.
 -- "prisma format" does more than just formatting -> it also applies some obvious fixes,
 -- like missing relationship fields. So it is kind of like compiling + formatting + fixing.
-prismaFormat :: Text -> IO (Maybe String, Text)
-prismaFormat prismaSchema = error "TODO"
+-- It works even for a prisma schema that has only model declarations!
+prismaFormat :: Text -> IO PrismaFormatResult
+prismaFormat prismaSchema = do
+  cp <- WP.getPackageProc WP.PrismaPackage ["format"]
+  (exitCode, response, stderr) <- P.readCreateProcessWithExitCode cp $ T.unpack prismaSchema
+  case exitCode of
+    ExitSuccess ->
+      return $
+        decodeFromString response
+          & either (error . ("Failed to parse response json from wasp's prisma ts package: " <>)) id
+    _exitFailure -> error $ "Failed while calling prisma format via wasp's prisma ts package: " <> stderr
 
--- "prisma format" works even without other declarations but models! That is great.
+data PrismaFormatResult = PrismaFormatResult
+  { _formattedSchemaPsl :: Text,
+    _schemaErrors :: Maybe Text
+  }
+  deriving (Show)
 
-prismaFormatModels :: [Model] -> IO (Maybe String, [Model])
-prismaFormatModels models = error "TODO: implement via prismaFormat"
+instance Aeson.FromJSON PrismaFormatResult where
+  parseJSON = Aeson.withObject "PrismaFormatResult" $ \obj -> do
+    formattedSchemaPsl <- obj .: "formattedSchemaPsl"
+    errors <- obj .:? "errors"
+    return (PrismaFormatResult {_formattedSchemaPsl = formattedSchemaPsl, _schemaErrors = errors})
+
+type PslModelText = Text
+
+type PslErrorsMsg = Text
+
+-- | Given a list of psl models in textual format (e.g. ["model User {\n...\n}", ...]),
+-- it returns back a list of those models but formatted, and also prisma format errors message,
+-- if there are any errors.
+prismaFormatModels :: [PslModelText] -> IO (Maybe PslErrorsMsg, [PslModelText])
+prismaFormatModels models = do
+  let schema = T.intercalate ("\n" <> delimiter <> "\n") models
+  result <- prismaFormat schema
+  let formattedModels = T.strip <$> T.splitOn delimiter (_formattedSchemaPsl result)
+  return (_schemaErrors result, formattedModels)
+  where
+    delimiter = "//==== WASP ====//"

waspc/src/Wasp/Util/Network/HTTP.hs

@@ -1,9 +1,15 @@
 module Wasp.Util.Network.HTTP
   ( catchRetryableHttpException,
+    getHttpExceptionStatusCode,
+    httpJSONThatThrowsIfNot2xx,
   )
 where
 
 import Control.Arrow ()
+import Control.Monad (void, when)
+import Control.Monad.IO.Class (MonadIO)
+import Data.Aeson (FromJSON)
+import qualified Data.Aeson as Aeson
 import qualified Network.HTTP.Conduit as HTTP.C
 import qualified Network.HTTP.Simple as HTTP
 import UnliftIO (MonadUnliftIO)
@@ -24,3 +30,21 @@ catchRetryableHttpException action handle =
       HTTP.getResponseStatusCode response `elem` retrayableHttpErrorStatusCodes
 
     retrayableHttpErrorStatusCodes = [503, 429, 408, 502, 504]
+
+getHttpExceptionStatusCode :: HTTP.HttpException -> Maybe Int
+getHttpExceptionStatusCode = \case
+  HTTP.HttpExceptionRequest _req (HTTP.C.StatusCodeException response _) ->
+    Just $ HTTP.getResponseStatusCode response
+  _otherwise -> Nothing
+
+-- | Throws an HttpException if status is not 2xx.
+-- Returns JSON parse error as Left if JSON parsing failed.
+httpJSONThatThrowsIfNot2xx :: (MonadIO m, FromJSON a) => HTTP.Request -> m (Either String a)
+httpJSONThatThrowsIfNot2xx request = do
+  response <- HTTP.httpLBS request
+
+  let statusCode = HTTP.getResponseStatusCode response
+  when (statusCode < 200 || statusCode >= 300) $
+    throwIO $ HTTP.HttpExceptionRequest request (HTTP.C.StatusCodeException (void response) "")
+
+  return $ Aeson.eitherDecode $ HTTP.getResponseBody response

waspc/waspc.cabal

@@ -341,6 +341,7 @@ library
     Wasp.Project.Env
     Wasp.Project.WebApp
     Wasp.Psl.Ast.Model
+    Wasp.Psl.Format
     Wasp.Psl.Generator.Model
     Wasp.Psl.Parser.Model
     Wasp.Psl.Util