momo_core.py

"""
Momo 1.0 — the decision brain of Babymomo.
Core contract shared by train.py / test.py / app.py / data generation.

All intelligence lives in the model weights: one JSON decision per user turn.
No agents.md, no skills files, no LangChain — just Momo.
"""
from __future__ import annotations

import json
import re
from typing import Any, Dict, List, Optional

# ---------------------------------------------------------------- constants --
MOMO_NAME = "Momo"
MOMO_VERSION = "1.0"
REPO_ID = "momo-1.0"

BASE_MODEL_ID = "Qwen/Qwen2-0.5B-Instruct"
EMBED_MODEL_ID = "BAAI/bge-small-en-v1.5"  # optional, memory search only

ACTIONS = (
    "STORE_MEMORY",
    "UPDATE_MEMORY",
    "DELETE_MEMORY",
    "MEMORY_ONLY",
    "WEB_ONLY",
    "HYBRID",
)

REQUIRED_KEYS = (
    "action",
    "memory_query",
    "web_query",
    "need_memory_search",
    "need_web",
    "memory_text",
)

# One compact system prompt. LoRA bakes the behavior into the weights, but we
# keep train/inference prompts IDENTICAL for maximum reliability.
SYSTEM_PROMPT = (
    "You are Momo 1.0, the decision brain of the Babymomo app. Read the user's "
    "message and reply with ONLY one JSON object, no other text:\n"
    '{"action": "STORE_MEMORY|UPDATE_MEMORY|DELETE_MEMORY|MEMORY_ONLY|WEB_ONLY|HYBRID", '
    '"memory_query": "", "web_query": "", "need_memory_search": false, "need_web": false, '
    '"memory_text": ""}\n'
    "Rules:\n"
    "- New personal fact shared -> STORE_MEMORY: put a clean third-person fact in memory_text.\n"
    "- Correction or changed fact -> UPDATE_MEMORY: memory_query finds the old record, memory_text holds the new fact.\n"
    "- Wants something forgotten -> DELETE_MEMORY: memory_query describes what to delete.\n"
    "- Question about the user's own life -> MEMORY_ONLY: memory_query searches memory.\n"
    "- General/world/external question -> WEB_ONLY: web_query is a short web search.\n"
    "- Personal + world mix, or save + search together -> HYBRID: fill memory_query and web_query "
    "(or memory_text when something must also be saved).\n"
    "- NEVER put personal details (private names, phone numbers, emails, ids) in web_query. Anonymize it.\n"
    "- If nothing is needed, leave queries empty and flags false."
)

# ----------------------------------------------------------------- builders --


def build_decision(
    action: str,
    memory_query: str = "",
    web_query: str = "",
    need_memory_search: bool = False,
    need_web: bool = False,
    memory_text: str = "",
) -> Dict[str, Any]:
    """Build a decision dict with a FIXED key order (training consistency)."""
    if action not in ACTIONS:
        raise ValueError(f"invalid action: {action!r}")
    return {
        "action": action,
        "memory_query": memory_query,
        "web_query": web_query,
        "need_memory_search": bool(need_memory_search),
        "need_web": bool(need_web),
        "memory_text": memory_text,
    }


def decision_to_json(decision: Dict[str, Any]) -> str:
    """Serialize decision dict -> canonical JSON string for training targets."""
    ordered = {k: decision.get(k, "" if k in ("memory_query", "web_query", "memory_text") else False)
               for k in REQUIRED_KEYS}
    if ordered["action"] not in ACTIONS:
        raise ValueError(f"invalid action: {ordered['action']!r}")
    ordered["need_memory_search"] = bool(ordered["need_memory_search"])
    ordered["need_web"] = bool(ordered["need_web"])
    return json.dumps(ordered, ensure_ascii=False)


_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)


def parse_decision(text: str) -> Optional[Dict[str, Any]]:
    """
    Parse the first JSON-looking block out of model output.
    Returns a validated decision dict, or None if unrecoverable.
    Tolerant: coerces bool-like values, trims stray text around the object.
    """
    if not text or not text.strip():
        return None
    m = _JSON_BLOCK_RE.search(text)
    if not m:
        return None
    raw = m.group(0)
    try:
        obj = json.loads(raw)
    except json.JSONDecodeError:
        try:  # single-quote fallback (model slip)
            obj = json.loads(raw.replace("'", '"'))
        except json.JSONDecodeError:
            return None
    if not isinstance(obj, dict) or "action" not in obj:
        return None
    obj["action"] = str(obj["action"]).upper().strip()
    if obj["action"] not in ACTIONS:
        return None
    for k in REQUIRED_KEYS:
        obj.setdefault(k, "" if k in ("memory_query", "web_query", "memory_text") else False)
    for k in ("memory_query", "web_query", "memory_text"):
        if obj[k] is None:
            obj[k] = ""
        obj[k] = str(obj[k])
    for k in ("need_memory_search", "need_web"):
        v = obj[k]
        if isinstance(v, str):
            v = v.strip().lower() in ("true", "1", "yes")
        obj[k] = bool(v)
    return obj


def validate_decision(decision: Dict[str, Any]) -> List[str]:
    """Semantic sanity checks used by test.py and the data validator."""
    problems: List[str] = []
    a = decision.get("action")
    mq, wq, mt = decision.get("memory_query", ""), decision.get("web_query", ""), decision.get("memory_text", "")
    ms, nw = bool(decision.get("need_memory_search")), bool(decision.get("need_web"))

    if a == "STORE_MEMORY" and not mt:
        problems.append("STORE_MEMORY requires memory_text")
    if a == "UPDATE_MEMORY" and (not mq or not mt):
        problems.append("UPDATE_MEMORY requires memory_query + memory_text")
    if a == "UPDATE_MEMORY" and not ms:
        problems.append("UPDATE_MEMORY must search memory to find the old record")
    if a == "DELETE_MEMORY" and not mq:
        problems.append("DELETE_MEMORY requires memory_query")
    if a == "DELETE_MEMORY" and not ms:
        problems.append("DELETE_MEMORY must search memory to find the record")
    if a == "MEMORY_ONLY" and ms and not mq:
        problems.append("MEMORY_ONLY with search=true requires memory_query")
    if a == "WEB_ONLY" and not wq:
        problems.append("WEB_ONLY requires web_query")
    if a == "HYBRID" and not ((mq or mt) and wq):
        problems.append("HYBRID requires (memory_query or memory_text) and web_query")
    if ms and not mq and a != "STORE_MEMORY":
        problems.append("need_memory_search=true but memory_query empty")
    if nw and not wq:
        problems.append("need_web=true but web_query empty")
    # MEMORY_ONLY with everything empty = pure conversation path (chit-chat) — allowed.
    if a != "MEMORY_ONLY" and not ms and not nw and not mt:
        problems.append("no retrieval requested and nothing to store (dead decision)")
    return problems