app.py

"""
Momo 1.0 — Babymomo brain API (FastAPI + optional Lightning App wrapper).

Pipeline:  user text -> Momo 1.0 JSON decision -> memory ops (store/update/delete)
           -> memory search (bge-small-en-v1.5) + Serper Google search -> final answer.

Memory is the Source of Truth; Google is extra. Personal data is anonymized
out of web_query by Momo before any web call (enforced again below as a safety net).

Run modes:
  python app.py                     # plain FastAPI+uvicorn (works in Lightning Studio terminal)
  python app.py --lightning         # Lightning App (lightning run app app.py --lightning? no: edit flag below)
  lightning run app app.py          # Lightning App mode (uses MomoFlow below)

Env vars:
  SERPER_API_KEY      google.serper.dev key          (required for web search)
  HF_TOKEN            huggingface token              (optional, adapter download)
  MOMO_HF_REPO        adapter repo (default Ansaribilal/momo-1.0)
  MOMO_ADAPTER        local adapter dir (default ./adapter if present)
  MOMO_MEMORY_PATH    memory json file (default ./memory_store.json)
  MOMO_ANSWER_LLM     "1" to generate answers with the model (slower), default template answers
"""
from __future__ import annotations

import json
import os
import re
import statistics
import sys
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

import httpx

sys_path = str(Path(__file__).resolve().parent)
if sys_path not in sys.path:
    sys.path.insert(0, sys_path)

from momo_core import (  # noqa: E402
    BASE_MODEL_ID,
    EMBED_MODEL_ID,
    SYSTEM_PROMPT,
    build_decision,
    parse_decision,
    validate_decision,
)

# ------------------------------------------------------------------ settings --

MOMO_HF_REPO = os.environ.get("MOMO_HF_REPO", "Ansaribilal/momo-1.0")
MOMO_ADAPTER = os.environ.get("MOMO_ADAPTER", "adapter" if Path("adapter").exists() else "")
MOMO_MEMORY_PATH = os.environ.get("MOMO_MEMORY_PATH", "memory_store.json")
MOMO_ANSWER_LLM = os.environ.get("MOMO_ANSWER_LLM", "0") == "1"
SERPER_URL = "https://google.serper.dev/search"

# ------------------------------------------------------------------ momo brain --


class MomoBrain:
    """Loads Qwen2-0.5B-Instruct (+ LoRA adapter when available) and emits decisions."""

    def __init__(self, base: str = BASE_MODEL_ID, adapter: str = MOMO_ADAPTER,
                 hf_repo: str = MOMO_HF_REPO):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.torch = torch
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if self.device == "cuda" else torch.float32

        self.tok = AutoTokenizer.from_pretrained(base, trust_remote_code=True)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(
            base, torch_dtype=dtype, trust_remote_code=True
        ).to(self.device).eval()

        adapter_used = None
        if adapter and Path(adapter).exists():
            adapter_used = adapter
        elif adapter and (os.path.sep not in adapter and ":" not in adapter):
            # treat as HF repo id if it looks like one, else skip
            adapter_used = adapter if "/" in adapter else None
        if not adapter_used and hf_repo and os.environ.get("HF_TOKEN"):
            adapter_used = hf_repo  # peft will pull from the hub

        if adapter_used:
            try:
                from peft import PeftModel
                self.model = PeftModel.from_pretrained(self.model, adapter_used)
                self.model = self.model.merge_and_unload()
                print(f"[momo] adapter loaded: {adapter_used}")
            except Exception as e:  # noqa: BLE001
                print(f"[momo] WARNING adapter load failed ({e}) — running on base weights")
        else:
            print("[momo] WARNING: no adapter — decisions come from base weights "
                  "(run train.py or set MOMO_HF_REPO+HF_TOKEN)")

        self.decision_latencies: List[float] = []

    def decide(self, user_text: str) -> Dict[str, Any]:
        """user text -> validated Momo decision dict (greedy, JSON-only target)."""
        t0 = time.perf_counter()
        prompt = self.tok.apply_chat_template(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_text},
            ],
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = self.tok(prompt, return_tensors="pt").to(self.device)
        with self.torch.inference_mode():
            out = self.model.generate(
                **inputs,
                max_new_tokens=140,
                do_sample=False,
                pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id,
                eos_token_id=self.tok.eos_token_id,
            )
        gen = out[0][inputs["input_ids"].shape[1]:]
        raw = self.tok.decode(gen, skip_special_tokens=True)
        dec = parse_decision(raw)
        dt = time.perf_counter() - t0
        self.decision_latencies.append(dt)
        if dec is None:
            # conservative fallback: try memory search, never fabricate an action
            dec = build_decision("MEMORY_ONLY", need_memory_search=True,
                                 memory_query=user_text[:120])
            dec["fallback"] = True
        problems = validate_decision(dec)
        if problems:
            dec["validation"] = problems
        dec["_latency_s"] = round(dt, 4)
        return dec

