import rumps
import requests
import threading
import subprocess
import os
import sys

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}.")


class ModelPickerApp(rumps.App):
    def __init__(self):
        super(ModelPickerApp, self).__init__("ModelPickerApp")

        # Dynamically create menu items from the MENUBAR_OPTIONS
        self.menu_items = {}
        for option in config.models:
            if option == 'default':
                continue
            self.menu_items[option] = rumps.MenuItem(
                title=option, callback=self.pick_model)

        self.menu = list(self.menu_items.values())
        self.menu_items[config.models['default']].state = True
        self.icon = "icon.png"

    def pick_model(self, sender):
        # Toggle the checked status of the clicked menu item
        sender.state = not sender.state

        # Send the choice to the local proxy app
        if sender.state:
            choice = sender.title
            try:
                response = requests.post(
                    "http://localhost:5001/set_target", json={"target": choice})
                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}.")

        # 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

    def run_server(self):
        subprocess.run(['python', 'proxy.py'])


if __name__ == '__main__':
    if '--setup' in sys.argv:
        setup()
    app = ModelPickerApp()
    print("Running server...")
    server_thread = threading.Thread(target=app.run_server)
    server_thread.start()
    app.run()

app.py
danielgross/localpilot
import os

models = {
    'GitHub': {
        'domain': 'https://copilot-proxy.githubusercontent.com',
        'type': 'remote',
    },
    'CodeLlama-7b': {
        'url': 'https://huggingface.co/TheBloke/CodeLlama-7B-GGUF/resolve/main/codellama-7b.Q5_K_S.gguf',
        'type': 'local',
        'filename': 'codellama-7b.Q5_K_S.gguf',
    },
    'Mistral-7b': {
        'url': 'https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q5_K_M.gguf',
        'type': 'local',
        'filename': 'mistral-7b-instruct-v0.1.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')

config.py
danielgross/localpilot
import config
import httpx
import os
import subprocess
import logging
from starlette import applications, responses, exceptions
from starlette.requests import Request
import config

app = applications.Starlette()
state = config.models[config.models['default']]
local_server_process = None
logging.basicConfig(level=logging.DEBUG)


def start_local_server(model_filename):
    global local_server_process
    if local_server_process:
        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
    logging.debug('Running: %s' % ' '.join(cmd))
    local_server_process = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)


@app.route('/set_target', methods=['POST'])
async def set_target(request: Request):
    global state
    response = await request.json()
    target = response['target']
    if target not in config.models:
        raise exceptions.HTTPException(
            status_code=400, detail=f'Invalid target: {target}')

    state = config.models[target]
    if config.models[target].get("type") == "local":
        start_local_server(os.path.join(
            config.model_folder, config.models[target]['filename']))

    message = f'Target set to {state}'
    return responses.JSONResponse({'message': message}, status_code=200)


@app.route('/{path:path}', methods=['GET', 'POST', 'PUT', 'DELETE'])
async def proxy(request: Request):
    global state
    path = request.url.path
    logging.debug(f'Current state: {state}')

    if state['type'] == 'remote':
        url = f"{state['domain']}{path}"
    elif state['type'] == 'local':
        url = f"http://localhost:8000{path}"

    data = await request.body()
    headers = dict(request.headers)
    r = None
    async with httpx.AsyncClient() as client:
        try:
            if request.method == 'GET':
                r = await client.get(url, params=request.query_params, headers=headers)
            elif request.method == 'POST':
                r = await client.post(url, data=data, headers=headers, timeout=30)
            elif request.method == 'PUT':
                r = await client.put(url, data=data, headers=headers)
            elif request.method == 'DELETE':
                r = await client.delete(url, headers=headers)
        except httpx.RemoteProtocolError as exc:
            logging.debug(f'Connection closed prematurely: {exc}')
    content = r.content if r else ''
    status_code = r.status_code if r else 204
    headers = dict(r.headers) if r else dict()
    return responses.Response(content=content, status_code=status_code, headers=headers)


@app.exception_handler(404)
async def not_found(request, exc):
    return responses.JSONResponse({"error": "Not found"}, status_code=404)


@app.exception_handler(500)
async def server_error(request, exc):
    return responses.JSONResponse({"error": "Server error"}, status_code=500)

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5001)

proxy.py
danielgross/localpilot