"""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)