FusionAGI/fusionagi/memory/service.py

"""Unified memory service: session, episodic, semantic, vector with tenant isolation."""

from typing import Any

from fusionagi.memory.episodic import EpisodicMemory
from fusionagi.memory.semantic import SemanticMemory
from fusionagi.memory.working import WorkingMemory


def _scoped_key(tenant_id: str, user_id: str, base: str) -> str:
    """Scope key by tenant and user."""
    parts = [tenant_id or "default", user_id or "anonymous", base]
    return ":".join(parts)


class VectorMemory:
    """
    Vector memory for embeddings retrieval.

    Uses in-memory cosine similarity search. For production, swap with
    pgvector, Pinecone, or Qdrant adapter behind the same interface.
    """

    def __init__(self, max_entries: int = 10000) -> None:
        self._store: list[dict[str, Any]] = []
        self._max_entries = max_entries

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        """Compute cosine similarity between two vectors."""
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = sum(x * x for x in a) ** 0.5
        norm_b = sum(x * x for x in b) ** 0.5
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def add(self, id: str, embedding: list[float], metadata: dict[str, Any] | None = None) -> None:
        """Add embedding to the vector store."""
        if len(self._store) >= self._max_entries:
            self._store.pop(0)
        self._store.append({"id": id, "embedding": embedding, "metadata": metadata or {}})

    def search(self, query_embedding: list[float], top_k: int = 10) -> list[dict[str, Any]]:
        """Search by cosine similarity, returning top-k results."""
        scored = []
        for entry in self._store:
            sim = self._cosine_similarity(query_embedding, entry["embedding"])
            scored.append({"id": entry["id"], "metadata": entry["metadata"], "score": sim})
        scored.sort(key=lambda x: x["score"], reverse=True)
        return scored[:top_k]

    def delete(self, id: str) -> bool:
        """Remove an entry by ID."""
        before = len(self._store)
        self._store = [e for e in self._store if e["id"] != id]
        return len(self._store) < before

    def count(self) -> int:
        """Return entry count."""
        return len(self._store)


class MemoryService:
    """
    Unified memory service with tenant isolation.
    Wraps WorkingMemory (session), EpisodicMemory, SemanticMemory, VectorMemory.
    """

    def __init__(
        self,
        tenant_id: str = "default",
        user_id: str | None = None,
    ) -> None:
        self._tenant_id = tenant_id
        self._user_id = user_id or "anonymous"
        self._working = WorkingMemory()
        self._episodic = EpisodicMemory()
        self._semantic = SemanticMemory()
        self._vector = VectorMemory()

    @property
    def session(self) -> WorkingMemory:
        """Short-term session memory."""
        return self._working

    @property
    def episodic(self) -> EpisodicMemory:
        """Episodic memory (what happened, decisions, outcomes)."""
        return self._episodic

    @property
    def semantic(self) -> SemanticMemory:
        """Semantic memory (facts, preferences)."""
        return self._semantic

    @property
    def vector(self) -> VectorMemory:
        """Vector memory (embeddings for retrieval)."""
        return self._vector

    def scope_session(self, session_id: str) -> str:
        """Return tenant/user scoped session key."""
        return _scoped_key(self._tenant_id, self._user_id, session_id)

    def get(self, session_id: str, key: str, default: Any = None) -> Any:
        """Get from session memory (scoped)."""
        scoped = self.scope_session(session_id)
        return self._working.get(scoped, key, default)

    def set(self, session_id: str, key: str, value: Any) -> None:
        """Set in session memory (scoped)."""
        scoped = self.scope_session(session_id)
        self._working.set(scoped, key, value)

    def append_episode(self, task_id: str, event: dict[str, Any], event_type: str | None = None) -> int:
        """Append to episodic memory (with tenant in metadata)."""
        event = dict(event)
        meta = event.setdefault("metadata", {})
        meta = dict(meta) if meta else {}
        meta["tenant_id"] = self._tenant_id
        meta["user_id"] = self._user_id
        event["metadata"] = meta
        return self._episodic.append(task_id, event, event_type)