# ------------------------------------------------------------------ memory --


class LocalJSONMemory:
    """
    Fallback memory backend (Source of Truth): JSON file + optional bge-small vectors.
    Swap with Babymomo's production backend by overriding these four methods.
    """

    def __init__(self, path: str = MOMO_MEMORY_PATH, use_embeddings: bool = True):
        self.path = Path(path)
        self.entries: List[Dict[str, Any]] = []
        self.embed_model = None
        self.embed_failed = False
        if self.path.exists():
            try:
                self.entries = json.loads(self.path.read_text(encoding="utf-8"))
            except Exception:  # noqa: BLE001
                self.entries = []
        if use_embeddings:
            self._try_load_embeddings()

    def _try_load_embeddings(self) -> None:
        try:
            from sentence_transformers import SentenceTransformer
            self.embed_model = SentenceTransformer(EMBED_MODEL_ID)
            for e in self.entries:
                if "vec" not in e:
                    e["vec"] = self.embed_model.encode(e["text"], normalize_embeddings=True).tolist()
        except Exception as e:  # noqa: BLE001
            self.embed_model = None
            self.embed_failed = True
            print(f"[memory] embeddings unavailable ({e}) — using token-overlap scoring")

    def _embed(self, text: str) -> Optional[List[float]]:
        if self.embed_model is None:
            return None
        return self.embed_model.encode(text, normalize_embeddings=True).tolist()

    @staticmethod
    def _overlap(a: str, b: str) -> float:
        sa = set(re.findall(r"\w+", a.lower()))
        sb = set(re.findall(r"\w+", b.lower()))
        if not sa or not sb:
            return 0.0
        return len(sa & sb) / (len(sa) | len(sb))

    def search(self, query: str, k: int = 4) -> List[Dict[str, Any]]:
        """Semantic search over memory. Target: <100ms (measured, returned by caller)."""
        t0 = time.perf_counter()
        qv = self._embed(query)
        scored = []
        for e in self.entries:
            if qv is not None and e.get("vec"):
                sim = sum(x * y for x, y in zip(qv, e["vec"]))
            else:
                sim = self._overlap(query, e["text"])
            scored.append((sim, e))
        scored.sort(key=lambda p: p[0], reverse=True)
        hits = [
            {"id": e["id"], "text": e["text"], "kind": e.get("kind", "fact"), "score": round(s, 4)}
            for s, e in scored[:k] if s > 0.15
        ]
        return hits

    def upsert(self, text: str, kind: str = "fact") -> Dict[str, Any]:
        """Smart save: skips exact duplicates, refreshes timestamp otherwise. Returns entry."""
        for e in self.entries:
            if e["text"].strip().lower() == text.strip().lower():
                e["ts"] = time.time()
                self._flush()
                return {"updated": True, "entry": e}
        entry = {"id": uuid.uuid4().hex[:12], "text": text, "kind": kind, "ts": time.time()}
        vec = self._embed(text)
        if vec is not None:
            entry["vec"] = vec
        self.entries.append(entry)
        self._flush()
        return {"created": True, "entry": entry}

    def delete_by_query(self, query: str, threshold: float = 0.55) -> List[Dict[str, Any]]:
        """Smart delete: removes best-matching entries above threshold."""
        qv = self._embed(query)
        keep, removed = [], []
        for e in self.entries:
            if qv is not None and e.get("vec"):
                sim = sum(x * y for x, y in zip(qv, e["vec"]))
            else:
                sim = self._overlap(query, e["text"])
            if sim >= threshold:
                removed.append(e)
            else:
                keep.append(e)
        if removed:
            self.entries = keep
            self._flush()
        return removed

    def _flush(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        slim = [{k: v for k, v in e.items() if k != "vec"} for e in self.entries]
        self.path.write_text(json.dumps(slim, ensure_ascii=False, indent=2), encoding="utf-8")

# --------------------------------------------------------------- web search --


def serper_search(query: str, num: int = 5) -> List[Dict[str, Any]]:
    """Google search via Serper.dev. Returns [{'title','link','snippet'}, ...]."""
    key = os.environ.get("SERPER_API_KEY")
    if not key:
        return [{"error": "SERPER_API_KEY not set — web search skipped"}]
    try:
        r = httpx.post(
            SERPER_URL,
            headers={"X-API-KEY": key, "Content-Type": "application/json"},
            json={"q": query, "num": num},
            timeout=8.0,
        )
        r.raise_for_status()
        organic = r.json().get("organic", [])
        return [{"title": o.get("title", ""), "link": o.get("link", ""),
                 "snippet": o.get("snippet", "")} for o in organic[:num]]
    except Exception as e:  # noqa: BLE001
        return [{"error": f"serper failed: {e}"}]

# ----------------------------------------------------- answer composition --

_PII_PATTERNS = [
    re.compile(r"\b(my|his|her|their)\s+(girlfriend|boyfriend|wife|husband|sister|brother|mother|father"
               r"|cousin|roommate|boss|colleague|friend|daughter|son|nephew|niece|grandma|grandpa"
               r"|neighbor|trainer|dentist|therapist)\s+([A-Z][a-z]+)"),
    re.compile(r"\b([A-Z][a-z]+)\s+(said|told|recommended|claims|thinks|advised|swears)"),
    re.compile(r"\bDr\.?\s+[A-Z][a-z]+"),
    re.compile(r"\b\+?\d[\d\s().-]{7,}\d\b"),  # phone-like numbers
    re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.]+\b"),  # emails
]


