server.py

"""
AAC Kids next-word backend.

POST /predict       {sentence:[...words], n:8|30, required:"book"|null}
                    -> {words: [...], done_prob: 0..1, debug: {...}}
POST /pick          {sentence, word, offered, mode}     -> {ok: true}
POST /add_phrase    {phrase:"i love octopus"}           -> {ok, tokens}
GET  /dictionary    ?q=oct&limit=50                     -> {words: [...]}
GET  /health        -> {ok, model}

Pipeline:
  1) Walk phrases.json (a kid-natural phrase trie) token-by-token from the
     sentence-so-far. Keys at the resting node are the next-word candidates.
  2) If the prefix dead-ends, top up from the trie's top-level openers.
  3) For "more words" (n > 8) supplement from the LLM only when the trie runs
     short. The 8-button grid almost never touches the LLM.
  4) Sort alphabetically (case-insensitive), dedupe, drop denylist + the
     just-tapped word + non-alphabetic junk. Keep `required` (game mode) first.
"""

from __future__ import annotations

import json
import logging
import os
import re
import threading
import time
import urllib.error
import urllib.request
from typing import Optional

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from merge_phrases import clean as _clean_trie, merge as _merge_trie

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("aac")

HERE = os.path.dirname(os.path.abspath(__file__))


def _load_env():
    path = os.path.join(HERE, ".env")
    if not os.path.exists(path):
        return
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip())


_load_env()

GROQ_API_KEY = os.environ.get("OPENROUTER_KEY") or os.environ.get("GROQ_API_KEY", "")
GROQ_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL_NAME = "qwen/qwen-2.5-7b-instruct"

DENY: set[str] = {
    "sexual", "sex", "porn", "drug", "drugs",
    "hate", "shit", "fuck", "bitch", "ass",
    "gun", "weapon", "naked",
}

FALLBACK = ["i", "you", "want", "more", "please", "yes", "no", "help"]

SYSTEM_PROMPT = (
    "You help a young child build sentences one tap at a time on an AAC "
    "speech device. The child shows you the partial sentence so far and you "
    "give a comma-separated list of candidate NEXT words.\n"
    "\n"
    "Hard rules:\n"
    "1. Each word must fit grammatically as the very next word.\n"
    "2. Output is a single line: comma-separated lowercase words. No "
    "numbers, quotes, punctuation other than commas, or explanation. "
    "Single words only — no multi-word phrases.\n"
    "3. Words must all be DIFFERENT and none can equal the last word in "
    "the partial sentence.\n"
    "4. Pick common kid-safe words. Avoid scary, adult, or rude words.\n"
    "5. Vary word types — mix nouns, verbs, function words.\n"
)


# ---------------------------------------------------------------------------
# Trie state — phrases.json is the source of truth for the regular grid.
# Loaded once at boot, mutated in-place when /add_phrase is called.
# ---------------------------------------------------------------------------
_PHRASES_PATH = os.path.join(HERE, "phrases.json")
with open(_PHRASES_PATH) as _f:
    _PHRASES: dict = json.load(_f)
_PHRASES_LOCK = threading.Lock()


# ---------------------------------------------------------------------------
# LLM response cache — persists `_ask_llm` outputs across restarts so we
# don't pay for the same (sentence, n) prompt twice. File is `llm_cache.json`,
# a flat dict { "<sentence>|<n>": ["word", ...] }.
# ---------------------------------------------------------------------------
_LLM_CACHE_PATH = os.path.join(HERE, "llm_cache.json")
_LLM_CACHE: dict[str, list[str]] = {}
_LLM_CACHE_LOCK = threading.Lock()
try:
    with open(_LLM_CACHE_PATH) as _f:
        _loaded = json.load(_f)
        if isinstance(_loaded, dict):
            _LLM_CACHE = {
                k: [w for w in v if isinstance(w, str)]
                for k, v in _loaded.items()
                if isinstance(v, list)
            }
except (FileNotFoundError, json.JSONDecodeError):
    _LLM_CACHE = {}


