server.py

"""Flask + SSE server for ColorGPT.

  /              → UI
  /api/speak     → SSE: speaker mode
  /api/read      → SSE: reader mode
  /api/dialogue  → SSE: dialogue mode

Run:  python server.py   →  http://127.0.0.1:5000
"""

import io
import json
from pathlib import Path

from flask import Flask, Response, render_template, request, send_file, stream_with_context

from histogram import render_palette, render_reading_map, render_transcript, render_vocab_map
from pacing import pace
from streams import dialogue, reader, speaker

app = Flask(__name__, template_folder="templates")
app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024  # 50 MB upload cap
app.config["TEMPLATES_AUTO_RELOAD"] = True  # always re-read template on each request

CORPUS_DIR = Path(__file__).parent / "corpus"
CORPUS_DIR.mkdir(exist_ok=True)
ALLOWED_CORPUS_EXT = {".txt", ".md"}


def _safe_corpus_path(name: str):
    """Resolve `name` strictly inside CORPUS_DIR. Returns Path or None."""
    if not name or "/" in name or "\\" in name or name.startswith("."):
        return None
    target = (CORPUS_DIR / name).resolve()
    try:
        target.relative_to(CORPUS_DIR.resolve())
    except ValueError:
        return None
    return target if target.exists() and target.is_file() else None


def _sse(events):
    for ev in events:
        yield f"data: {json.dumps(ev)}\n\n"


@app.route("/")
def index():
    return render_template("index.html")


SSE_HEADERS = {"X-Accel-Buffering": "no", "Cache-Control": "no-cache"}


def _stream(events, paced: bool, base_tps: float):
    if paced:
        events = pace(events, base_tps=base_tps)
    return Response(
        stream_with_context(_sse(events)),
        mimetype="text/event-stream",
        headers=SSE_HEADERS,
    )


def _pacing_args():
    paced = request.args.get("pace", "1") != "0"
    base_tps = float(request.args.get("tps", 2.0))
    return paced, base_tps


@app.route("/api/speak")
def api_speak():
    prompt = request.args.get("prompt", "")
    max_new = int(request.args.get("max", 500))
    temperature = float(request.args.get("temperature", 0.95))
    paced, base_tps = _pacing_args()
    return _stream(speaker(prompt, max_new_tokens=max_new, temperature=temperature), paced, base_tps)


@app.route("/api/corpus", methods=["GET"])
def api_corpus_list():
    files = sorted(
        f.name for f in CORPUS_DIR.iterdir()
        if f.is_file() and f.suffix.lower() in ALLOWED_CORPUS_EXT
    )
    return {"files": files}


@app.route("/api/corpus", methods=["POST"])
def api_corpus_upload():
    if "file" not in request.files:
        return {"error": "no file"}, 400
    f = request.files["file"]
    name = Path(f.filename or "").name  # strip path components
    if not name or Path(name).suffix.lower() not in ALLOWED_CORPUS_EXT:
        return {"error": "must be .txt or .md"}, 400
    f.save(CORPUS_DIR / name)
    return {"saved": name}


def _render_corpus_image(name: str, kind: str):
    target = _safe_corpus_path(name)
    if target is None:
        return {"error": f"corpus '{name}' not found"}, 404
    text = target.read_text(encoding="utf-8")
    img = render_transcript(text) if kind == "transcript" else render_palette(text)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return send_file(buf, mimetype="image/png", download_name=f"{name}-{kind}.png")


@app.route("/api/corpus/<name>/transcript.png")
def api_corpus_transcript(name):
    return _render_corpus_image(name, "transcript")


@app.route("/api/corpus/<name>/palette.png")
def api_corpus_palette(name):
    return _render_corpus_image(name, "palette")


