Added workaround for API prompting and saving

Author: KTS-o7

requirements.txt

@@ -7,7 +7,6 @@ attrs==24.2.0
 certifi==2024.8.30
 charset-normalizer==3.3.2
 click==8.1.7
-colorama==0.4.6
 distro==1.9.0
 docstring_parser==0.16
 frozenlist==1.4.1

setup.py

@@ -1,54 +1,6 @@
 from setuptools import setup, find_packages
-from setuptools.command.install import install
 import os
 
-DB_PATH = os.path.expanduser('~/.snapshell/system_info.db')
-
-class CustomInstallCommand(install):
-    def run(self):
-        
-        # ANSI escape codes for colors
-        GREEN = '\033[92m'
-        YELLOW = '\033[93m'
-        RESET = '\033[0m'
-        
-        # Prompt user for GROQ API key
-        print(GREEN + "Please enter your GROQ API key:" + RESET)
-        groq_api_key = input(YELLOW + "> " + RESET)
-        
-        # Detect the user's shell
-        user_shell = os.environ.get('SHELL', '/bin/bash')
-        shell_config_file = {
-            '/bin/bash': '~/.bashrc',
-            '/bin/zsh': '~/.zshrc',
-            '/bin/fish': '~/.config/fish/config.fish'
-        }.get(user_shell)
-        
-        if shell_config_file:
-            shell_config_path = os.path.expanduser(shell_config_file)
-            
-            # Set the GROQ API key in the user's shell config file
-            with open(shell_config_path, "a") as shell_config:
-                shell_config.write(f'\nexport HELPER_GROQ_API_KEY="{groq_api_key}"\n')
-                print(GREEN + f"Added GROQ API key to {shell_config_path}" + RESET)
-            
-            # Reload shell configuration
-            os.system(f'source {shell_config_path}')
-            from snapshell.database import create_database, update_database
-            print(YELLOW + f"Setting up the database... at {DB_PATH}" + RESET)
-            create_database()
-            install.run(self)
-            update_database()
-            
-            print(GREEN + "Database successfully created\n" + RESET)
-            print(GREEN + "use the tool as snapshell\n" + RESET)
-            # Restart the shell
-            os.execvp(user_shell, [user_shell])
-            
-        else:
-            print(f"Unrecognized shell: {user_shell}. Please manually add the following lines to your shell configuration file:")
-            print(f'export HELPER_GROQ_API_KEY="{groq_api_key}"')
-
 setup(
     name='snapshell',
     version='1.0.0',
@@ -81,7 +33,4 @@ def run(self):
         'Operating System :: OS Independent',
     ],
     python_requires='>=3.10',
-    cmdclass={
-        'install': CustomInstallCommand,
-    },
 )

snapshell/cli.py

@@ -1,8 +1,8 @@
 import argparse
-from .llm_api import suggest_command
-from .database import update_database, DB_PATH
-import sqlite3
 import os
+import sqlite3
+from snapshell.database import create_database, update_database, DB_PATH
+from .llm_api import suggest_command, set_api_key, load_api_key
 
 # ANSI escape codes for colors
 RESET = "\033[0m"
@@ -13,6 +13,20 @@
 YELLOW = "\033[33m"
 RED = "\033[31m"
 
+def initial_setup():
+    # Prompt user for GROQ API key
+    print(GREEN + "Please enter your GROQ API key:" + RESET)
+    groq_api_key = input(YELLOW + "> " + RESET)
+    
+    # Set the API key
+    set_api_key(groq_api_key)
+    
+    print(YELLOW + f"Setting up the database... at {DB_PATH}" + RESET)
+    create_database()
+    update_database()
+    print(GREEN + "Database successfully created\n" + RESET)
+    print(GREEN + "Use the tool as snapshell\n" + RESET)
+
 def view_history():
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
@@ -46,8 +60,22 @@ def main():
     parser.add_argument('--update-db', action='store_true', help="Update the database with installed packages")
     parser.add_argument('--view-history', action='store_true', help="View command history")
     parser.add_argument('--clear-history', action='store_true', help="Clear command history")
+    parser.add_argument('--set-api-key', type=str, help="Set the GROQ API key")
     args = parser.parse_args()
 
+    if args.set_api_key:
+        set_api_key(args.set_api_key)
+        print(GREEN + "API key set successfully." + RESET)
+        return
+
+    # Load the API key from the configuration file
+    API_KEY =  load_api_key()
+
+
+    if not API_KEY:
+        print("This is working")
+        initial_setup()
+    
     if args.update_db:
         print(f"{YELLOW}Updating database...{RESET}")
         update_database()
@@ -60,7 +88,7 @@ def main():
     if args.clear_history:
         clear_history()
         return
-
+    
     print(f"{CYAN}Welcome to the Linux Command Tool. Type 'exit' to quit.{RESET}")
 
     conversation_history = []

snapshell/llm_api.py

@@ -1,16 +1,18 @@
 # llm_api.py
 import os
+import json
 from groq import Groq
 from pydantic import BaseModel, Field, ValidationError
 from .system_info import fetch_system_info
 from .package_managers import detect_package_manager
 import sqlite3
 from .database import DB_PATH, save_command_suggestion
 
+# Configuration file path
+CONFIG_FILE = os.path.expanduser("~/.snapshell_config.json")
 
-API_KEY = os.getenv('HELPER_GROQ_API_KEY')
-
-groq = Groq(api_key=API_KEY)
+# Global variable to store the API key
+API_KEY = None
 
 # Data model for LLM to generate
 class CommandSuggestion(BaseModel):
@@ -21,9 +23,26 @@ class CommandSuggestion(BaseModel):
 class SQLQuery(BaseModel):
     query: str = Field(description="The SQL query to be executed")
 
+def load_api_key():
+    global API_KEY
+    if os.path.exists(CONFIG_FILE):
+        with open(CONFIG_FILE, "r") as config_file:
+            config = json.load(config_file)
+            API_KEY = config.get("api_key")
+    return API_KEY
+
+def set_api_key(api_key):
+    global API_KEY
+    API_KEY = api_key
+    config = {"api_key": api_key}
+    with open(CONFIG_FILE, "w") as config_file:
+        json.dump(config, config_file)
+
 def suggest_command(user_input, conversation_history):
     if not API_KEY:
-        raise ValueError("API key not found. Please set the GROQ_API_KEY environment variable.")
+        raise ValueError("API key not set. Please set the API key using set_api_key function.")
+
+    groq = Groq(api_key=API_KEY)
 
     system_info = fetch_system_info()
     package_manager = detect_package_manager()
@@ -103,6 +122,8 @@ def fallback_to_llm(user_input, package_manager, fallback_message, conversation_
 
     messages.append({"role": "user", "content": user_input})
 
+    groq = Groq(api_key=API_KEY)
+
     chat_completion = groq.chat.completions.create(
         messages=messages,
         model="llama-3.2-90b-text-preview",
@@ -131,6 +152,8 @@ def formulate_sql_query(user_input):
         {"role": "user", "content": user_input},
     ]
 
+    groq = Groq(api_key=API_KEY)
+
     chat_completion = groq.chat.completions.create(
         messages=messages,
         model="llama-3.2-90b-text-preview",
@@ -170,4 +193,7 @@ def query_database(sql_query):
         return relevant_packages
     except sqlite3.Error as e:
         # If there's any issue with the database query, return an empty result
-        return []
+        return []
+
+# Load the API key from the configuration file on module import
+load_api_key()