Usability and UI Improvements

app.py

@@ -7,66 +7,205 @@
 
 import config
 
-
 def setup():
     if not os.path.exists(config.model_folder):
         if input(f"Model folder {config.model_folder} does not exist. Create it? (y/n) ").lower() == 'y':
             os.mkdir(config.model_folder)
-    current = os.listdir(config.model_folder)
     for model in config.models:
-        if model == 'default':
-            continue
-        if config.models[model]['type'] == 'local':
-            if config.models[model]['filename'] not in current:
-                if input(f'Model {model} not found in {config.model_folder}. Would you like to download it? (y/n) ').lower() == 'y':
-                    url = config.models[model]['url']
-                    print(f"Downloading {model} from {url}...")
-                    subprocess.run(['curl', '-L', url, '-o', os.path.join(
-                        config.model_folder, config.models[model]['filename'])])
-            else:
-                print(f"Model {model} found in {config.model_folder}.")
+        if is_installed(model):
+            print(f"Model {model} found in {config.model_folder}.")   
+        else:
+            if input(f'Model {model} not found in {config.model_folder}. Would you like to download it? (y/n) ').lower() == 'y':
+                    install_model(model, True)
+
+
+def is_installed(model):
+    if config.models[model]['type'] == 'local':
+        return os.path.exists(config.model_folder + "/" + config.models[model]['filename'])
+    else:
+        try:
+            response = requests.get("http://google.com", timeout=1.0)
+            return response.status_code == 200
+        except requests.RequestException as e:
+            return False
+
+
+def install_model(model, verbose = False, app = None):
+    url = config.models[model]['url']
+    if verbose:
+        print(f"Downloading {model} from {url}...")
 
+    if app is not None:
+        lock_model(model)
+        subprocess.run([
+            'osascript',
+            '-e',
+            # install the model in a new terminal window, then delete the lock file
+            f'tell application "Terminal" to do script "curl -L {url} -o {os.path.join(config.model_folder, config.models[model]["filename"])} && rm {config.model_folder}/{model}.lock"'])
+        return
+    else:
+        subprocess.run(['curl', '-L', url, '-o', os.path.join(
+            config.model_folder, config.models[model]['filename'])])
+
+def lock_model(model):
+    with open(f"{config.model_folder}/{model}.lock", 'w') as f:
+        f.write("")
+
+def is_installing(model):
+    if config.models[model]['type'] == 'local':
+        return os.path.exists(f"{config.model_folder}/{model}.lock")
+    else:
+        return False
 
 class ModelPickerApp(rumps.App):
     def __init__(self):
-        super(ModelPickerApp, self).__init__("ModelPickerApp")
+        super(ModelPickerApp, self).__init__("ModelPickerApp", quit_button=None)
 
-        # Dynamically create menu items from the MENUBAR_OPTIONS
+        self.rebuild_menu()
+        self.icon = config.ICON
+        rumps.Timer(self.update_menu, 5).start()
+
+    def rebuild_menu(self):
+        self.menu.clear()
         self.menu_items = {}
+        show_uninstalled = config.get_settings(config.SETTINGS_SHOW_UNINSTALLED)
         for option in config.models:
-            if option == 'default':
+            if not show_uninstalled and config.models[option]['type'] == 'local' and not is_installed(option):
                 continue
             self.menu_items[option] = rumps.MenuItem(
-                title=option, callback=self.pick_model)
+                title=option, callback=self.pick_model, icon=None)
+
+        self.menu_items['Settings'] = rumps.MenuItem(
+            title='Settings', icon=None)
+
+        self.add_bool_setting(config.SETTINGS_SHOW_UNINSTALLED, True)
+        self.add_bool_setting(config.SETTINGS_SHOW_STATUS_ICONS, True)
+        self.add_bool_setting(config.SETTINGS_SHOW_CURRENT_MODEL, True)
+        self.add_settings_menu(config.SETTINGS_SWITCH, ['Automatic', 'Trigger Offline', 'Manual'])
+        self.add_settings_menu(config.SETTINGS_DEFAULT_ONLINE, filter(lambda x: config.models[x]['type'] == 'remote', config.models))
+        self.add_settings_menu(config.SETTINGS_DEFAULT_OFFLINE, filter(lambda x: config.models[x]['type'] == 'local' and (show_uninstalled or is_installed(x)), config.models))
+
+        self.menu_items['Quit'] = rumps.MenuItem(title='Quit', callback=rumps.quit_application, icon=None)
 
         self.menu = list(self.menu_items.values())
-        self.menu_items[config.models['default']].state = True
-        self.icon = "icon.png"
+        self.menu_items[config.get_settings(config.SETTINGS_CURRENT_MODEL)].state = True
+        self.update_menu(None)
 