@app.route("/api/corpus/<name>/reading-map.png")
def api_corpus_reading_map(name):
    target = _safe_corpus_path(name)
    if target is None:
        return {"error": f"corpus '{name}' not found"}, 404
    text = target.read_text(encoding="utf-8")
    cell_mm = request.args.get("cell_mm", type=float)
    dpi = int(request.args.get("dpi", 300))
    img = render_reading_map(text, title=name, cell_mm=cell_mm, dpi=dpi)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return send_file(buf, mimetype="image/png", download_name=f"{name}-reading-map.png")


@app.route("/api/lut")
def api_lut():
    """Raw LUT bytes for the playground canvas. ~445KB, cacheable."""
    from engine import LUT_PATH
    resp = send_file(LUT_PATH, mimetype="application/octet-stream")
    resp.headers["Cache-Control"] = "public, max-age=86400"
    return resp


def _collect_corpus_token_ids() -> set:
    """Union of token IDs that appear in any corpus file."""
    from engine import load_tokenizer
    tokenizer = load_tokenizer()
    used = set()
    for path in CORPUS_DIR.iterdir():
        if not path.is_file() or path.suffix.lower() not in ALLOWED_CORPUS_EXT:
            continue
        try:
            text = path.read_text(encoding="utf-8")
        except Exception:
            continue
        used.update(tokenizer.encode(text, add_special_tokens=False))
    return used


@app.route("/api/filter/<kind>.bin")
def api_filter(kind):
    """Return a bitmap of which tokens pass the named filter.
    1 bit per token, packed LSB-first per byte. ~19KB for 151936 tokens.

      latin  — tokens whose decoded text is pure ASCII (≤ U+007F)
      corpus — tokens that appear in any file in corpus/
    """
    from engine import get_non_latin_token_ids, load_lut
    _, meta = load_lut()
    vocab_size = meta["vocab_size"]

    if kind == "latin":
        bad = {item[0] for item in get_non_latin_token_ids()}
        passing = (tid for tid in range(vocab_size) if tid not in bad)
    elif kind == "corpus":
        from engine import load_tokenizer
        corpus_name = request.args.get("corpus", "").strip()
        if corpus_name:
            target = _safe_corpus_path(corpus_name)
            if target is None:
                return {"error": f"corpus '{corpus_name}' not found"}, 404
            text = target.read_text(encoding="utf-8")
            ids = set(load_tokenizer().encode(text, add_special_tokens=False))
            passing = iter(ids)
        else:
            passing = iter(_collect_corpus_token_ids())
    else:
        return {"error": f"unknown filter '{kind}' (use latin or corpus)"}, 400

    n_bytes = (vocab_size + 7) // 8
    bitmap = bytearray(n_bytes)
    for tid in passing:
        if 0 <= tid < vocab_size:
            bitmap[tid >> 3] |= 1 << (tid & 7)
    return Response(bytes(bitmap), mimetype="application/octet-stream")


@app.route("/api/vocab.png")
def api_vocab_map():
    sort = request.args.get("sort", "id")
    cell = int(request.args.get("cell", 6))
    if sort not in ("id", "hue"):
        return {"error": "sort must be 'id' or 'hue'"}, 400
    img = render_vocab_map(cell=cell, sort=sort)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return send_file(buf, mimetype="image/png", download_name=f"vocab-{sort}.png")


@app.route("/api/read")
def api_read():
    text = request.args.get("text", "")
    corpus = request.args.get("corpus", "")
    paced, base_tps = _pacing_args()
    if corpus:
        target = _safe_corpus_path(corpus)
        if target is None:
            return {"error": f"corpus '{corpus}' not found"}, 404
        events = reader(path=target)
    else:
        events = reader(text=text)
    return _stream(events, paced, base_tps)


@app.route("/api/dialogue")
def api_dialogue():
    seed = request.args.get("seed", "In the silence")
    tokens_per_turn = int(request.args.get("turn", 20))
    total = int(request.args.get("total", 2000))
    paced, base_tps = _pacing_args()
    return _stream(
        dialogue(seed=seed, tokens_per_turn=tokens_per_turn, total_tokens=total), paced, base_tps
    )


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000, threaded=True, debug=False)