From a06bb9b042f75f213bbbf0047e22dc4564bfa930 Mon Sep 17 00:00:00 2001
From: Jason Schleifer <jason.schleifer@gmail.com>
Date: Sat, 21 Dec 2024 12:46:06 +1300
Subject: [PATCH] updated defaults

---
 nodes/drivers/gtUIOpenAiChatPromptDriver.py      | 5 +++--
 nodes/drivers/gtUIOpenAiImageGenerationDriver.py | 6 ++++--
 nodes/utils/openai_utils.py                      | 2 +-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/nodes/drivers/gtUIOpenAiChatPromptDriver.py b/nodes/drivers/gtUIOpenAiChatPromptDriver.py
index ad16349..2a6850b 100644
--- a/nodes/drivers/gtUIOpenAiChatPromptDriver.py
+++ b/nodes/drivers/gtUIOpenAiChatPromptDriver.py
@@ -4,6 +4,7 @@ from ..utils.openai_utils import get_available_models
 from .gtUIBasePromptDriver import gtUIBasePromptDriver
 
 models = get_available_models("ChatModel")
+DEFAULT_MODEL = "gpt-4o-mini"
 DEFAULT_API_KEY = "OPENAI_API_KEY"
 
 
@@ -27,7 +28,7 @@ class gtUIOpenAiChatPromptDriver(gtUIBasePromptDriver):
         inputs["optional"].update(base_optional_inputs)
 
         # Set model default
-        inputs["optional"]["model"] = (models, {"default": models[0]})
+        inputs["optional"]["model"] = (models, {"default": DEFAULT_MODEL})
         inputs["optional"].update(
             {
                 "response_format": (
@@ -50,7 +51,7 @@ class gtUIOpenAiChatPromptDriver(gtUIBasePromptDriver):
 
     def build_params(self, **kwargs):
         api_key = self.getenv(kwargs.get("api_key_env_var", DEFAULT_API_KEY))
-        model = kwargs.get("model", None)
+        model = kwargs.get("model", DEFAULT_MODEL)
         response_format = kwargs.get("response_format", None)
         seed = kwargs.get("seed", None)
         stream = kwargs.get("stream", False)
diff --git a/nodes/drivers/gtUIOpenAiImageGenerationDriver.py b/nodes/drivers/gtUIOpenAiImageGenerationDriver.py
index c5b4bcf..c5cd563 100644
--- a/nodes/drivers/gtUIOpenAiImageGenerationDriver.py
+++ b/nodes/drivers/gtUIOpenAiImageGenerationDriver.py
@@ -6,6 +6,8 @@ from .gtUIBaseImageDriver import gtUIBaseImageGenerationDriver
 DEFAULT_API_KEY = "OPENAI_API_KEY"
 # models = ["dall-e-3", "dall-e-2"]
 models = get_available_models("ImageModel")
+DEFAULT_MODEL = "dall-e-3"
+
 sizes = ["256x256", "512x512", "1024x1024", "1024x1792", "1792x1024"]
 
 
@@ -20,7 +22,7 @@ class gtUIOpenAiImageGenerationDriver(gtUIBaseImageGenerationDriver):
                 "image_generation_model": (
                     models,
                     {
-                        "default": models[0],
+                        "default": DEFAULT_MODEL,
                         "tooltip": "Select the image generation model.",
                     },
                 ),
@@ -52,7 +54,7 @@ class gtUIOpenAiImageGenerationDriver(gtUIBaseImageGenerationDriver):
 
     def build_params(self, **kwargs):
         size_from_args = kwargs.get("size", sizes[2])
-        model = kwargs.get("image_generation_model", models[0])
+        model = kwargs.get("image_generation_model", DEFAULT_MODEL)
         size = self.adjust_size_based_on_model(model, size_from_args)
         api_key = self.getenv(kwargs.get("api_key_env_var", DEFAULT_API_KEY))
 
diff --git a/nodes/utils/openai_utils.py b/nodes/utils/openai_utils.py
index 2ab72b4..5d63619 100644
--- a/nodes/utils/openai_utils.py
+++ b/nodes/utils/openai_utils.py
@@ -32,7 +32,7 @@ def get_available_models(model_type: ModelTypes) -> list:
     }
 
     settings = GriptapeSettings()
-    api_key = settings.get(API_KEY)
+    api_key = settings.get_settings_key_or_use_env(API_KEY)
     if api_key is None:
         return []
     client = OpenAI()