def _llm_cache_key(sentence: str, n: int) -> str:
    return f"{sentence.strip().lower()}|{int(n)}"


def _llm_cache_get(sentence: str, n: int) -> Optional[list[str]]:
    with _LLM_CACHE_LOCK:
        hit = _LLM_CACHE.get(_llm_cache_key(sentence, n))
        return list(hit) if hit is not None else None


def _llm_cache_put(sentence: str, n: int, words: list[str]) -> None:
    if not words:
        return  # don't cache empty results — leaves room for a retry later.
    key = _llm_cache_key(sentence, n)
    with _LLM_CACHE_LOCK:
        _LLM_CACHE[key] = list(words)
        try:
            tmp = _LLM_CACHE_PATH + ".tmp"
            with open(tmp, "w") as f:
                json.dump(_LLM_CACHE, f, ensure_ascii=False, indent=1)
            os.replace(tmp, _LLM_CACHE_PATH)
        except OSError as e:
            log.warning("failed to persist llm_cache.json: %s", e)


def _walk_trie(sentence: list[str]) -> dict:
    """Walk phrases.json token-by-token. Returns the resting node (a dict
    whose keys are the candidate next words), or {} if the prefix is unknown."""
    node: dict = _PHRASES
    for raw in sentence:
        tok = (raw or "").strip().lower()
        if not tok:
            continue
        nxt = node.get(tok)
        if not isinstance(nxt, dict):
            return {}
        node = nxt
    return node


def _leaf_count(node) -> int:
    """How many phrases live under this trie node. Empty dict = 1 leaf (a
    completed phrase). Used to rank candidates by popularity."""
    if not isinstance(node, dict) or not node:
        return 1
    return sum(_leaf_count(v) for v in node.values())


# ---------------------------------------------------------------------------
# Common-dictionary lookup. dictionary.json is a flat list of kid-safe words.
# ---------------------------------------------------------------------------
_DICT_PATH = os.path.join(HERE, "dictionary.json")
try:
    with open(_DICT_PATH) as _df:
        _raw_dict = json.load(_df)
    _DICT: list[str] = sorted({
        str(w).strip().lower() for w in _raw_dict
        if w and isinstance(w, str)
    })
except (OSError, json.JSONDecodeError) as _e:
    log.warning("dictionary.json missing or invalid: %s", _e)
    _DICT = []


# ---------------------------------------------------------------------------
# JSONL log of every /predict + /pick + /add_phrase. Used for offline analysis.
# ---------------------------------------------------------------------------
_LOG_PATH = os.path.join(HERE, "predictions.jsonl")


def _jsonl(record: dict) -> None:
    try:
        with open(_LOG_PATH, "a") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
    except OSError as e:
        log.warning("failed to write jsonl log: %s", e)


# ---------------------------------------------------------------------------
# FastAPI
# ---------------------------------------------------------------------------
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class PredictReq(BaseModel):
    sentence: list[str] = []
    n: int = 8
    required: Optional[str] = None


class PickReq(BaseModel):
    sentence: list[str] = []
    word: str
    offered: list[str] = []
    mode: Optional[str] = None


class AddPhraseReq(BaseModel):
    phrase: str


@app.get("/health")
def health():
    return {
        "ok": bool(GROQ_API_KEY) or bool(_PHRASES),
        "loading": False,
        "error": None,
        "model": MODEL_NAME,
        "trie_openers": len(_PHRASES),
        "dict_size": len(_DICT),
        "llm_cache_size": len(_LLM_CACHE),
    }


