app.py

"""Vaaani engine — CPU, OpenAI-compatible inference for the Vaaani RAG backend.

This is a HEADLESS model engine (no standalone UI). It is the drop-in for the RAG
product's VAAANI_LLM_BASE_URL, and it routes by the request's `model` field:

    vaaani-base      -> base Qwen2.5-3B GGUF, NO adapter   (general RAG / chat / ingest)
    vaaani-flagship  -> base + curriculum LoRA              (the Root-Bridge tutor)

The "no symbol before Grade 5" firewall is applied to sub-G5 tutor output only
(detected from "Grade N" in the system prompt); general RAG and G5 are untouched.

Endpoints:
    GET  /                     health JSON
    POST /v1/chat/completions  OpenAI-compatible, model-routed  (the RAG calls this)
    POST /chat, /chat/stream   simple test helpers
"""
import os
import re
import json
from typing import List, Optional

from fastapi import FastAPI
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from huggingface_hub import hf_hub_download

# ── config (override via Space Variables; secrets like HF_TOKEN stay secrets) ──
REPO          = os.environ.get("VAAANI_MODEL_REPO", "Shaankar39/vaaani-flagship-gguf")
BASE_FILE     = os.environ.get("VAAANI_BASE_FILE", "vaaani-base-q4_k_m.gguf")
LORA_FILE     = os.environ.get("VAAANI_LORA_FILE", "vaaani-flagship-lora-f16.gguf")
BASE_NAME     = os.environ.get("VAAANI_LLM_MODEL", "vaaani-base")
FLAGSHIP_NAME = os.environ.get("VAAANI_FLAGSHIP_MODEL", "vaaani-flagship")
N_THREADS = int(os.environ.get("N_THREADS", "2"))
N_CTX     = int(os.environ.get("N_CTX", "2048"))
N_BATCH   = int(os.environ.get("N_BATCH", "256"))
MAX_TOK   = int(os.environ.get("MAX_TOKENS", "512"))

app = FastAPI(title="Vaaani Engine (CPU)")
_llms = {}  # "base" | "flagship" -> Llama (lazy)


def get_llm(use_lora: bool):
    key = "flagship" if use_lora else "base"
    if key not in _llms:
        from llama_cpp import Llama
        kw = dict(model_path=hf_hub_download(REPO, BASE_FILE),
                  n_ctx=N_CTX, n_threads=N_THREADS, n_batch=N_BATCH, verbose=False)
        if use_lora:
            kw["lora_path"] = hf_hub_download(REPO, LORA_FILE)
        _llms[key] = Llama(**kw)
    return _llms[key]


def _use_lora(model_name: Optional[str]) -> bool:
    """Apply the curriculum adapter only when the flagship model is requested."""
    return (model_name or "").strip() == FLAGSHIP_NAME


class Msg(BaseModel):
    role: str
    content: str


class ChatReq(BaseModel):
    messages: List[Msg]
    model: Optional[str] = None
    temperature: float = 0.2
    max_tokens: Optional[int] = None
    stream: bool = False


def _msgs(req: ChatReq):
    return [{"role": m.role, "content": m.content} for m in req.messages]


# ── "no symbol before Grade 5" firewall (deterministic, serving-layer) ───────
_SLASH_PHONEME = re.compile(r"/[A-Za-zθðŋʃʒʧʤ]+/")
_IPA_CHARS = re.compile(r"[θðŋʃʒʧʤæɪʊəɔɑːʰˈˌ]")
_GRADE_RE = re.compile(r"Grade\s+(\d+)")


def _system_text(msgs):
    for m in msgs:
        if m["role"] == "system":
            return m["content"]
    return ""


def _is_sub_g5(system_text: str) -> bool:
    m = _GRADE_RE.search(system_text or "")
    return bool(m) and int(m.group(1)) < 5


def _firewall(text: str) -> str:
    text = _SLASH_PHONEME.sub("that sound", text)
    text = _IPA_CHARS.sub("", text)
    text = re.sub(r"\(\s*[,;]?\s*\)", "", text)
    text = re.sub(r"\s{2,}", " ", text).replace(" )", ")").replace("( ", "(")
    return text.strip()