def anonymize_web_query(q: str) -> str:
    """Safety net: Momo should already anonymize; this scrubs leaks before Serper."""
    cleaned = q
    for pat in _PII_PATTERNS:
        cleaned = pat.sub(" ", cleaned)
    cleaned = re.sub(r"\s{2,}", " ", cleaned).strip(" -,")
    return cleaned


def compose_answer(user_text: str, dec: Dict[str, Any], mem_hits: List[Dict[str, Any]],
                   web_hits: List[Dict[str, Any]], memory_ops: Dict[str, Any]) -> str:
    """Deterministic grounded answer (fast path). Set MOMO_ANSWER_LLM=1 for model answers."""
    parts: List[str] = []
    if dec.get("fallback"):
        parts.append("I couldn't fully parse that — searching your memory to be safe.")
    if memory_ops.get("stored"):
        parts.append(f"Saved to memory: \"{memory_ops['stored']}\"")
    if memory_ops.get("updated"):
        parts.append("Memory updated accordingly.")
    if memory_ops.get("deleted"):
        names = ", ".join(f"\"{d['text']}\"" for d in memory_ops["deleted"][:3])
        parts.append(f"Removed from memory: {names}." if names else "Nothing matched that well enough to delete.")
    if mem_hits:
        lines = "\n".join(f"- {h['text']} (relevance {h['score']})" for h in mem_hits[:3])
        parts.append(f"From your memory:\n{lines}")
    elif dec.get("need_memory_search") and not memory_ops.get("deleted"):
        parts.append("Nothing relevant found in your memory yet.")
    if web_hits and not web_hits[0].get("error"):
        lines = "\n".join(f"- {w['title']}: {w['snippet']} ({w['link']})" for w in web_hits[:3])
        parts.append(f"From the web:\n{lines}")
    elif dec.get("need_web") and web_hits and web_hits[0].get("error"):
        parts.append(f"Web search unavailable: {web_hits[0]['error']}")
    if not parts:
        return "Got it. Anything else?"
    return "\n\n".join(parts)

