rag.py

"""Long-term memory for LifeOS — fully local RAG.

Embeddings run through llama.cpp (nomic-embed-text-v1.5 GGUF, 146MB) — the
same runtime as the chat model, so the whole stack stays on-device. Notes
live in data/longterm.json with their vectors; retrieval is numpy cosine
top-k. No vector DB, no cloud.
"""

import hashlib
import json
import os
import threading
import time
import uuid

import numpy as np

import config

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
STORE_PATH = os.path.join(DATA_DIR, "longterm.json")

EMBED_REPO = config.EMBED_REPO
EMBED_FILE = config.EMBED_FILE

_lock = threading.Lock()
_embedder = None


def warmup() -> None:
    """Load the embedder eagerly. Called at startup so nomic-embed is resident
    before the vision model ever loads — loading a model *after* the VLM can
    fail on small GPUs, and every chat needs the embedder for recall anyway."""
    _get_embedder()


def _get_embedder():
    global _embedder
    if _embedder is None:
        import cuda_bootstrap

        cuda_bootstrap.ensure()
        from llama_cpp import Llama

        _embedder = Llama.from_pretrained(
            repo_id=EMBED_REPO,
            filename=EMBED_FILE,
            embedding=True,
            n_ctx=2048,
            n_gpu_layers=config.GPU_LAYERS,
            verbose=False,
        )
    return _embedder


def embed(text: str, is_query: bool = False) -> np.ndarray:
    # nomic-embed expects task prefixes
    prefixed = ("search_query: " if is_query else "search_document: ") + text
    with _lock:
        out = _get_embedder().create_embedding(prefixed)
    vec = np.asarray(out["data"][0]["embedding"], dtype=np.float32)
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


def _load_store() -> list[dict]:
    if not os.path.exists(STORE_PATH):
        return []
    with open(STORE_PATH, "r", encoding="utf-8") as f:
        notes = json.load(f)
    # Backfill ids on notes written before note management existed.
    if any(not n.get("id") for n in notes):
        for n in notes:
            if not n.get("id"):
                n["id"] = uuid.uuid4().hex[:8]
        _save_store(notes)
    return notes


def _save_store(notes: list[dict]) -> None:
    os.makedirs(DATA_DIR, exist_ok=True)
    with open(STORE_PATH, "w", encoding="utf-8") as f:
        json.dump(notes, f, ensure_ascii=False)


def remember(text: str, kind: str = "fact", meta: dict | None = None) -> dict:
    """Embed and persist a long-term note. kind: fact|event|preference."""
    note = {
        "id": hashlib.md5(f"{text}{time.time()}".encode()).hexdigest()[:12],
        "text": text,
        "kind": kind,
        "ts": time.time(),
        "meta": meta or {},
        "vec": embed(text).tolist(),
    }
    with _lock:
        notes = _load_store()
        # skip exact-duplicate texts
        if any(n["text"] == text for n in notes):
            return note
        notes.append(note)
        _save_store(notes)
    return note


def recall(query: str, k: int = 5, kind: str | None = None) -> list[dict]:
    """Top-k notes by cosine similarity; returns [{text, kind, score}]."""
    with _lock:
        notes = _load_store()
    if kind:
        notes = [n for n in notes if n["kind"] == kind]
    if not notes:
        return []
    q = embed(query, is_query=True)
    # Self-heal: notes embedded under a different embedder (e.g. a store from
    # an older build) have the wrong vector dimension — re-embed them.
    stale = [n for n in notes if len(n["vec"]) != len(q)]
    if stale:
        for n in stale:
            n["vec"] = embed(n["text"]).tolist()
        with _lock:
            current = _load_store()
            by_id = {n["id"]: n["vec"] for n in notes}
            for n in current:
                if n["id"] in by_id:
                    n["vec"] = by_id[n["id"]]
            _save_store(current)
    mat = np.asarray([n["vec"] for n in notes], dtype=np.float32)
    scores = mat @ q
    order = np.argsort(-scores)[:k]
    return [
        {"text": notes[i]["text"], "kind": notes[i]["kind"], "score": float(scores[i])}
        for i in order
        if scores[i] > 0
    ]


def list_notes() -> list[dict]:
    """All long-term notes without vectors: [{id, text, kind}]."""
    with _lock:
        notes = _load_store()
    return [{"id": n["id"], "text": n["text"], "kind": n["kind"]} for n in notes]


def update_note(note_id: str, text: str) -> bool:
    """Replace a note's text (re-embeds; kind/meta preserved). True if found."""
    vec = embed(text).tolist()
    with _lock:
        notes = _load_store()
        for n in notes:
            if n["id"] == note_id:
                n["text"] = text
                n["vec"] = vec
                n["ts"] = time.time()
                _save_store(notes)
                return True
    return False


def delete_note(note_id: str) -> bool:
    with _lock:
        notes = _load_store()
        kept = [n for n in notes if n["id"] != note_id]
        if len(kept) != len(notes):
            _save_store(kept)
            return True
    return False


SEED_NOTES = [
    ("Awais is lactose-sensitive — fine with yogurt and hard cheese, avoid milk-heavy dishes", "preference"),
    ("Awais eats halal only and skips pork entirely", "preference"),
    ("Awais hates cilantro", "preference"),
    ("Awais aims for at least 120g of protein per day", "preference"),
    ("Awais's weekly grocery budget is $80 and he shops on Saturdays", "fact"),
    ("Awais shares streaming accounts with his roommate Hamza — Disney+ was Hamza's idea", "fact"),
    ("Awais signed up for FitnessPal Pro during a January resolution and stopped using it in March", "event"),
    ("Awais is training for a 10K in September; longest run so far is 6km", "fact"),
    ("Awais's left knee gets sore if he runs two days in a row — needs a rest day between runs", "preference"),
    ("Awais prefers lifting in the evening after classes, runs in the morning", "preference"),
    ("Awais's go-to cheap protein is chicken thighs and canned chickpeas", "preference"),
    ("Awais cooked too much pasta lately and wants more variety in meals", "event"),
    ("Awais gets paid around the 1st of each month — roughly $1400 from his part-time job", "fact"),
    ("Awais wants to save $150 a month toward a climbing gym membership", "fact"),
    ("Awais usually meal-preps rice in bulk on Sundays", "preference"),
]


def ensure_seeded() -> int:
    """Seed the long-term store on first run. Returns note count."""
    with _lock:
        existing = _load_store()
    if existing:
        return len(existing)
    for text, kind in SEED_NOTES:
        remember(text, kind=kind, meta={"seed": True})
    with _lock:
        return len(_load_store())


def reset_to_seed() -> int:
    with _lock:
        if os.path.exists(STORE_PATH):
            os.remove(STORE_PATH)
    return ensure_seeded()


def reset_to_empty() -> int:
    """Delete all long-term notes (real-user reset). Returns 0."""
    with _lock:
        if os.path.exists(STORE_PATH):
            os.remove(STORE_PATH)
    return 0