NexusOS/synapse/memory/curator.py

"""Curated memory extraction — read a finished conversation, save what's new.

This is the whole memory-writing path: pick up the messages the curator has not
read yet, ask the model which permanent facts they contain, merge each result
against what is already stored, and move the conversation's watermark.

It runs in-process. Extraction used to sit behind an HTTP call to a separate
service on :8001 holding a second, smaller model, because that model could not
share the GPU with the chat model. The curator is the chat model now — already
resident, already warm — so the extra process bought nothing but a hop.
"""

from __future__ import annotations

import asyncio
import math
import uuid

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

_EXTRACT_TIMEOUT = 300.0


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


async def extract_for_conversation(conversation_id: str, project_id: str = "") -> list[dict]:
    """Extract and save facts from a conversation's unread messages.

    Returns the saved/updated items. Safe to call more than once: the watermark
    means a second run with no new messages does nothing and never reaches the
    model. Never raises — memory extraction must not take the caller down.
    """
    messages, last_id = store.pending_extraction(conversation_id)
    if not messages:
        return []

    # Only global facts and this project's are in scope: another project's facts
    # must not be shown to the curator as "already known", or as a merge target.
    existing = [m for m in store.all() if m.project_id in ("", project_id)]
    existing_sections = list({i.section for i in existing})
    existing_texts = [i.text for i in existing]

    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", -1), 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(
                messages, existing_sections, existing_texts, mgr,
                model=model, num_gpu=num_gpu,
            ),
            timeout=_EXTRACT_TIMEOUT,
        )
    except Exception:
        # Leave the watermark alone so the next sweep retries these messages.
        return []

    saved = []
    try:
        # 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}

        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,
                                     project_id=target.project_id)
                store.update(updated)
                if new_vec:
                    existing_embeds[updated.id] = new_vec
                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"], project_id=project_id)
                store.add(item)
                if new_vec:
                    existing_embeds[item.id] = new_vec
                saved.append({"id": item.id, "section": item.section, "text": item.text})
    except Exception:
        pass

    # Watermark even when nothing was saved — the model has read these messages
    # and re-reading them would just spend another call to reach the same "no".
    store.set_extracted_through(conversation_id, last_id)
    return saved