# ------------------------------------------------------------------ fastapi --


def build_app() -> "FastAPI":  # noqa: F821
    from fastapi import FastAPI
    from pydantic import BaseModel

    app = FastAPI(title="Momo 1.0 — Babymomo Brain", version="1.0")
    brain: Optional[MomoBrain] = None

    class TextIn(BaseModel):
        text: str
        user_id: str = "default"
        k_memory: int = 4
        k_web: int = 5

    @app.on_event("startup")
    def _startup() -> None:
        nonlocal brain
        brain = MomoBrain()

    @app.get("/health")
    def health() -> Dict[str, Any]:
        lat = brain.decision_latencies[-20:] if brain else []
        return {
            "status": "ok", "model": "momo-1.0", "version": "1.0",
            "device": brain.device if brain else None,
            "decision_avg_ms": round(statistics.mean(lat) * 1000, 1) if lat else None,
            "web_enabled": bool(os.environ.get("SERPER_API_KEY")),
        }

    @app.post("/decide")
    def decide(payload: TextIn) -> Dict[str, Any]:
        """Fast path: user text -> Momo decision JSON only (no memory/web side effects)."""
        assert brain is not None
        return {"decision": brain.decide(payload.text)}

    @app.post("/chat")
    def chat(payload: TextIn) -> Dict[str, Any]:
        """Full pipeline: decide -> memory ops -> searches -> final answer."""
        assert brain is not None
        memory = get_backend(payload.user_id)
        dec = brain.decide(payload.text)

        # 1) memory write operations (memory backend is smart: keep/delete/update)
        memory_ops: Dict[str, Any] = {}
        if dec["action"] == "STORE_MEMORY" and dec["memory_text"]:
            memory_ops["stored"] = memory.upsert(dec["memory_text"], kind="fact")["entry"]["text"]
        elif dec["action"] == "UPDATE_MEMORY" and dec["memory_text"]:
            memory.delete_by_query(dec["memory_query"])
            memory.upsert(dec["memory_text"], kind="fact")
            memory_ops["updated"] = dec["memory_text"]
        elif dec["action"] == "DELETE_MEMORY" and dec["memory_query"]:
            memory_ops["deleted"] = memory.delete_by_query(dec["memory_query"])

        # 2) memory search (source of truth)
        t0 = time.perf_counter()
        mem_hits = memory.search(dec["memory_query"], k=payload.k_memory) \
            if dec["need_memory_search"] and dec["memory_query"] else []
        memory_ms = (time.perf_counter() - t0) * 1000

        # 3) web search (extra, anonymized)
        wq = anonymize_web_query(dec["web_query"])
        web_hits = serper_search(wq, num=payload.k_web) \
            if dec["need_web"] and wq else []

        # 4) final answer
        if MOMO_ANSWER_LLM and (mem_hits or web_hits):
            answer = llm_answer(brain, payload.text, dec, mem_hits, web_hits)
        else:
            answer = compose_answer(payload.text, dec, mem_hits, web_hits, memory_ops)

        return {
            "answer": answer,
            "decision": {k: v for k, v in dec.items() if not k.startswith("_")},
            "web_query_sent": wq or None,
            "memory_hits": mem_hits,
            "web_hits": web_hits,
            "memory_ops": memory_ops,
            "timing_ms": {"memory_filter": round(memory_ms, 2),
                          "decision": round(dec.get("_latency_s", 0) * 1000, 1)},
        }

    return app


