{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# THOX Mesh Node - Colab template\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ttracx/thoxllm-cloud/blob/main/templates/colab/THOX_Mesh_Node.ipynb)\n",
    "\n",
    "Run a THOX model on a Colab VM and register it as **your own** node on\n",
    "ThoxLLM Cloud.\n",
    "\n",
    "**You supply your own Hugging Face token. That is the only requirement.**\n",
    "\n",
    "Your token is entered into *this* Colab runtime with `getpass` - it is never\n",
    "printed, never written to a file, and never sent to THOX. It is used only to\n",
    "download model weights from Hugging Face.\n",
    "\n",
    "> **Opening this notebook.** The badge above resolves once `ttracx/thoxllm-cloud`\n",
    "> is public. While it is private, download the notebook from the template Space\n",
    "> (`https://huggingface.co/spaces/tommytracx/thox-mesh-node/blob/main/colab/THOX_Mesh_Node.ipynb`)\n",
    "> and use **File -> Upload notebook** in Colab. Nothing else changes.\n",
    "\n",
    "---\n",
    "\n",
    "### What this notebook does\n",
    "\n",
    "1. Installs the serving runtime and the THOX registration helper.\n",
    "2. Prompts for your HF token and your ThoxLLM Cloud API key.\n",
    "3. Downloads a CPU-sized THOX model and serves it on `127.0.0.1:8000`\n",
    "   behind an OpenAI-compatible API.\n",
    "4. Opens a Cloudflare tunnel, giving this VM a public URL.\n",
    "5. Registers that URL so the gateway can route to it by alias.\n",
    "6. Heartbeats so it keeps being routed to.\n",
    "\n",
    "Run the cells top to bottom."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install\n",
    "\n",
    "`thoxmesh_node` is the thin registration helper. It is published as a wheel\n",
    "alongside the template Space on Hugging Face, so this cell needs no GitHub\n",
    "access and no credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q llama-cpp-python huggingface_hub\n",
    "%pip install -q \"thoxmesh-node @ https://huggingface.co/tommytracx/thoxmesh-node-dist/resolve/main/dist/thoxmesh_node-0.1.0-py3-none-any.whl\"\n",
    "\n",
    "import thoxmesh_node\n",
    "print('thoxmesh_node', thoxmesh_node.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Credentials\n",
    "\n",
    "`getpass` keeps these out of the notebook output and out of your saved copy.\n",
    "\n",
    "- **HF token** - https://huggingface.co/settings/tokens (`read` scope is enough).\n",
    "- **ThoxLLM Cloud URL + API key** - from the portal, *Deploy your own mesh*.\n",
    "  The key needs the `mesh:register` scope. **It is not your HF token**, and\n",
    "  registration rejects `hf_*` values.\n",
    "\n",
    "Leave the cloud fields blank to serve the model without registering it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['HF_TOKEN'] = getpass('Your Hugging Face token: ').strip()\n",
    "os.environ['THOX_CLOUD_URL'] = input('ThoxLLM Cloud URL (blank to skip): ').strip()\n",
    "if os.environ['THOX_CLOUD_URL']:\n",
    "    os.environ['THOX_API_KEY'] = getpass('ThoxLLM Cloud API key (mesh:register): ').strip()\n",
    "\n",
    "# Sanity checks only - never print either credential.\n",
    "assert os.environ['HF_TOKEN'].startswith('hf_'), 'That does not look like an HF token'\n",
    "if os.environ.get('THOX_API_KEY', '').startswith('hf_'):\n",
    "    raise SystemExit('That is your Hugging Face token, not a ThoxLLM Cloud API '\n",
    "                     'key. Your HF token never leaves this runtime.')\n",
    "print('HF token accepted (%d chars)' % len(os.environ['HF_TOKEN']))\n",
    "print('Cloud:', os.environ['THOX_CLOUD_URL'] or '(not registering)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Choose a model\n",
    "\n",
    "The default is the smallest model that still produces coherent output on a\n",
    "free CPU runtime. Swap in a larger THOX model if you have a GPU runtime.\n",
    "\n",
    "| Model | File | Size | Notes |\n",
    "|---|---|---|---|\n",
    "| `Thox-ai/ThoxMini-3B` | `thoxmini-3b-Q4_K_M.gguf` | 2.0 GB | **default** - coherent output |\n",
    "| `Thox-ai/thox-micro-125m-GGUF` | `thox-micro-125m.q4_k_m.gguf` | 85 MB | loads in seconds, but output is degenerate - plumbing tests only |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_REPO = 'Thox-ai/ThoxMini-3B'        # @param {type:'string'}\n",
    "MODEL_FILE = 'thoxmini-3b-Q4_K_M.gguf'    # @param {type:'string'}\n",
    "ALIAS      = 'thoxmini-3b'                # @param {type:'string'}\n",
    "PORT       = 8000\n",
    "CONTEXT    = 2048\n",
    "\n",
    "os.environ['THOX_MODEL_ID'] = MODEL_REPO\n",
    "os.environ['THOX_ALIAS'] = ALIAS\n",
    "\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)\n",
    "print('Downloaded to', model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Serve the model\n",
    "\n",
    "`llama_cpp.server` exposes an OpenAI-compatible API. It binds to loopback;\n",
    "the next cell puts a tunnel in front of it so the gateway can reach it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess, sys, time, urllib.request\n",
    "\n",
    "server = subprocess.Popen(\n",
    "    [sys.executable, '-m', 'llama_cpp.server',\n",
    "     '--model', model_path,\n",
    "     '--host', '127.0.0.1',\n",
    "     '--port', str(PORT),\n",
    "     '--n_ctx', str(CONTEXT),\n",
    "     '--model_alias', ALIAS],\n",
    "    stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,\n",
    ")\n",
    "\n",
    "def wait_for_server(url, timeout=300):\n",
    "    deadline = time.time() + timeout\n",
    "    while time.time() < deadline:\n",
    "        if server.poll() is not None:\n",
    "            raise RuntimeError('Model server exited:\\n' + (server.stdout.read() or ''))\n",
    "        try:\n",
    "            with urllib.request.urlopen(url, timeout=5) as response:\n",
    "                if response.status == 200:\n",
    "                    return True\n",
    "        except Exception:\n",
    "            time.sleep(2)\n",
    "    raise TimeoutError('Model server did not become ready in %ss' % timeout)\n",
    "\n",
    "wait_for_server(f'http://127.0.0.1:{PORT}/v1/models')\n",
    "print('Model server ready on 127.0.0.1:%d' % PORT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quick check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, urllib.request\n",
    "\n",
    "payload = json.dumps({\n",
    "    'model': ALIAS,\n",
    "    'messages': [{'role': 'user', 'content': 'Say hello in five words.'}],\n",
    "    'max_tokens': 32,\n",
    "}).encode()\n",
    "request = urllib.request.Request(\n",
    "    f'http://127.0.0.1:{PORT}/v1/chat/completions',\n",
    "    data=payload, headers={'content-type': 'application/json'})\n",
    "with urllib.request.urlopen(request, timeout=120) as response:\n",
    "    reply = json.load(response)\n",
    "print(reply['choices'][0]['message']['content'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Tunnel - how the gateway reaches this node\n",
    "\n",
    "A Colab VM has no inbound address, so it needs a public URL before the\n",
    "ThoxLLM Cloud gateway can proxy to it. `cloudflared` provides one.\n",
    "\n",
    "> **This URL is unauthenticated** - anyone holding it can call your model.\n",
    "> Set `THOX_UPSTREAM_AUTH=bearer` with a credential you mint if that matters.\n",
    "> It is also ephemeral: a new tunnel means re-registering.\n",
    "\n",
    "Skip only if you are registering a private address into your own mesh with\n",
    "`THOX_REGISTRAR=device_v2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RUN_TUNNEL = True  # @param {type:'boolean'}\n",
    "\n",
    "public_url = None\n",
    "if RUN_TUNNEL:\n",
    "    !wget -q -O cloudflared https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64\n",
    "    !chmod +x cloudflared\n",
    "    import re, threading\n",
    "    tunnel = subprocess.Popen(\n",
    "        ['./cloudflared', 'tunnel', '--no-autoupdate', '--url', f'http://127.0.0.1:{PORT}'],\n",
    "        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n",
    "\n",
    "    def read_tunnel_url():\n",
    "        global public_url\n",
    "        for line in tunnel.stdout:\n",
    "            match = re.search(r'https://[-a-z0-9]+\\.trycloudflare\\.com', line)\n",
    "            if match:\n",
    "                public_url = match.group(0)\n",
    "                print('Tunnel (yours, not the mesh address):', public_url)\n",
    "                return\n",
    "\n",
    "    threading.Thread(target=read_tunnel_url, daemon=True).start()\n",
    "    time.sleep(15)\n",
    "\n",
    "os.environ['THOX_PUBLIC_URL'] = public_url or ''\n",
    "print('Public URL:', public_url or '(none)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Register your node\n",
    "\n",
    "Two paths, wanting **opposite** addresses:\n",
    "\n",
    "| Registrar | Advertises | Use when |\n",
    "|---|---|---|\n",
    "| `spine` *(default)* | the public tunnel URL | Normal BYO. The gateway proxies to you. |\n",
    "| `device_v2` | a private address | Node-to-node routing in your own mesh. |\n",
    "\n",
    "| Attach mode | Advertises | Valid for |\n",
    "|---|---|---|\n",
    "| `public` | the tunnel URL | `spine` |\n",
    "| `mesh` | assigned `100.96.x.x` | `device_v2`, WireGuard up |\n",
    "| `lan` / `loopback` | private address | `device_v2` |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from thoxmesh_node import AttachError, NodeAgent, NodeConfig\n",
    "\n",
    "REGISTRAR   = 'spine'   # @param ['spine', 'device_v2']\n",
    "ATTACH_MODE = 'public'  # @param ['public', 'loopback', 'mesh', 'lan']\n",
    "\n",
    "agent = None\n",
    "if os.environ.get('THOX_CLOUD_URL') and os.environ.get('THOX_API_KEY'):\n",
    "    config = NodeConfig.from_env(\n",
    "        alias=ALIAS,\n",
    "        port=PORT,\n",
    "        context_window=CONTEXT,\n",
    "        registrar=REGISTRAR,\n",
    "        attach_mode=ATTACH_MODE,\n",
    "        display_name='THOX Colab Node',\n",
    "        product_type='other',\n",
    "        platform='linux',\n",
    "        role='compute',\n",
    "        capabilities={'chat': True, 'completions': True},\n",
    "        state_dir=__import__('pathlib').Path('/content/.thoxmesh'),\n",
    "        app_version='colab',\n",
    "    )\n",
    "    agent = NodeAgent(config)\n",
    "    try:\n",
    "        endpoint_id = agent.join()\n",
    "        print('Registered via', REGISTRAR)\n",
    "        print('  endpoint', endpoint_id)\n",
    "        print('  address ', agent.attachment.base_url)\n",
    "        if agent.device is not None:\n",
    "            print('  device  ', agent.device.device_id)\n",
    "    except AttachError as error:\n",
    "        print('Could not advertise an address:\\n', error)\n",
    "else:\n",
    "    print('Not registering - serving locally only.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Verify\n",
    "\n",
    "On the `device_v2` path this asks the control plane to resolve your model\n",
    "and confirms it selects this node - the same path `ThoxRoute` uses. On the\n",
    "`spine` path it prints how to call your alias through the gateway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if agent is not None and agent.endpoint_id and agent.client and agent.device:\n",
    "    peers = agent.client.list_peers(agent.device.device_id)\n",
    "    print('Peers in mesh:', len(peers['peers']))\n",
    "    for peer in peers['peers']:\n",
    "        marker = ' <- this node' if peer['id'] == agent.device.device_id else ''\n",
    "        print('  %-38s %-8s %s%s' % (peer['id'], peer['status'], peer.get('mesh_ip'), marker))\n",
    "\n",
    "    route = agent.client.device_action(\n",
    "        agent.device.device_id, 'resolve_route', {'model_id': os.environ['THOX_MODEL_ID']})\n",
    "    endpoint = route['route']['endpoint']\n",
    "    print('\\nresolve_route selected:', endpoint['id'], 'at', endpoint['base_url'])\n",
    "    assert endpoint['id'] == agent.endpoint_id, 'Mesh routed to a different node'\n",
    "    print('OK - this node is registered and routable.')\n",
    "elif agent is not None and agent.endpoint_id:\n",
    "    print('Registered with the gateway. Address it by alias:')\n",
    "    print('  POST %s/v1/chat/completions' % os.environ['THOX_CLOUD_URL'])\n",
    "    print('  authorization: Bearer <your API key>')\n",
    "    print('  body: {\"model\": \"%s\", \"messages\": [...]}' % ALIAS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Keep the node alive\n",
    "\n",
    "The agent heartbeats on a background thread. Leave this cell running to keep\n",
    "the Colab session (and the node) up. Interrupt it to stop.\n",
    "\n",
    "When the session ends the node stops heartbeating and the mesh stops\n",
    "routing to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "try:\n",
    "    while True:\n",
    "        time.sleep(60)\n",
    "        if agent is not None:\n",
    "            print(time.strftime('%H:%M:%S'), 'heartbeat ok')\n",
    "except KeyboardInterrupt:\n",
    "    print('Stopping...')\n",
    "finally:\n",
    "    if agent is not None:\n",
    "        agent.stop()\n",
    "        print('Node unregistered.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Teardown\n",
    "\n",
    "Run this to shut down cleanly without ending the runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if agent is not None:\n",
    "    agent.stop()\n",
    "try:\n",
    "    server.terminate()\n",
    "except Exception:\n",
    "    pass\n",
    "print('Stopped.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Troubleshooting\n",
    "\n",
    "| Symptom | Cause | Fix |\n",
    "|---|---|---|\n",
    "| `invalid_or_expired_pairing_token` | Join keys are single-use and time-limited. | Issue a new one in the portal. |\n",
    "| `identity_public_key_already_registered` | This VM already paired; `/content/.thoxmesh` was kept. | Reuse it, or delete the folder to pair as a new device. |\n",
    "| `Model endpoint must use a private...` | Registered a public URL via `device_v2`. | Use `THOX_REGISTRAR=spine`, or a private attach mode. |\n",
    "| `gateway cannot reach` | Registered a private address via `spine`. | Start the tunnel, use `THOX_ATTACH_MODE=public`. |\n",
    "| `looks like a Hugging Face token` | HF token pasted as the API key. | Use a portal API key with scope `mesh:register`. |\n",
    "| `request_timestamp_out_of_window` | Runtime clock drifted >5 min. | Restart the runtime. |\n",
    "| `replay_detected` | A nonce was reused. | Do not re-send a captured request; the agent handles this. |\n",
    "\n",
    "Security model: `templates/SECURITY.md`."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true,
   "name": "THOX_Mesh_Node.ipynb"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}