# ---------------------------------------------------------------------------
# LLM helper — used only when n > 8 ("more words") AND the trie ran short.
# ---------------------------------------------------------------------------
def _ask_llm(sentence: str, n: int) -> list[str]:
    if not GROQ_API_KEY:
        return []

    cached = _llm_cache_get(sentence, n)
    if cached is not None:
        log.info("llm cache hit: sentence=%r n=%d (%d words)", sentence, n, len(cached))
        return list(cached)

    if sentence:
        user_msg = (
            f"Partial sentence: \"{sentence}\"\n"
            f"List exactly {n} different lowercase next-word options that "
            f"would each grammatically and naturally follow as the very "
            f"next word. Pick varied, kid-friendly, common everyday "
            f"single words. Comma-separated, no other text."
        )
    else:
        user_msg = (
            f"The child has not started a sentence yet. List exactly {n} "
            f"different lowercase first-word options to start a "
            f"kid-friendly sentence. Comma-separated, no other text."
        )

    body = json.dumps({
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
        ],
        "max_tokens": 200,
        "temperature": 0.7,
        "top_p": 0.9,
    }).encode()

    req = urllib.request.Request(GROQ_URL, data=body, headers={
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json",
        "User-Agent": "aac-kids/1.0",
        "HTTP-Referer": "http://localhost",
        "X-Title": "aac-kids",
    })

    try:
        with urllib.request.urlopen(req, timeout=10) as r:
            data = json.loads(r.read())
        text = data["choices"][0]["message"]["content"]
    except (urllib.error.URLError, urllib.error.HTTPError, KeyError, TimeoutError) as e:
        log.warning("llm call failed: %s", e)
        return []

    out: list[str] = []
    for chunk in text.replace("\n", ",").split(","):
        w = chunk.strip().lower().strip("\"'.!?-_:;()[] ")
        if not w:
            continue
        w = w.split()[0]
        out.append(w)
    _llm_cache_put(sentence, n, out)
    return out


def _ok_word(w: str) -> bool:
    return bool(w) and w.isalpha() and 1 <= len(w) <= 14 and w not in DENY


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------
@app.post("/predict")
def predict(req: PredictReq):
    n = max(1, min(req.n or 8, 30))
    t0 = time.time()
    out: list[str] = []
    seen: set[str] = set()

    # 1. required word (game mode) goes first, never sorted.
    required = (req.required or "").strip().lower()
    if required and _ok_word(required):
        out.append(required)
        seen.add(required)

    # 2. don't repeat the just-tapped word.
    last = (req.sentence[-1].strip().lower() if req.sentence else "")
    if last:
        seen.add(last)

    # 3. trie walk — bulk of candidates. Track each candidate's weight
    #    (number of phrases descending from it) so we can rank popular
    #    next-words first.
    node = _walk_trie(req.sentence or [])
    weights: dict[str, int] = {}
    for w, sub in node.items():
        wl = w.strip().lower()
        if _ok_word(wl) and wl not in seen:
            weights[wl] = max(weights.get(wl, 0), _leaf_count(sub))

    # 4. true dead-end — only when the trie had zero hits, top up from
    #    top-level openers so the grid is never empty. Don't do this when
    #    the prefix matched (even partially): mixing global openers into
    #    valid continuations produces ungrammatical buttons.
    if not weights and not out:
        for w, sub in _PHRASES.items():
            wl = w.strip().lower()
            if _ok_word(wl) and wl not in seen:
                weights[wl] = max(weights.get(wl, 0), _leaf_count(sub))

    # 5. for "more words" (n > 8) supplement with LLM if still short.
    #    The trie often has fewer than 30 grammatical continuations; let
    #    the LLM fill the tail rather than padding with global openers.
    used_llm = False
    if n > 8 and len(weights) + len(out) < n:
        used_llm = True
        sentence_str = " ".join((req.sentence or [])).strip()
        for w in _ask_llm(sentence_str, n=n + 4):
            wl = w.strip().lower()
            if _ok_word(wl) and wl not in seen and wl not in weights:
                weights[wl] = 0  # LLM tail — sorts last alphabetically.

    # 6. ranked head + alphabetical tail.
    #    Top 10 by descending leaf-count (most popular continuations first),
    #    then everything else sorted alphabetically.
    HEAD = 10
    by_freq = sorted(weights.items(), key=lambda kv: (-kv[1], kv[0]))
    head = [w for w, _ in by_freq[:HEAD]]
    tail = sorted([w for w, _ in by_freq[HEAD:]], key=lambda s: s.lower())
    for w in head + tail:
        if len(out) >= n:
            break
        if w in seen:
            continue
        out.append(w)
        seen.add(w)

    # 7. last-resort static fallback.
    if len(out) < n:
        for w in FALLBACK:
            if len(out) >= n:
                break
            if w in seen:
                continue
            out.append(w)
            seen.add(w)

    # 8. done_prob from trie shape.
    if isinstance(node, dict) and "" in node:
        done_prob = 0.6
    elif isinstance(node, dict) and not node:
        done_prob = 0.3 if req.sentence else 0.0
    else:
        done_prob = 0.0

    elapsed_ms = round((time.time() - t0) * 1000)
    sentence_str = " ".join(req.sentence or []).strip()
    _jsonl({
        "ts": time.time(),
        "kind": "predict",
        "sentence": sentence_str,
        "n": n,
        "required": req.required,
        "words": out[:n],
        "plan_source": "trie",
        "used_llm": used_llm,
        "elapsed_ms": elapsed_ms,
    })

    return {
        "words": out[:n],
        "done_prob": done_prob,
        "debug": {
            "sentence": sentence_str,
            "plan_source": "trie",
            "used_llm": used_llm,
            "trie_hits": len([w for w in node.keys() if _ok_word(w.strip().lower())]),
            "elapsed_ms": elapsed_ms,
        },
    }


