From 53dfaaa5da75e6455ddb8e33289e80d4602cbd62 Mon Sep 17 00:00:00 2001
From: "ajaysi (aider)" <ajay.calsoft@gmail.com>
Date: Sun, 6 Oct 2024 13:56:06 +0530
Subject: [PATCH] feat: Collect all missing API keys using a Streamlit form.

---
 alwrity.py                   | 21 ++++++---------------
 lib/utils/api_key_manager.py | 25 +++++++++++++------------
 2 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/alwrity.py b/alwrity.py
index 1a56841b..90f85aab 100644
--- a/alwrity.py
+++ b/alwrity.py
@@ -218,21 +218,12 @@ def main():
     setup_environment_paths()
     sidebar_configuration()
 
-    missing_keys = check_api_keys()
-    if missing_keys:
-        st.warning(f"API keys not found: {', '.join(missing_keys)}. Please provide them below. Restart the app after saving the keys.")
-        with st.form(key='api_keys_form'):
-            for key in missing_keys:
-                st.text_input(f"{key}:", type="password", key=key)
-            if st.form_submit_button("Save Keys"):
-                with open(".env", "a") as env_file:
-                    for key in missing_keys:
-                        key_value = st.session_state[key]
-                        env_file.write(f"{key}={key_value}\n")
-                st.success("API keys saved successfully! Please restart the application.")
-    else:
-        setup_tabs()
-        modify_prompts_sidebar()
+    if not check_api_keys():
+        st.stop()
+    if not check_llm_environs():
+        st.stop()
+    setup_tabs()
+    modify_prompts_sidebar()
 
 
 def setup_environment_paths():
diff --git a/lib/utils/api_key_manager.py b/lib/utils/api_key_manager.py
index 8ffe050b..3d7a2da6 100644
--- a/lib/utils/api_key_manager.py
+++ b/lib/utils/api_key_manager.py
@@ -20,19 +20,20 @@ def check_api_keys():
         key: url for key, url in api_keys.items() if os.getenv(key) is None
     }
     if missing_keys:
-        st.error("🚨 Some API keys are missing! Please provide them below:")
-        for key, url in missing_keys.items():
-            api_key = st.text_input(f"Enter 🔏 {key}: 👉[Get it here]({url})👈")
-            if api_key:
-                os.environ[key] = api_key
-                try:
-                    with open(".env", "a") as env_file:
-                        env_file.write(f"{key}={api_key}\n")
-                except IOError as e:
-                    st.error(f"Failed to write {key} to .env file: {e}")
-                st.success(f"✅ {key} added successfully!")
+        st.warning(f"API keys not found: {', '.join(missing_keys)}. Please provide them below. Restart the app after saving the keys.")
+        with st.form(key='api_keys_form'):
+            for key, url in missing_keys.items():
+                st.text_input(f"{key}: 👉[Get it here]({url})👈", type="password", key=key)
+            if st.form_submit_button("Save Keys"):
+                with open(".env", "a") as env_file:
+                    for key in missing_keys:
+                        key_value = st.session_state[key]
+                        env_file.write(f"{key}={key_value}\n")
+                st.success("API keys saved successfully! Please restart the application.")
+                st.stop()
         return False
     return True
+
 @st.cache_data
 def check_llm_environs():
     """
@@ -43,7 +44,7 @@ def check_llm_environs():
     supported_providers = {
         'google': "GEMINI_API_KEY",
         'openai': "OPENAI_API_KEY",
-        'mistralai': "MISTRAL_API_KEY"
+        'mistral': "MISTRAL_API_KEY"
     }
     if not gpt_provider or gpt_provider.lower() not in supported_providers:
         gpt_provider = st.selectbox(