From 54aceb28a847cd80f0570a15576625abd2aa81da Mon Sep 17 00:00:00 2001
From: ajaysi <ajay.calsoft@gmail.com>
Date: Sun, 6 Oct 2024 13:56:02 +0530
Subject: [PATCH] feat: Add API key management and LLM provider selection

---
 lib/utils/api_key_manager.py | 75 ++++++++++++++++++++++++++++++------
 1 file changed, 63 insertions(+), 12 deletions(-)

diff --git a/lib/utils/api_key_manager.py b/lib/utils/api_key_manager.py
index 5bd4acd2..8ffe050b 100644
--- a/lib/utils/api_key_manager.py
+++ b/lib/utils/api_key_manager.py
@@ -2,20 +2,71 @@ import os
 import streamlit as st
 from dotenv import load_dotenv
 
-import os
-import streamlit as st
-
-import os
-import streamlit as st
 
+@st.cache_data
 def check_api_keys():
-    """Checks for API keys and returns a list of missing keys."""
-    required_keys = ["GOOGLE_API_KEY", "OPENAI_API_KEY", "ANTHROPIC_API_KEY"]
-    missing_keys = []
-    for key in required_keys:
-        if not os.getenv(key):
-            missing_keys.append(key)
-    return missing_keys
+    """
+    Checks if the required API keys are present in the environment variables.
+    Prompts the user to enter missing keys and saves them in the .env file.
+    """
+    api_keys = {
+        "METAPHOR_API_KEY": "https://dashboard.exa.ai/login",
+        "TAVILY_API_KEY": "https://tavily.com/#api",
+        "SERPER_API_KEY": "https://serper.dev/signup",
+        "STABILITY_API_KEY": "https://platform.stability.ai/",
+        "FIRECRAWL_API_KEY": "https://www.firecrawl.dev/account"
+    }
+    missing_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!")
+        return False
+    return True
+@st.cache_data
+def check_llm_environs():
+    """
+    Ensures that the LLM provider and corresponding API key are set.
+    Prompts the user to select a provider and enter the API key if missing.
+    """
+    gpt_provider = os.getenv("GPT_PROVIDER")
+    supported_providers = {
+        'google': "GEMINI_API_KEY",
+        'openai': "OPENAI_API_KEY",
+        'mistralai': "MISTRAL_API_KEY"
+    }
+    if not gpt_provider or gpt_provider.lower() not in supported_providers:
+        gpt_provider = st.selectbox(
+            "Select your LLM Provider", options=list(supported_providers.keys())
+        )
+        os.environ["GPT_PROVIDER"] = gpt_provider
+        try:
+            with open(".env", "a") as env_file:
+                env_file.write(f"GPT_PROVIDER={gpt_provider}\n")
+        except IOError as e:
+            st.error(f"Failed to write GPT_PROVIDER to .env file: {e}")
+        st.success(f"GPT Provider set to {gpt_provider}")
+    api_key_var = supported_providers[gpt_provider.lower()]
+    if not os.getenv(api_key_var):
+        api_key = st.text_input(f"Enter {api_key_var}:")
+        if api_key:
+            os.environ[api_key_var] = api_key
+            with open(".env", "a") as env_file:
+                env_file.write(f"{api_key_var}={api_key}\n")
+            st.success(f"{api_key_var} added successfully!")
+        return False
+    return True
+