diff --git a/app.py b/app.py index 9d06a3a..69938ec 100644 --- a/app.py +++ b/app.py @@ -1,9 +1,9 @@ +from argparse import ArgumentParser import rumps import requests import threading import subprocess import os -import sys import config @@ -28,9 +28,12 @@ def setup(): class ModelPickerApp(rumps.App): - def __init__(self): + def __init__(self, proxy_port=5001, local_server_port=8000): super(ModelPickerApp, self).__init__("ModelPickerApp") + self.proxy_port = proxy_port + self.local_server_port = local_server_port + # Dynamically create menu items from the MENUBAR_OPTIONS self.menu_items = {} for option in config.models: @@ -52,7 +55,8 @@ def pick_model(self, sender): choice = sender.title try: response = requests.post( - "http://localhost:5001/set_target", json={"target": choice}) + f"http://localhost:{self.proxy_port}/set_target", + json={"target": choice}) if response.status_code == 200: print(f"Successfully sent selection: {choice}.") else: @@ -69,13 +73,25 @@ def pick_model(self, sender): self.menu_items[item].state = False def run_server(self): - subprocess.run(['python', 'proxy.py']) + subprocess.run([ + 'python', 'proxy.py', '--port', str(self.proxy_port), + '--local-server-port', str(self.local_server_port), + ]) if __name__ == '__main__': - if '--setup' in sys.argv: + parser = ArgumentParser() + parser.add_argument('--proxy-port', type=int, default=5001) + parser.add_argument('--local-server-port', type=int, default=8000) + parser.add_argument('--setup', action='store_true') + args = parser.parse_args() + + if args.setup: setup() - app = ModelPickerApp() + app = ModelPickerApp( + proxy_port=args.proxy_port, + local_server_port=args.local_server_port, + ) print("Running server...") server_thread = threading.Thread(target=app.run_server) server_thread.start() diff --git a/proxy.py b/proxy.py index 09d488c..a25b9fe 100644 --- a/proxy.py +++ b/proxy.py @@ -1,3 +1,4 @@ +from argparse import ArgumentParser import config import httpx import os @@ -9,6 +10,7 @@ app = applications.Starlette() state = config.models[config.models['default']] +local_server_port = 8000 local_server_process = None logging.basicConfig(level=logging.DEBUG) @@ -19,7 +21,8 @@ def start_local_server(model_filename): local_server_process.terminate() local_server_process.wait() cmd = ["python3", "-m", "llama_cpp.server", "--model", model_filename, - "--n_gpu_layers", "1", "--n_ctx", "4096"] # TODO: set this more correctly + "--n_gpu_layers", "1", "--n_ctx", "4096", # TODO: set this more correctly + "--port", str(local_server_port)] logging.debug('Running: %s' % ' '.join(cmd)) local_server_process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) @@ -52,7 +55,7 @@ async def proxy(request: Request): if state['type'] == 'remote': url = f"{state['domain']}{path}" elif state['type'] == 'local': - url = f"http://localhost:8000{path}" + url = f"http://localhost:{local_server_port}{path}" data = await request.body() headers = dict(request.headers) @@ -86,4 +89,11 @@ async def server_error(request, exc): if __name__ == '__main__': import uvicorn - uvicorn.run(app, host="0.0.0.0", port=5001) + + parser = ArgumentParser() + parser.add_argument('--port', type=int, default=5001) + parser.add_argument('--local-server-port', type=int, default=8000) + args = parser.parse_args() + + local_server_port = args.local_server_port + uvicorn.run(app, host="0.0.0.0", port=args.port)