diff --git a/main.py b/main.py index 5d368b9..63d5d88 100644 --- a/main.py +++ b/main.py @@ -2,8 +2,10 @@ Endpoints --------- -- ``POST /api/{agent_name}/save`` — Save a memory -- ``POST /api/{agent_name}/query`` — Query memories by semantic similarity +- ``POST /api/{agent_name}/save`` — Save a memory (shared by default, + or private if ``private: true``) +- ``POST /api/{agent_name}/query`` — Query memories visible to the agent + (all shared + its own private ones) - ``GET /health`` — Health check """ @@ -25,8 +27,9 @@ from store import MemoryStore app = FastAPI( title="Vector Memory Server", description="Semantic memory retrieval via REST. " - "Per-agent collections, powered by ChromaDB + all-MiniLM-L6-v2.", - version="0.1.0", + "Shared-by-default, private-opt-in. " + "ChromaDB + all-MiniLM-L6-v2.", + version="0.2.0", ) store = MemoryStore(Config.DATA_DIR) @@ -38,6 +41,7 @@ store = MemoryStore(Config.DATA_DIR) class SaveRequest(BaseModel): text: str + private: bool = False type: Optional[str] = "fact" tags: Optional[list[str]] = None date: Optional[str] = None @@ -63,9 +67,15 @@ class QueryResponse(BaseModel): @app.post("/api/{agent_name}/save", response_model=SaveResponse) async def save_memory(agent_name: str, request: SaveRequest) -> SaveResponse: - """Embed and store a memory for the given agent.""" + """Embed and store a memory for *agent_name*. + + By default the memory is **shared** (visible to all agents). + Set ``private: true`` to restrict visibility to *agent_name* only. + """ try: - metadata: dict = {} + metadata: dict = { + "private": request.private, + } if request.type: metadata["type"] = request.type if request.tags: @@ -83,7 +93,11 @@ async def save_memory(agent_name: str, request: SaveRequest) -> SaveResponse: async def query_memory( agent_name: str, request: QueryRequest ) -> QueryResponse: - """Retrieve semantically similar memories for the given agent.""" + """Retrieve memories visible to *agent_name*. + + Returns all **shared** memories from any agent, plus **private** + memories owned by *agent_name*. + """ try: results = store.query(agent_name, request.query, request.top_n) return QueryResponse(results=results) diff --git a/store.py b/store.py index 13ce082..93771c7 100644 --- a/store.py +++ b/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=``). + +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({