@app.post("/pick")
def pick(req: PickReq):
    """Frontend reports which word was tapped — pure logging."""
    word = (req.word or "").strip().lower()
    sentence = list(req.sentence or [])
    _jsonl({
        "ts": time.time(),
        "kind": "pick",
        "sentence": " ".join(sentence).strip(),
        "word": word,
        "offered": [str(w).strip().lower() for w in (req.offered or [])],
        "mode": req.mode,
    })
    return {"ok": True}


@app.post("/add_phrase")
def add_phrase(req: AddPhraseReq):
    """User typed a word or phrase. Tokenize, validate, deep-merge into the
    trie, persist. Future /predict calls will see it."""
    raw = (req.phrase or "").strip().lower()
    tokens = [t for t in re.split(r"\s+", raw) if t]
    tokens = [t for t in tokens if _ok_word(t) or (t.replace("'", "").isalpha() and 1 <= len(t) <= 20 and t not in DENY)]
    # Allow contractions ("i'm", "let's") for full phrases — strip-then-check.
    tokens = [t for t in tokens if t.replace("'", "").isalpha() and 1 <= len(t) <= 20 and t not in DENY]
    if not tokens:
        return {"ok": False, "reason": "empty"}

    branch: dict = {}
    cur = branch
    for tok in tokens:
        cur[tok] = {}
        cur = cur[tok]

    with _PHRASES_LOCK:
        merged = _merge_trie(_PHRASES, branch)
        cleaned = _clean_trie(merged)
        _PHRASES.clear()
        _PHRASES.update(cleaned)
        try:
            with open(_PHRASES_PATH, "w") as f:
                json.dump(_PHRASES, f, ensure_ascii=False, indent=1)
        except OSError as e:
            log.warning("failed to persist phrases.json: %s", e)
            return {"ok": False, "reason": "write-failed"}

    _jsonl({"ts": time.time(), "kind": "add_phrase", "tokens": tokens})
    return {"ok": True, "tokens": tokens}


@app.get("/dictionary")
def dictionary(q: str = "", limit: int = 50):
    """Filter the curated kid-safe wordlist as the user types.
    Prefix matches first, then substring matches; both alphabetical."""
    q = (q or "").strip().lower()
    limit = max(1, min(limit, 200))
    if not q:
        return {"words": _DICT[:limit]}
    pref = [w for w in _DICT if w.startswith(q)]
    sub = [w for w in _DICT if (q in w) and not w.startswith(q)]
    return {"words": (pref + sub)[:limit]}


if __name__ == "__main__":  # pragma: no cover
    import uvicorn
    uvicorn.run("server:app", host="127.0.0.1", port=8900, reload=False)