Enhancements to Secret Management

Pipfile

@@ -49,6 +49,7 @@ h2o-wave = "0.26"
 tiktoken = "0.4.0"
 hf-transfer = "0.1.3"
 peft = "0.4.0"
+keyring = "24.2.0"
 
 [dev-packages]
 pytest = "==7.4.0"

README.md

@@ -53,6 +53,7 @@ Using CLI for fine-tuning LLMs:
 
 ## What's New
 
+- [PR 364](https://github.com/h2oai/h2o-llmstudio/pull/364) User secrets are now handled more securely and flexible. Support for handling secrets using the 'keyring' library was added. User settings are tried to be migrated automatically using the keyring library. Please monitor the logs for more information.
 - [PR 328](https://github.com/h2oai/h2o-llmstudio/pull/328) RLHF is now a separate problem type. Note that starting a new RLHF experiment from an old experiment that used RLHF is no longer supported. To continue from a previous experiment, please start a new experiment and enter the settings from the previous experiment manually. 
 - [PR 308](https://github.com/h2oai/h2o-llmstudio/pull/308) Sequence to sequence models have been added as a new problem type.
 - [PR 152](https://github.com/h2oai/h2o-llmstudio/pull/152) Add RLHF functionality for fine-tuning LLMs.

app_utils/config.py

@@ -84,6 +84,7 @@ def get_size(x):
     ],
     "user_settings": {
         "theme_dark": True,
+        "credential_saver": "Keyring",
         "default_aws_bucket_name": f"{os.getenv('AWS_BUCKET', 'bucket_name')}",
         "default_aws_access_key": os.getenv("AWS_ACCESS_KEY_ID", ""),
         "default_aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY", ""),

app_utils/handlers.py

@@ -42,7 +42,12 @@
     list_current_experiments,
 )
 from app_utils.sections.settings import settings
-from app_utils.utils import add_model_type, load_user_settings, save_user_settings
+from app_utils.utils.setting_utils import (
+    load_default_user_settings,
+    load_user_settings_and_secrets,
+    save_user_settings_and_secrets,
+)
+from app_utils.utils.utils import add_model_type
 from app_utils.wave_utils import report_error, wave_utils_handle_error
 
 logger = logging.getLogger(__name__)
@@ -73,13 +78,13 @@ async def handle(q: Q) -> None:
             await settings(q)
         elif q.args["save_settings"]:
             logger.info("Saving user settings")
-            save_user_settings(q)
+            await save_user_settings_and_secrets(q)
             await settings(q)
         elif q.args["load_settings"]:
-            load_user_settings(q)
+            load_user_settings_and_secrets(q)
             await settings(q)
         elif q.args["restore_default_settings"]:
-            load_user_settings(q, force_defaults=True)
+            load_default_user_settings(q)
             await settings(q)
 
         elif q.args["report_error"]:

app_utils/initializers.py

@@ -7,20 +7,20 @@
 from h2o_wave import Q
 
 from app_utils.sections.common import interface
-from llm_studio.src.utils.config_utils import load_config_py, save_config_yaml
-
-from .config import default_cfg
-from .db import Database, Dataset
-from .utils import (
+from app_utils.utils.utils import (
     get_data_dir,
     get_database_dir,
     get_download_dir,
     get_output_dir,
     get_user_db_path,
     get_user_name,
-    load_user_settings,
     prepare_default_dataset,
 )
+from llm_studio.src.utils.config_utils import load_config_py, save_config_yaml
+
+from .config import default_cfg
+from .db import Database, Dataset
+from .utils.setting_utils import load_user_settings_and_secrets
 
 logger = logging.getLogger(__name__)
 
@@ -97,7 +97,7 @@ async def initialize_client(q: Q) -> None:
 
         import_data(q)
 
-        load_user_settings(q)
+        load_user_settings_and_secrets(q)
 
         await interface(q)
 

app_utils/sections/chat.py

@@ -14,7 +14,12 @@
 from h2o_wave import ui
 from transformers import AutoTokenizer, TextStreamer
 
-from app_utils.utils import get_experiments, get_ui_elements, parse_ui_elements, set_env
+from app_utils.utils.utils import (
+    get_experiments,
+    get_ui_elements,
+    parse_ui_elements,
+    set_env,
+)
 from llm_studio.src.datasets.text_utils import get_tokenizer
 from llm_studio.src.models.text_causal_language_modeling_model import Model
 from llm_studio.src.utils.config_utils import load_config_yaml

app_utils/sections/dataset.py

@@ -12,7 +12,7 @@
 from app_utils.config import default_cfg
 from app_utils.db import Dataset
 from app_utils.sections.experiment import experiment_start
-from app_utils.utils import (
+from app_utils.utils.utils import (
     add_model_type,
     check_valid_upload_content,
     clean_error,

app_utils/sections/experiment.py

@@ -21,7 +21,7 @@
 from app_utils.config import default_cfg
 from app_utils.sections.chat import chat_tab, load_cfg_model_tokenizer
 from app_utils.sections.common import clean_dashboard
-from app_utils.utils import (
+from app_utils.utils.utils import (
     add_model_type,
     flatten_dict,
     get_cfg_list_items,

app_utils/sections/home.py

@@ -4,7 +4,7 @@
 
 from app_utils.config import default_cfg
 from app_utils.sections.common import clean_dashboard
-from app_utils.utils import (
+from app_utils.utils.utils import (
     get_datasets,
     get_experiments,
     get_gpu_usage,

app_utils/sections/project.py

@@ -10,7 +10,7 @@
     experiment_rename_form,
     get_table_and_message_item_indices,
 )
-from app_utils.utils import get_experiments_status
+from app_utils.utils.utils import get_experiments_status
 
 logger = logging.getLogger(__name__)
 

app_utils/sections/settings.py

@@ -4,6 +4,7 @@
 from h2o_wave import Q, ui
 
 from app_utils.sections.common import clean_dashboard
+from app_utils.utils.setting_utils import Secrets
 from llm_studio.src.loggers import Loggers
 
 
@@ -24,6 +25,37 @@ async def settings(q: Q) -> None:
                 'Save settings persistently' button below. To reload \
                 the persistently saved settings, use the 'Load settings' button.",
             ),
+            ui.separator("Credential Storage"),
+            ui.inline(
+                items=[
+                    ui.label("Credential Handler", width=label_width),
+                    ui.dropdown(
+                        name="credential_saver",
+                        value=q.client["credential_saver"],
+                        choices=[ui.choice(name, name) for name in Secrets.names()],
+                        trigger=False,
+                        width="300px",
+                    ),
+                ]
+            ),
+            ui.message_bar(
+                type="info",
+                text="""Method used to save credentials (passwords) \
+                for 'Save settings persistently'. \
+                The recommended approach for saving credentials (passwords) is to \
+                use either Keyring or to avoid permanent storage \
+                (requiring re-entry upon app restart). \
+                If Keyring isn't set up on the system, \
+                a pop-up will notify the user of the failure. \
+                In this scenario, you can choose not to save \
+                credentials long-term or utilize a local .env file. \
+                Only resort to the latter if your machine's \
+                accessibility is restricted to you.\
+                When you select 'Save settings persistently', \
+                credentials will be removed from all non-selected methods. \
+                'Restore Default Settings' will clear credentials from all methods.
+                """,
+            ),
             ui.separator("Appearance"),
             ui.inline(
                 items=[