def llm_answer(brain: MomoBrain, user_text: str, dec: Dict[str, Any],
               mem_hits: List[Dict[str, Any]], web_hits: List[Dict[str, Any]]) -> str:
    """Optional: generate a natural answer with Momo itself (MOMO_ANSWER_LLM=1)."""
    mem_ctx = "\n".join(f"- {h['text']}" for h in mem_hits) or "(none)"
    web_ctx = "\n".join(f"- {w['title']}: {w['snippet']}" for w in web_hits if not w.get("error")) or "(none)"
    msgs = [
        {"role": "system", "content": "You are Momo 1.0. Answer the user using their memory "
                                      "(source of truth) and web results. Be concise and personal."},
        {"role": "user", "content": f"User said: {user_text}\n\nMemory:\n{mem_ctx}\n\nWeb:\n{web_ctx}\n\nAnswer:"},
    ]
    try:
        prompt = brain.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        inputs = brain.tok(prompt, return_tensors="pt").to(brain.device)
        with brain.torch.inference_mode():
            out = brain.model.generate(**inputs, max_new_tokens=180, do_sample=False,
                                       pad_token_id=brain.tok.pad_token_id or brain.tok.eos_token_id)
        gen = out[0][inputs["input_ids"].shape[1]:]
        return brain.tok.decode(gen, skip_special_tokens=True).strip()
    except Exception as e:  # noqa: BLE001
        return compose_answer(user_text, dec, mem_hits, web_hits, {}) + f"\n(llm answer failed: {e})"

# ----------------------------------------------------------------- backend --


_BACKEND: Optional[LocalJSONMemory] = None


def get_backend(user_id: str) -> LocalJSONMemory:
    """Backend hook. Replace with Babymomo's production memory backend when ready.
    Multi-user ready: append user_id to the storage path."""
    global _BACKEND
    if _BACKEND is None:
        base, ext = os.path.splitext(MOMO_MEMORY_PATH)
        _BACKEND = LocalJSONMemory(f"{base}.{user_id}{ext}" if user_id != "default" else MOMO_MEMORY_PATH)
    return _BACKEND

# ---------------------------------------------------------------- lightning --


def run_uvicorn(host: str = "0.0.0.0", port: int = 8000) -> None:
    import uvicorn
    uvicorn.run(build_app(), host=host, port=port)


try:
    import lightning as L
    HAS_LIGHTNING = True
except ImportError:  # plain python app.py still works without the lightning SDK
    HAS_LIGHTNING = False

if HAS_LIGHTNING:

    class MomoWork(L.LightningWork):
        """Runs the FastAPI server inside a Lightning Work process.
        Add cloud_compute=L.CloudCompute('gpu') when serving on a T4."""

        def __init__(self):
            super().__init__(parallel=True)

        def run(self):
            run_uvicorn(self.host, self.port)

    class MomoFlow(L.LightningFlow):
        """Lightning App entry: `lightning run app app.py`"""

        def __init__(self):
            super().__init__()
            self.work = MomoWork()

        def run(self):
            self.work.run()
            if self.work.url:
                print(f"[momo] public URL: {self.work.url}")

    # `lightning run app app.py` discovers this variable; harmless under `python app.py`
    app = L.LightningApp(MomoFlow())


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser(description="Momo 1.0 API")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=int(os.environ.get("PORT", 8000)))
    args = ap.parse_args()
    print(f"[momo] starting API on {args.host}:{args.port}")
    run_uvicorn(args.host, args.port)