faggot-chan/colab_setup.py

#!/usr/bin/env python3
"""
run this in a colab cell to stand up main.py + a public wss:// tunnel to it.
session is ephemeral — restart the notebook, get a new url, re-paste it
into your bot.py-side files/config/connection.txt (or rerun setup_wizard.py).
"""
import os
import re
import secrets
import subprocess
import time
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_DIR = SCRIPT_DIR / "config"

def model_from_main(name):
    # read the model names straight out of main.py so they can't drift apart
    text = (SCRIPT_DIR / "main.py").read_text(encoding="utf-8")
    return re.search(rf'^{name}\s*=\s*["\'](.+?)["\']', text, re.M).group(1)

CHATMODEL = model_from_main("chatmodel")
EMBEDMODEL = model_from_main("embedmodel")


def run(cmd, **kwargs):
    print(f"$ {cmd}")
    return subprocess.run(cmd, shell=True, check=True, **kwargs)


def install_ollama():
    print("\n-- installing ollama --")
    run("curl -fsSL https://ollama.com/install.sh | sh")
    subprocess.Popen("ollama serve", shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(5)  # give the server a sec to come up before we start pulling models


def pull_models():
    print("\n-- pulling models (this'll take a while) --")
    run(f"ollama pull {CHATMODEL}")
    run(f"ollama pull {EMBEDMODEL}")


def install_python_deps():
    print("\n-- installing python deps --")
    run("pip install -q ollama websockets")


def ensure_auth_token() -> str:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    token_file = CONFIG_DIR / "auth_token.txt"
    if token_file.exists():
        token = token_file.read_text(encoding="utf-8").strip()
        print(f"\nusing existing auth token from {token_file}")
    else:
        token = secrets.token_urlsafe(32)
        token_file.write_text(f"{token}\n", encoding="utf-8")
        print(f"\ngenerated new auth token, wrote {token_file}")
    return token


def ensure_params():
    # main.py is stateless — system_prompt comes over the wire in every payload from bot.py.
    # params.json is the one thing it actually reads off its own disk (ollama.chat options
    # aren't sent in the payload), so that's the only file that has to exist here.
    params_file = CONFIG_DIR / "params.json"
    if not params_file.exists():
        raise SystemExit(f"missing params.json in {CONFIG_DIR} — upload it before running this")


def start_main():
    print("\n-- starting main.py --")
    # same deal as the tunnel below, main.py prints a lot and raises BrokenPipe once this cell's stdout is gone
    log_file = open(SCRIPT_DIR / "main.log", "w")
    return subprocess.Popen(
        f"python3 {SCRIPT_DIR / 'main.py'}",
        shell=True,
        stdout=log_file,
        stderr=subprocess.STDOUT,
    )


def start_tunnel() -> str:
    print("\n-- installing cloudflared --")
    run("curl -fsSL -o cloudflared.deb https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64.deb")
    run("dpkg -i cloudflared.deb")

    print("\n-- opening tunnel to localhost:8765 --")
    # logging to a real file instead of a PIPE — a pipe's read end lives in this python
    # process, so once this script exits the pipe's reader disappears and cloudflared
    # gets SIGPIPE on its next log write and dies. a file has no reader to lose.
    log_path = SCRIPT_DIR / "cloudflared.log"
    log_file = open(log_path, "w")
    subprocess.Popen(
        "cloudflared tunnel --url http://localhost:8765",
        shell=True,
        stdout=log_file,
        stderr=subprocess.STDOUT,
    )

    url = None
    deadline = time.time() + 30
    while time.time() < deadline:
        text = log_path.read_text(encoding="utf-8", errors="ignore")
        match = re.search(r"https://[a-zA-Z0-9\-]+\.trycloudflare\.com", text)
        if match:
            url = match.group(0)
            break
        time.sleep(1)

    if not url:
        raise SystemExit(f"couldn't find the tunnel url, check {log_path}")

    return url.replace("https://", "wss://")


def main():
    install_ollama()
    pull_models()
    install_python_deps()
    ensure_params()
    token = ensure_auth_token()
    start_main()
    time.sleep(3)  # let main.py bind its port before the tunnel tries to reach it
    wss_url = start_tunnel()

    print("\n" + "=" * 60)
    print("paste these into setup_wizard.py on the bot.py side:")
    print(f"  address: {wss_url}")
    print(f"  token:   {token}")
    print("=" * 60)


if __name__ == "__main__":
    main()