From 16bcd86bb7239edf11903053756378d8a9459435 Mon Sep 17 00:00:00 2001
From: ajaysi <ajay.calsoft@gmail.com>
Date: Sun, 6 Oct 2024 16:22:12 +0530
Subject: [PATCH] API keys setup improvements & housekeeping

---
 alwrity.py                   | 12 ++++++------
 lib/utils/api_key_manager.py | 10 ++++++++--
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/alwrity.py b/alwrity.py
index b6484322..b190361a 100644
--- a/alwrity.py
+++ b/alwrity.py
@@ -210,14 +210,14 @@ def sidebar_configuration():
 
 def main():
     #load_environment
-    from dotenv import load_dotenv
     load_dotenv()
     setup_ui()
-    setup_environment_paths()
-    sidebar_configuration()
-    check_all_api_keys()
-    setup_tabs()
-    modify_prompts_sidebar()
+
+    if check_all_api_keys():
+        setup_environment_paths()
+        sidebar_configuration()
+        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 5730cc2a..8be43195 100644
--- a/lib/utils/api_key_manager.py
+++ b/lib/utils/api_key_manager.py
@@ -2,14 +2,15 @@ import os
 import streamlit as st
 from dotenv import load_dotenv
 
-
-@st.cache_data
 def check_all_api_keys():
     """
     Checks if all required API keys are present in the environment variables.
     Prompts the user to enter missing keys and saves them in the .env file.
     This includes general API keys and the LLM provider key.
     """
+    # Load environment variables from .env (MUST COME FIRST)
+    load_dotenv()
+
     api_keys = {
         "METAPHOR_API_KEY": "https://dashboard.exa.ai/login",
         "TAVILY_API_KEY": "https://tavily.com/#api",
@@ -18,6 +19,7 @@ def check_all_api_keys():
         "FIRECRAWL_API_KEY": "https://www.firecrawl.dev/account"
     }
 
+    # Check for missing keys AFTER loading environment variables
     missing_keys = {
         key: url for key, url in api_keys.items() if os.getenv(key) is None
     }
@@ -44,14 +46,18 @@ def check_all_api_keys():
     if not os.getenv(api_key_var):
         missing_keys[api_key_var] = ''
 
+    # If there are missing keys, prompt the user to enter them
     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'):
+            # Gather all missing keys in one go
             for key, url in missing_keys.items():
                 if url:
                     st.text_input(f"{key}: 👉[Get it here]({url})👈", type="password", key=key)
                 else:
                     st.text_input(f"{key}:", type="password", key=key)
+
+            # Save all keys at once when the button is clicked
             if st.form_submit_button("Save Keys"):
                 with open(".env", "a") as env_file:
                     for key in missing_keys: