memory-server/store.py

"""ChromaDB-backed memory store with shared/private visibility.

All memories live in a single collection. Each memory has two key
metadata fields:

- ``agent`` (str) — the agent that created it
- ``private`` (bool) — ``False`` by default (visible to all agents)

When querying, the store returns:

1. **All shared memories** (``private=False``) — every agent sees these.
2. **Only the querying agent's own private memories** (``private=True``
   and ``agent=<asker>``).

This gives a clean hierarchy: shared knowledge by default, with an
opt-in privacy mechanism for per-agent internal notes.
"""

from __future__ import annotations

import uuid

import chromadb
from chromadb.config import Settings

from embedder import Embedder

COLLECTION_NAME = "memories"


class MemoryStore:
    """Persistent vector store using ChromaDB.

    … with an optional ``private`` flag (default ``False``) so agents
    can keep sensitive memories to themselves.
    """

    def __init__(self, data_dir: str) -> None:
        self.client = chromadb.PersistentClient(
            path=data_dir,
            settings=Settings(anonymized_telemetry=False),
        )
        self.embedder = Embedder.get_instance()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _collection(self):
        """Return the single shared collection (create if missing)."""
        return self.client.get_or_create_collection(
            name=COLLECTION_NAME,
            metadata={"description": "Shared + private agent memories"},
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def save(
        self,
        agent_name: str,
        text: str,
        metadata: dict | None = None,
    ) -> str:
        """Embed and store *text*.

        Parameters
        ----------
        agent_name
            The agent creating this memory (stored in metadata).
        text
            The memory content.
        metadata
            Optional extra fields.  If it contains a ``private`` key
            that will be passed to ChromaDB as-is; otherwise ``private``
            defaults to ``False``.  The ``agent`` key is always
            overwritten with *agent_name*.

        Returns
        -------
        The auto-generated memory UUID.
        """
        collection = self._collection()
        memory_id = str(uuid.uuid4())
        embedding = self.embedder.embed(text)

        meta: dict = dict(metadata or {})
        meta.setdefault("private", False)
        meta["agent"] = agent_name

        collection.add(
            embeddings=[embedding],
            documents=[text],
            metadatas=[meta],
            ids=[memory_id],
        )
        return memory_id

    def query(
        self,
        agent_name: str,
        query_text: str,
        top_n: int = 5,
    ) -> list[dict]:
        """Return semantically similar memories visible to *agent_name*.

        Visibility rules
        ----------------
        - All non‑private memories (``private=False``) are returned
          regardless of which agent created them.
        - Private memories (``private=True``) are only returned for the
          agent that owns them (``agent == agent_name``).

        Each result dict contains
        - ``text`` — the stored memory text
        - ``metadata`` — metadata (minus the raw ``text`` field)
        - ``distance`` — cosine distance from the query
        """
        collection = self._collection()
        query_embedding = self.embedder.embed(query_text)

        # ── ChromaDB where filter ──────────────────────────────────
        # Show shared memories + the asking agent's own private ones.
        where_filter = {
            "$or": [
                {"private": {"$eq": False}},
                {"agent": {"$eq": agent_name}},
            ]
        }

        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_n,
            where=where_filter,
        )

        formatted: list[dict] = []
        if results["documents"] and results["documents"][0]:
            for i, doc in enumerate(results["documents"][0]):
                distance = (
                    results["distances"][0][i]
                    if results.get("distances")
                    else None
                )
                meta = (
                    {k: v for k, v in results["metadatas"][0][i].items()}
                    if results.get("metadatas")
                    else {}
                )
                meta.pop("text", None)

                formatted.append({
                    "text": doc,
                    "metadata": meta,
                    "distance": distance,
                })

        return formatted