class StreamScrubber:
    """Streaming-safe firewall for sub-G5: scrubs /phoneme/ and IPA across token
    boundaries while still streaming token-by-token. Holds text after a '/' until
    it can tell a phoneme (/b/) from a real slash (on/off)."""

    def __init__(self):
        self.hold = ""

    def feed(self, s: str) -> str:
        out = []
        for ch in s:
            if self.hold:
                if ch == "/":
                    out.append("that sound")
                    self.hold = ""
                elif (ch.isalpha() or ch in "θðŋʃʒʧʤ") and len(self.hold) <= 6:
                    self.hold += ch
                else:
                    out.append(self.hold + ch)
                    self.hold = ""
            elif ch == "/":
                self.hold = "/"
            else:
                out.append(ch)
        return _IPA_CHARS.sub("", "".join(out))

    def flush(self) -> str:
        rest, self.hold = self.hold, ""
        return _IPA_CHARS.sub("", rest)


@app.get("/")
def health():
    return {"status": "ok", "engine": "vaaani", "base": BASE_FILE, "lora": LORA_FILE,
            "models": [BASE_NAME, FLAGSHIP_NAME], "loaded": list(_llms.keys())}


@app.post("/chat")
def chat(req: ChatReq):
    sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
    out = get_llm(_use_lora(req.model)).create_chat_completion(
        messages=_msgs(req), temperature=req.temperature,
        max_tokens=req.max_tokens or MAX_TOK, stream=False)
    reply = out["choices"][0]["message"]["content"]
    if sub_g5:
        reply = _firewall(reply)
    return {"reply": reply}


@app.post("/chat/stream")
def chat_stream(req: ChatReq):
    sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
    llm = get_llm(_use_lora(req.model))

    def gen():
        scrub = StreamScrubber() if sub_g5 else None
        for chunk in llm.create_chat_completion(
                messages=_msgs(req), temperature=req.temperature,
                max_tokens=req.max_tokens or MAX_TOK, stream=True):
            delta = chunk["choices"][0]["delta"].get("content")
            if not delta:
                continue
            piece = scrub.feed(delta) if scrub else delta
            if piece:
                yield f"data: {json.dumps({'delta': piece})}\n\n"
        if scrub:
            tail = scrub.flush()
            if tail:
                yield f"data: {json.dumps({'delta': tail})}\n\n"
        yield "data: [DONE]\n\n"
    return StreamingResponse(gen(), media_type="text/event-stream")


@app.post("/v1/chat/completions")
def openai_compat(req: ChatReq):
    """OpenAI-compatible — the Vaaani RAG backend points VAAANI_LLM_BASE_URL here and
    calls /v1/chat/completions with model=vaaani-base or vaaani-flagship."""
    sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
    llm = get_llm(_use_lora(req.model))
    if req.stream:
        def gen():
            scrub = StreamScrubber() if sub_g5 else None
            for chunk in llm.create_chat_completion(
                    messages=_msgs(req), temperature=req.temperature,
                    max_tokens=req.max_tokens or MAX_TOK, stream=True):
                if sub_g5:
                    delta = chunk["choices"][0]["delta"].get("content")
                    if delta:
                        chunk["choices"][0]["delta"]["content"] = scrub.feed(delta)
                yield f"data: {json.dumps(chunk)}\n\n"
            if scrub:
                tail = scrub.flush()
                if tail:
                    yield f"data: {json.dumps({'choices':[{'index':0,'delta':{'content':tail},'finish_reason':None}]})}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(gen(), media_type="text/event-stream")
    out = llm.create_chat_completion(
        messages=_msgs(req), temperature=req.temperature,
        max_tokens=req.max_tokens or MAX_TOK, stream=False)
    if sub_g5:
        out["choices"][0]["message"]["content"] = _firewall(
            out["choices"][0]["message"]["content"])
    return JSONResponse(out)