Refactor: single collection with shared/private visibility
Breaking change: per-agent collections replaced by a single 'memories' collection with metadata-based access control. - All memories live in one collection - 'agent' metadata tracks the creator - 'private' (bool, default False) controls visibility - Query returns: shared (private=False) + own private (agent=asker) - Save endpoint accepts optional 'private' field - Version bumped to 0.2.0
This commit is contained in:
92
store.py
92
store.py
@@ -1,4 +1,22 @@
|
||||
"""ChromaDB-backed memory store with per-agent collections."""
|
||||
"""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
|
||||
|
||||
@@ -7,12 +25,14 @@ from chromadb.config import Settings
|
||||
|
||||
from embedder import Embedder
|
||||
|
||||
COLLECTION_NAME = "memories"
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
"""Persistent vector store using ChromaDB.
|
||||
|
||||
Each agent gets its own collection (``agent_{name}``) so memories
|
||||
are cleanly separated and can be queried independently.
|
||||
… with an optional ``private`` flag (default ``False``) so agents
|
||||
can keep sensitive memories to themselves.
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str) -> None:
|
||||
@@ -26,11 +46,11 @@ class MemoryStore:
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _get_collection(self, agent_name: str):
|
||||
"""Get or create a ChromaDB collection for *agent_name*."""
|
||||
def _collection(self):
|
||||
"""Return the single shared collection (create if missing)."""
|
||||
return self.client.get_or_create_collection(
|
||||
name=f"agent_{agent_name}",
|
||||
metadata={"agent": agent_name},
|
||||
name=COLLECTION_NAME,
|
||||
metadata={"description": "Shared + private agent memories"},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -43,18 +63,36 @@ class MemoryStore:
|
||||
text: str,
|
||||
metadata: dict | None = None,
|
||||
) -> str:
|
||||
"""Embed *text* and store it in the agent's collection.
|
||||
"""Embed and store *text*.
|
||||
|
||||
Returns the auto-generated memory ID.
|
||||
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._get_collection(agent_name)
|
||||
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=[metadata or {}],
|
||||
metadatas=[meta],
|
||||
ids=[memory_id],
|
||||
)
|
||||
return memory_id
|
||||
@@ -65,20 +103,36 @@ class MemoryStore:
|
||||
query_text: str,
|
||||
top_n: int = 5,
|
||||
) -> list[dict]:
|
||||
"""Return the *top_n* most semantically similar memories as text.
|
||||
"""Return semantically similar memories visible to *agent_name*.
|
||||
|
||||
Each result dict contains:
|
||||
- ``text`` – the stored memory text
|
||||
- ``metadata`` – any additional metadata (original ``text``
|
||||
field in metadata is stripped)
|
||||
- ``distance`` – cosine distance from the query
|
||||
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._get_collection(agent_name)
|
||||
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] = []
|
||||
@@ -94,8 +148,6 @@ class MemoryStore:
|
||||
if results.get("metadatas")
|
||||
else {}
|
||||
)
|
||||
# Don't return the raw document text in metadata (it's
|
||||
# already the top-level "text" field).
|
||||
meta.pop("text", None)
|
||||
|
||||
formatted.append({
|
||||
|
||||
Reference in New Issue
Block a user