NexusOS/synapse/memory/service.py

"""Dedicated memory curator service — run alongside Synapse on port 8001.

Endpoints:
  GET  /            health check
  GET  /memories    list all memory items (optional ?section= filter)
  POST /memories    direct write — no LLM, saves immediately
  PATCH /memories/{id}  update a memory item
  DELETE /memories/{id} delete a memory item
  POST /memories/extract  LLM-curated: evaluate a conversation exchange and
                          optionally save a new permanent fact
"""

from __future__ import annotations

import asyncio
import math
import uuid
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException, Body
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from .store import store, MemoryItem
from .extractor import extract_memory
from ..ollama_manager import get_ollama_manager

app = FastAPI(title="Nexus Memory Service", version="1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def _warm_curator():
    """Preload the curator model (in RAM, num_gpu=0 by default) so the first
    extraction isn't a cold load that blows the timeout. Runs in the background
    so it never delays startup. keep_alive then holds it warm between messages."""
    async def _bg():
        try:
            mgr = get_ollama_manager()
            # Ollama is started by the Synapse backend (a separate process), so
            # at our startup it usually isn't reachable yet. Wait for it before
            # warming instead of failing with "All connection attempts failed" —
            # which leaves the curator cold and makes the first extraction slow.
            for _ in range(60):  # up to ~2 min
                if await asyncio.to_thread(mgr.is_running):
                    break
                await asyncio.sleep(2)
            else:
                return
            settings = store.get_settings()
            model = settings.get("memory_model") or await mgr.select_best_model()
            num_gpu = await mgr.resolve_num_gpu(settings.get("memory_gpu_offload", 0), model)
            await mgr.warm(model, num_gpu=num_gpu)
        except Exception:
            pass
    asyncio.create_task(_bg())


@app.get("/")
async def health():
    return {"status": "ok", "count": len(store.all())}


@app.get("/memories")
async def list_memories(section: Optional[str] = None):
    items = store.all()
    if section:
        items = [i for i in items if i.section.lower() == section.lower()]
    return {"items": [{"id": i.id, "section": i.section, "text": i.text, "tags": i.tags} for i in items]}


@app.post("/memories")
async def add_memory(payload: Dict[str, Any] = Body(...)):
    text = (payload.get("text") or "").strip()
    if not text:
        raise HTTPException(status_code=400, detail="Missing 'text'")
    item = MemoryItem(
        id=str(uuid.uuid4()),
        section=(payload.get("section") or "General").strip(),
        text=text,
        tags=payload.get("tags", []),
    )
    store.add(item)
    return {"id": item.id, "section": item.section, "text": item.text, "tags": item.tags}


@app.patch("/memories/{item_id}")
async def update_memory(item_id: str, payload: Dict[str, Any] = Body(...)):
    existing = store.get(item_id)
    if not existing:
        raise HTTPException(status_code=404, detail="Not found")
    updated = MemoryItem(
        id=item_id,
        section=(payload.get("section") or existing.section).strip(),
        text=(payload.get("text") or existing.text).strip(),
        tags=payload.get("tags", existing.tags),
    )
    store.update(updated)
    return {"id": updated.id, "section": updated.section, "text": updated.text, "tags": updated.tags}


@app.delete("/memories/{item_id}")
async def delete_memory(item_id: str):
    if not store.get(item_id):
        raise HTTPException(status_code=404, detail="Not found")
    store.delete(item_id)
    return {"status": "deleted"}


def _cosine(a: list, b: list) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb) if na and nb else 0.0


class ExtractRequest(BaseModel):
    user_message: str
    assistant_response: str


@app.post("/memories/extract")
async def extract_and_save(req: ExtractRequest):
    """LLM-curated extraction — asks Mistral to evaluate the exchange against all
    existing memory items and save only new, permanent personal facts."""
    existing = store.all()
    existing_sections = list({i.section for i in existing})
    existing_texts = [i.text for i in existing]

    # Adaptable per-machine curator config (see store _SETTINGS_DEFAULTS):
    # which model does extraction, and whether it runs on CPU/RAM or the GPU.
    settings = store.get_settings()
    mgr = get_ollama_manager()
    model = settings.get("memory_model") or await mgr.select_best_model()
    num_gpu = await mgr.resolve_num_gpu(settings.get("memory_gpu_offload", 0), model)
    try:
        merge_threshold = float(settings.get("memory_merge_threshold", 0.88))
    except (TypeError, ValueError):
        merge_threshold = 0.88

    try:
        results = await asyncio.wait_for(
            extract_memory(
                req.user_message,
                req.assistant_response,
                existing_sections,
                existing_texts,
                mgr,
                model=model,
                num_gpu=num_gpu,
            ),
            timeout=300.0,
        )

        # Embed existing facts once so each new fact can be matched against them.
        # A near-duplicate UPDATES the matched fact in place (edit with new info)
        # rather than appending a copy. Best effort: if embeddings are down we
        # fall back to plain append. Merge disabled unless 0 < threshold < 1.
        existing_embeds: dict = {}
        if results and 0 < merge_threshold < 1:
            vecs = await asyncio.gather(*(mgr.embed(it.text) for it in existing))
            existing_embeds = {it.id: v for it, v in zip(existing, vecs) if v}

        saved = []
        for result in results:
            new_vec = await mgr.embed(result["text"]) if existing_embeds else None
            match_id, best = None, 0.0
            if new_vec:
                for eid, ev in existing_embeds.items():
                    sim = _cosine(new_vec, ev)
                    if sim > best:
                        best, match_id = sim, eid
                if best < merge_threshold:
                    match_id = None

            target = store.get(match_id) if match_id else None
            if target:
                # Near-duplicate of an existing fact — overwrite with the newer
                # statement, keeping the original id/section/position.
                updated = MemoryItem(id=target.id, section=target.section,
                                     text=result["text"], tags=target.tags)
                store.update(updated)
                if new_vec:
                    existing_embeds[updated.id] = new_vec  # keep cache fresh for later facts in this batch
                saved.append({"id": updated.id, "section": updated.section,
                              "text": updated.text, "updated": True})
            else:
                item = MemoryItem(id=str(uuid.uuid4()),
                                  section=result["section"], text=result["text"])
                store.add(item)
                if new_vec:
                    existing_embeds[item.id] = new_vec
                saved.append({"id": item.id, "section": item.section, "text": item.text})
        if saved:
            first = {k: saved[0][k] for k in ("id", "section", "text")}
            return {"saved": True, "items": saved, **first}
    except asyncio.TimeoutError:
        pass
    except Exception:
        pass
    return {"saved": False}