-    def pick_model(self, sender):
-        # Toggle the checked status of the clicked menu item
-        sender.state = not sender.state
+        if config.get_settings(config.SETTINGS_SHOW_CURRENT_MODEL):
+            self.title = config.get_settings(config.SETTINGS_CURRENT_MODEL)
+        else:
+            self.title = None
 
-        # Send the choice to the local proxy app
+    def add_settings_menu(self, name, options, triggerRebuild = False):
+        self.menu_items["Settings"].add(rumps.MenuItem(title=name, icon=None))
+        selected_option = config.get_settings(name)
+
+        for option in options:
+            self.menu_items["Settings"][name].add(
+                rumps.MenuItem(title=option, callback=lambda sender: self.set_setting(sender, name, triggerRebuild), icon=None))
+            if option == selected_option:
+                self.menu_items["Settings"][name][option].state = True
+
+    def add_bool_setting(self, name, triggerRebuild = False):
+        self.menu_items["Settings"].add(
+                rumps.MenuItem(title=name, callback=lambda sender: self.set_bool_setting(sender, name, triggerRebuild), icon=None))
+        self.menu_items["Settings"][name].state = config.get_settings(name)
+
+    def set_setting(self, sender, setting, triggerRebuild = False):
         if sender.state:
-            choice = sender.title
+            return
+
+        config.set_settings(setting, sender.title)
+
+        if triggerRebuild:
+            self.rebuild_menu()
+        else:
+            for item in self.menu['Settings'][setting]:
+                self.menu_items['Settings'][setting][item].state = item == sender.title
+
+    def set_bool_setting(self, sender, setting, triggerRebuild = False):
+        config.set_settings(setting, not sender.state)
+
+        if triggerRebuild:
+            self.rebuild_menu()
+        else:
+            sender.state = config.get_settings(setting)
+
+
+    def update_menu(self, sender):
+        status_icons = config.get_settings(config.SETTINGS_SHOW_STATUS_ICONS)
+        for option in self.menu_items:
+            if not option in config.models:
+                continue
+            if status_icons:
+                if is_installing(option):
+                    self.menu_items[option].icon = config.ICON_INSTALLING
+                elif is_installed(option):
+                    self.menu_items[option].icon = config.ICON_INSTALLED
+                else:
+                    self.menu_items[option].icon = config.ICON_UNINSTALLED
+            else:
+                self.menu_items[option].icon = None
+
+        currently_online = config.models[config.get_settings(config.SETTINGS_CURRENT_MODEL)]['type'] == 'remote'
+
+        if currently_online and config.get_settings(config.SETTINGS_SWITCH) == 'Automatic' or config.get_settings(config.SETTINGS_SWITCH) == 'Trigger Offline':
+            try:
+                response = requests.get("http://google.com", timeout=1.0)
+                if response.status_code != 200:
+                    self.pick_model(self.menu_items[config.get_settings(config.SETTINGS_DEFAULT_OFFLINE)])
+            except requests.RequestException as e:
+                self.pick_model(self.menu_items[config.get_settings(config.SETTINGS_DEFAULT_OFFLINE)])
+
+        if not currently_online and config.get_settings(config.SETTINGS_SWITCH) == 'Automatic':
             try:
-                response = requests.post(
-                    "http://localhost:5001/set_target", json={"target": choice})
+                response = requests.get("http://google.com", timeout=1.0)
                 if response.status_code == 200:
-                    print(f"Successfully sent selection: {choice}.")
-                else:
-                    rumps.alert(
-                        "Error", f"Failed to send selection. Server responded with: {response.status_code}.")
+                    self.pick_model(self.menu_items[config.get_settings(config.SETTINGS_DEFAULT_ONLINE)])
             except requests.RequestException as e:
-                rumps.alert("Error", f"Failed to send selection. Error: {e}.")
+                pass
+
+    def pick_model(self, sender):
+        if (sender.state):
+            return
+
+        # check if the model is installed
+        if is_installing(sender.title):
+            rumps.alert("Model Installing", f"{sender.title} is currently installing.")
+            return
+        elif not is_installed(sender.title):
+            if config.models[sender.title]['type'] == 'remote':
+                return
+            if (rumps.alert("Install Model", f"Install {sender.title}?", cancel = True) == 1):
+                install_model(sender.title, app = self)
+                return
+            else:
+                return
+
+        # Send the choice to the local proxy app
+        choice = sender.title
+        try:
+            response = requests.post(
+                "http://localhost:5001/set_target", json={"target": choice}, timeout=1.0)
+            if response.status_code == 200:
+                print(f"Successfully sent selection: {choice}.")
+            else:
+                rumps.alert(
+                    "Error", f"Failed to send selection. Server responded with: {response.status_code}.")
+        except requests.RequestException as e:
+            rumps.alert("Error", f"Failed to send selection. Error: {e}.")
+            return
+
+        # Toggle the checked status of the clicked menu item
+        if config.get_settings(config.SETTINGS_SHOW_CURRENT_MODEL):
+            self.title = sender.title
+        config.set_settings(config.SETTINGS_CURRENT_MODEL, sender.title)
 
         # If other options were previously selected, deselect them
         for item in self.menu:
             if item == 'Quit':
                 continue
