Add option to specify Plan GPT model from CLI

wasp-lang · Dec 11, 2023 · 2559d30 · 2559d30
1 parent 745a02e
commit 2559d30
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/waspc/src/Wasp/AI/GenerateNewProject/Common.hs b/waspc/src/Wasp/AI/GenerateNewProject/Common.hs
@@ -9,6 +9,8 @@ module Wasp.AI.GenerateNewProject.Common
     queryChatGPTForJSON,
     defaultChatGPTParams,
     defaultChatGPTParamsForFixing,
+    defaultChatGPTParamsForPlan,
+    defaultChatGPTParamsForPlanFixing,
     writeToWaspFileEnd,
   )
 where
@@ -34,6 +36,7 @@ data NewProjectConfig = NewProjectConfig
     -- One of the Tailwind color names: https://tailwindcss.com/docs/customizing-colors
     projectPrimaryColor :: !(Maybe String),
     projectDefaultGptModel :: !(Maybe GPT.Model),
+    projectPlanGptModel :: !(Maybe GPT.Model),
     projectDefaultGptTemperature :: !(Maybe Float)
   }
   deriving (Show)
@@ -43,12 +46,14 @@ instance Aeson.FromJSON NewProjectConfig where
     auth <- obj .:? "auth"
     primaryColor <- obj .:? "primaryColor"
     defaultGptModel <- obj .:? "defaultGptModel"
+    planGptModel <- obj .:? "defaultGptModel"
     defaultGptTemperature <- obj .:? "defaultGptTemperature"
     return
       ( NewProjectConfig
           { projectAuth = auth,
             projectPrimaryColor = primaryColor,
             projectDefaultGptModel = defaultGptModel,
+            projectPlanGptModel = planGptModel,
             projectDefaultGptTemperature = defaultGptTemperature
           }
       )
@@ -59,6 +64,7 @@ emptyNewProjectConfig =
     { projectAuth = Nothing,
       projectPrimaryColor = Nothing,
       projectDefaultGptModel = Nothing,
+      projectPlanGptModel = Nothing,
       projectDefaultGptTemperature = Nothing
     }
 
@@ -145,6 +151,16 @@ defaultChatGPTParamsForFixing projectDetails =
   let params = defaultChatGPTParams projectDetails
    in params {GPT._temperature = subtract 0.2 <$> GPT._temperature params}
 
+defaultChatGPTParamsForPlan :: NewProjectDetails -> ChatGPTParams
+defaultChatGPTParamsForPlan projectDetails =
+  let params = defaultChatGPTParams projectDetails
+   in params {GPT._model = fromMaybe GPT.GPT_4 (projectPlanGptModel $ _projectConfig projectDetails)}
+
+defaultChatGPTParamsForPlanFixing :: NewProjectDetails -> ChatGPTParams
+defaultChatGPTParamsForPlanFixing projectDetails =
+  let params = defaultChatGPTParamsForPlan projectDetails
+   in params {GPT._temperature = subtract 0.2 <$> GPT._temperature params}
+
 writeToWaspFileEnd :: FilePath -> Text -> CodeAgent ()
 writeToWaspFileEnd waspFilePath text = do
   writeToFile waspFilePath $

diff --git a/waspc/src/Wasp/AI/GenerateNewProject/Plan.hs b/waspc/src/Wasp/AI/GenerateNewProject/Plan.hs
@@ -25,13 +25,13 @@ import qualified Text.Parsec as Parsec
 import Wasp.AI.CodeAgent (CodeAgent, writeToLog)
 import Wasp.AI.GenerateNewProject.Common
   ( NewProjectDetails (..),
-    defaultChatGPTParams,
-    defaultChatGPTParamsForFixing,
+    defaultChatGPTParamsForPlan,
+    defaultChatGPTParamsForPlanFixing,
     queryChatGPTForJSON,
   )
 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 Wasp.AI.OpenAI.ChatGPT (ChatMessage (..), ChatRole (..))
 import qualified Wasp.Psl.Format as Prisma
 import qualified Wasp.Psl.Parser.Model as Psl.Parser
 import qualified Wasp.Util.Aeson as Util.Aeson
@@ -42,7 +42,7 @@ type PlanRule = String
 generatePlan :: NewProjectDetails -> [PlanRule] -> CodeAgent Plan
 generatePlan newProjectDetails planRules = do
   writeToLog "Generating plan (slowest step, usually takes 30 to 90 seconds)..."
-  initialPlan <- queryChatGPTForJSON ((defaultChatGPTParams newProjectDetails) {_model = planGptModel}) chatMessages
+  initialPlan <- queryChatGPTForJSON (defaultChatGPTParamsForPlan newProjectDetails) chatMessages
   writeToLog $ "Initial plan generated!\n" <> summarizePlan initialPlan
   writeToLog "Fixing initial plan..."
   fixedPlan <- fixPlanRepeatedly 3 initialPlan
@@ -170,7 +170,7 @@ generatePlan newProjectDetails planRules = do
                     |]
           writeToLog "Sending plan to GPT for fixing..."
           fixedPlan <-
-            queryChatGPTForJSON ((defaultChatGPTParamsForFixing newProjectDetails) {_model = planGptModel}) $
+            queryChatGPTForJSON (defaultChatGPTParamsForPlanFixing newProjectDetails) $
               chatMessages
                 <> [ ChatMessage {role = Assistant, content = Util.Aeson.encodeToText plan'},
                      ChatMessage
@@ -190,8 +190,6 @@ generatePlan newProjectDetails planRules = do
                    ]
           return (False, fixedPlan)
 
-    planGptModel = GPT_4
-
 checkPlanForEntityIssues :: Plan -> [String]
 checkPlanForEntityIssues plan =
   checkNumEntities