aboutsummaryrefslogtreecommitdiff
path: root/colab_setup.py
blob: b99068cf10caa9235da5f2700e8c2be913eb8eb3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#!/usr/bin/env python3
#all of this is vibecoded but i read it and it's okay 
"""
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 config/connection.txt.
"""

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"

CHATMODEL = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M"
EMBEDMODEL = "hf.co/CompendiumLabs/bge-base-en-v1.5-gguf"


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


def install_ollama():
    print("\n-- installing ollama --")
    run("apt install zstd | sh") #hi david here, colab doesn't come with zstd and since this script is meant for colab only it's okay to put apt 
    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_and_prompt():
    # main.py needs both of these to exist — if this is a fresh colab instance
    # and you haven't uploaded them yet, stop here and add them first.
    params_file = CONFIG_DIR / "params.json"
    prompt_file = CONFIG_DIR / "system_prompt.txt"
    missing = [f.name for f in (params_file, prompt_file) if not f.exists()]
    if missing:
        raise SystemExit(
            f"missing {', '.join(missing)} in {CONFIG_DIR} — upload them before running this"
        )


def start_main():
    print("\n-- starting main.py --")
    return subprocess.Popen(
        f"python3 {SCRIPT_DIR / 'main.py'}",
        shell=True,
    )


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 --")
    proc = subprocess.Popen(
        "cloudflared tunnel --url http://localhost:8765",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    url = None
    deadline = time.time() + 30
    for line in proc.stdout:
        print(line, end="")
        match = re.search(r"https://[a-zA-Z0-9\-]+\.trycloudflare\.com", line)
        if match:
            url = match.group(0)
            break
        if time.time() > deadline:
            break

    if not url:
        raise SystemExit("couldn't find the tunnel url in cloudflared's output, check the log above")

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


def main():
    install_ollama()
    pull_models()
    install_python_deps()
    ensure_params_and_prompt()
    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) #wtf is this shit claude
    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()