diff --git a/app.py b/app.py index 9d06a3a..9fb9335 100644 --- a/app.py +++ b/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']) diff --git a/config.py b/config.py index 9b11232..7c154d0 100644 --- a/config.py +++ b/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] diff --git a/proxy.py b/proxy.py index 09d488c..5d16843 100644 --- a/proxy.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index 980b630..ea55803 100644 --- a/requirements.txt +++ b/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 diff --git a/icon.png b/resources/icon.png similarity index 100% rename from icon.png rename to resources/icon.png diff --git a/resources/installed.png b/resources/installed.png new file mode 100644 index 0000000..7deacc8 Binary files /dev/null and b/resources/installed.png differ diff --git a/resources/installing.png b/resources/installing.png new file mode 100644 index 0000000..37b5a2a Binary files /dev/null and b/resources/installing.png differ diff --git a/resources/uninstalled.png b/resources/uninstalled.png new file mode 100644 index 0000000..87c8ad5 Binary files /dev/null and b/resources/uninstalled.png differ diff --git a/settings.json b/settings.json new file mode 100644 index 0000000..114b19b --- /dev/null +++ b/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" +} \ No newline at end of file