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