-            if item != sender.title:
-                self.menu_items[item].state = False
+            self.menu_items[item].state = item == sender.title
 
     def run_server(self):
         subprocess.run(['python', 'proxy.py'])

config.py

@@ -1,4 +1,22 @@
 import os
+import json
+
+ICON = "resources/icon.png"
+ICON_INSTALLED = "resources/installed.png"
+ICON_INSTALLING = "resources/installing.png"
+ICON_UNINSTALLED = "resources/uninstalled.png"
+
+SETTINGS_CURRENT_MODEL = "Current Model"
+SETTINGS_SHOW_UNINSTALLED = "Show Uninstalled"
+SETTINGS_SHOW_STATUS_ICONS = "Show Status Icons"
+SETTINGS_SHOW_CURRENT_MODEL = "Show Current Model"
+SETTINGS_DEFAULT_ONLINE = "Default Online"
+SETTINGS_DEFAULT_OFFLINE = "Default Offline"
+SETTINGS_SWITCH = "Switch"
+
+SWITCH_AUTOMATIC = "Automatic"
+SWITCH_TRIGGER_OFFLINE = "Trigger Offline"
+SWITCH_MANUAL = "Manual"
 
 models = {
     'GitHub': {
@@ -15,12 +33,31 @@
         'type': 'local',
         'filename': 'mistral-7b-instruct-v0.1.Q5_K_M.gguf',
     },
+    "stable-code-3b": {
+        "url": "https://huggingface.co/stabilityai/stable-code-3b/resolve/main/stable-code-3b-Q5_K_M.gguf",
+        "type": "local",
+        "filename": "stable-code-3b-Q5_K_M.gguf",
+    },
     'CodeLlama-34b': {
         'url': 'https://huggingface.co/TheBloke/CodeLlama-34B-Instruct-GGUF/resolve/main/codellama-34b-instruct.Q4_K_M.gguf',
         'type': 'local',
         'filename': 'codellama-34b-instruct.Q4_K_M.gguf',
-    },
-    'default': 'GitHub',
+    }
 }
 
 model_folder = os.path.expanduser('~/models')
+
+def set_settings(setting, value):
+    with open('settings.json', 'r') as f:
+        settings = json.load(f)
+
+    settings[setting] = value
+
+    with open('settings.json', 'w') as f:
+        json.dump(settings, f, indent=4)
+
+def get_settings(setting):
+    with open('settings.json', 'r') as f:
+        settings = json.load(f)
+
+    return settings[setting]

proxy.py

@@ -5,10 +5,11 @@
 import logging
 from starlette import applications, responses, exceptions
 from starlette.requests import Request
+
 import config
 
 app = applications.Starlette()
-state = config.models[config.models['default']]
+state = config.models[config.get_settings(config.SETTINGS_DEFAULT_ONLINE)]
 local_server_process = None
 logging.basicConfig(level=logging.DEBUG)
 
@@ -86,4 +87,19 @@ async def server_error(request, exc):
 
 if __name__ == '__main__':
     import uvicorn
+    import psutil
+
+    # kill any existing local server on 5001 or 8000
+    for proc in psutil.process_iter():
+        try:
+            for conns in proc.connections(kind='inet'):
+                if conns.laddr.port == 5001:
+                    print(f"Killing process {proc.name()} on port 5001")
+                    proc.kill()
+                if conns.laddr.port == 8000:
+                    print(f"Killing process {proc.name()} on port 8000")
+                    proc.kill()
+        except:
+            continue
+
     uvicorn.run(app, host="0.0.0.0", port=5001)

requirements.txt

@@ -11,6 +11,7 @@ httpx==0.25.0
 idna==3.4
 llama_cpp_python==0.2.11
 numpy==1.26.0
+psutil==5.9.8
 pydantic==2.4.2
 pydantic-settings==2.0.3
 pydantic_core==2.10.1

settings.json

@@ -0,0 +1,9 @@
+{
+    "Current Model": "GitHub",
+    "Show Uninstalled": false,
+    "Show Status Icons": true,
+    "Show Current Model": true,
+    "Default Online": "GitHub",
+    "Default Offline": "CodeLlama-7b",
+    "Switch": "Automatic"
+}