Model/tools/memory.py

"""
Frox AI — Memory Tool

A local-first, self-contained memory store: JSON-persisted facts per
user, retrieved by cosine similarity over embeddings from Morph's own
model (ctx.engine.embed() — see inference/engine/morph_engine.py).
No external vector DB required, so this works standalone in this repo.

The production backend architecture document specs a Postgres +
Qdrant version of this same idea for multi-instance deployments
(Section 5: Memory System) — this module is the reference
implementation / local-dev equivalent, not a competing design.
"""
from __future__ import annotations

import json
import math
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

from tools.registry import tool, ToolContext


@dataclass
class MemoryItem:
    id: str
    user_id: str
    text: str
    category: str
    embedding: List[float]
    created_at: float

    def to_dict(self, include_embedding: bool = False) -> dict:
        d = {"id": self.id, "user_id": self.user_id, "text": self.text,
             "category": self.category, "created_at": self.created_at}
        if include_embedding:
            d["embedding"] = self.embedding
        return d


def _cosine(a: List[float], b: List[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


class MemoryStore:
    """
    Per-process memory store, JSON-persisted to disk. Embeddings come
    from whatever MorphInferenceEngine is passed in at save time — the
    store itself has no model dependency, so it's easy to unit test
    with fake embeddings.
    """

    def __init__(self, path: Optional[str] = None):
        self.path = Path(path) if path else None
        self._items: Dict[str, MemoryItem] = {}
        if self.path and self.path.exists():
            self._load()

    def _load(self):
        try:
            data = json.loads(self.path.read_text())
            for item in data:
                self._items[item["id"]] = MemoryItem(**item)
        except (json.JSONDecodeError, TypeError, KeyError):
            pass   # start fresh rather than crash on a corrupted file

    def _persist(self):
        if self.path is None:
            return
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.path.write_text(json.dumps(
            [i.to_dict(include_embedding=True) for i in self._items.values()], indent=2
        ))

    def add(self, user_id: str, text: str, embedding: List[float], category: str = "fact") -> MemoryItem:
        item = MemoryItem(
            id=str(uuid.uuid4()), user_id=user_id, text=text,
            category=category, embedding=embedding, created_at=time.time(),
        )
        self._items[item.id] = item
        self._persist()
        return item

    def search(self, user_id: str, query_embedding: List[float], k: int = 5) -> List[MemoryItem]:
        candidates = [i for i in self._items.values() if i.user_id == user_id]
        scored = sorted(candidates, key=lambda i: _cosine(i.embedding, query_embedding), reverse=True)
        return scored[:k]

    def list_all(self, user_id: str) -> List[MemoryItem]:
        return [i for i in self._items.values() if i.user_id == user_id]

    def delete(self, memory_id: str) -> bool:
        if memory_id in self._items:
            del self._items[memory_id]
            self._persist()
            return True
        return False

    def clear(self, user_id: str):
        to_delete = [i.id for i in self._items.values() if i.user_id == user_id]
        for mid in to_delete:
            del self._items[mid]
        self._persist()


@tool(
    name="memory_save",
    description="Save a fact about the user for later conversations",
    timeout=10.0,
)
def memory_save(ctx: ToolContext, text: str, category: str = "fact") -> dict:
    """
    Args:
        text: The fact to remember, e.g. "User prefers Python over JavaScript".
        category: One of fact | preference | project | person (freeform, for organization).

    Plain `def`, not `async def`: engine.embed() is synchronous and
    GPU-bound — thread-offloaded by the registry.
    """
    if ctx.memory_store is None:
        raise RuntimeError("No memory_store configured in ToolContext")
    if ctx.engine is None:
        raise RuntimeError("No engine configured in ToolContext (needed to embed the memory)")

    embedding = ctx.engine.embed(text)
    item = ctx.memory_store.add(ctx.user_id, text, embedding, category)
    return {"saved": True, "memory_id": item.id, "text": text}


@tool(
    name="memory_search",
    description="Search the user's saved memories for relevant facts",
    timeout=10.0,
)
def memory_search(ctx: ToolContext, query: str, k: int = 5) -> dict:
    """
    Args:
        query: What to search for, e.g. "what programming language does the user like".
        k: Max number of memories to return.

    Plain `def`, not `async def`: engine.embed() is synchronous and
    GPU-bound — thread-offloaded by the registry.
    """
    if ctx.memory_store is None:
        raise RuntimeError("No memory_store configured in ToolContext")
    if ctx.engine is None:
        raise RuntimeError("No engine configured in ToolContext (needed to embed the query)")

    query_embedding = ctx.engine.embed(query)
    results = ctx.memory_store.search(ctx.user_id, query_embedding, k=k)
    return {"query": query, "memories": [r.to_dict() for r in results]}


@tool(
    name="memory_forget",
    description="Delete a previously saved memory by its id",
    timeout=5.0,
)
def memory_forget(ctx: ToolContext, memory_id: str) -> dict:
    """
    Args:
        memory_id: The id returned by memory_save.

    Plain `def` for consistency with memory_save/memory_search in this
    module — no async I/O here either.
    """
    if ctx.memory_store is None:
        raise RuntimeError("No memory_store configured in ToolContext")
    deleted = ctx.memory_store.delete(memory_id)
    return {"deleted": deleted, "memory_id": memory_id}