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:
28
main.py
28
main.py
@@ -2,8 +2,10 @@
|
|||||||
|
|
||||||
Endpoints
|
Endpoints
|
||||||
---------
|
---------
|
||||||
- ``POST /api/{agent_name}/save`` — Save a memory
|
- ``POST /api/{agent_name}/save`` — Save a memory (shared by default,
|
||||||
- ``POST /api/{agent_name}/query`` — Query memories by semantic similarity
|
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
|
- ``GET /health`` — Health check
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -25,8 +27,9 @@ from store import MemoryStore
|
|||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="Vector Memory Server",
|
title="Vector Memory Server",
|
||||||
description="Semantic memory retrieval via REST. "
|
description="Semantic memory retrieval via REST. "
|
||||||
"Per-agent collections, powered by ChromaDB + all-MiniLM-L6-v2.",
|
"Shared-by-default, private-opt-in. "
|
||||||
version="0.1.0",
|
"ChromaDB + all-MiniLM-L6-v2.",
|
||||||
|
version="0.2.0",
|
||||||
)
|
)
|
||||||
|
|
||||||
store = MemoryStore(Config.DATA_DIR)
|
store = MemoryStore(Config.DATA_DIR)
|
||||||
@@ -38,6 +41,7 @@ store = MemoryStore(Config.DATA_DIR)
|
|||||||
|
|
||||||
class SaveRequest(BaseModel):
|
class SaveRequest(BaseModel):
|
||||||
text: str
|
text: str
|
||||||
|
private: bool = False
|
||||||
type: Optional[str] = "fact"
|
type: Optional[str] = "fact"
|
||||||
tags: Optional[list[str]] = None
|
tags: Optional[list[str]] = None
|
||||||
date: Optional[str] = None
|
date: Optional[str] = None
|
||||||
@@ -63,9 +67,15 @@ class QueryResponse(BaseModel):
|
|||||||
|
|
||||||
@app.post("/api/{agent_name}/save", response_model=SaveResponse)
|
@app.post("/api/{agent_name}/save", response_model=SaveResponse)
|
||||||
async def save_memory(agent_name: str, request: SaveRequest) -> 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:
|
try:
|
||||||
metadata: dict = {}
|
metadata: dict = {
|
||||||
|
"private": request.private,
|
||||||
|
}
|
||||||
if request.type:
|
if request.type:
|
||||||
metadata["type"] = request.type
|
metadata["type"] = request.type
|
||||||
if request.tags:
|
if request.tags:
|
||||||
@@ -83,7 +93,11 @@ async def save_memory(agent_name: str, request: SaveRequest) -> SaveResponse:
|
|||||||
async def query_memory(
|
async def query_memory(
|
||||||
agent_name: str, request: QueryRequest
|
agent_name: str, request: QueryRequest
|
||||||
) -> QueryResponse:
|
) -> 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:
|
try:
|
||||||
results = store.query(agent_name, request.query, request.top_n)
|
results = store.query(agent_name, request.query, request.top_n)
|
||||||
return QueryResponse(results=results)
|
return QueryResponse(results=results)
|
||||||
|
|||||||
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
|
import uuid
|
||||||
|
|
||||||
@@ -7,12 +25,14 @@ from chromadb.config import Settings
|
|||||||
|
|
||||||
from embedder import Embedder
|
from embedder import Embedder
|
||||||
|
|
||||||
|
COLLECTION_NAME = "memories"
|
||||||
|
|
||||||
|
|
||||||
class MemoryStore:
|
class MemoryStore:
|
||||||
"""Persistent vector store using ChromaDB.
|
"""Persistent vector store using ChromaDB.
|
||||||
|
|
||||||
Each agent gets its own collection (``agent_{name}``) so memories
|
… with an optional ``private`` flag (default ``False``) so agents
|
||||||
are cleanly separated and can be queried independently.
|
can keep sensitive memories to themselves.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, data_dir: str) -> None:
|
def __init__(self, data_dir: str) -> None:
|
||||||
@@ -26,11 +46,11 @@ class MemoryStore:
|
|||||||
# Internal helpers
|
# Internal helpers
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def _get_collection(self, agent_name: str):
|
def _collection(self):
|
||||||
"""Get or create a ChromaDB collection for *agent_name*."""
|
"""Return the single shared collection (create if missing)."""
|
||||||
return self.client.get_or_create_collection(
|
return self.client.get_or_create_collection(
|
||||||
name=f"agent_{agent_name}",
|
name=COLLECTION_NAME,
|
||||||
metadata={"agent": agent_name},
|
metadata={"description": "Shared + private agent memories"},
|
||||||
)
|
)
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
@@ -43,18 +63,36 @@ class MemoryStore:
|
|||||||
text: str,
|
text: str,
|
||||||
metadata: dict | None = None,
|
metadata: dict | None = None,
|
||||||
) -> str:
|
) -> 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())
|
memory_id = str(uuid.uuid4())
|
||||||
embedding = self.embedder.embed(text)
|
embedding = self.embedder.embed(text)
|
||||||
|
|
||||||
|
meta: dict = dict(metadata or {})
|
||||||
|
meta.setdefault("private", False)
|
||||||
|
meta["agent"] = agent_name
|
||||||
|
|
||||||
collection.add(
|
collection.add(
|
||||||
embeddings=[embedding],
|
embeddings=[embedding],
|
||||||
documents=[text],
|
documents=[text],
|
||||||
metadatas=[metadata or {}],
|
metadatas=[meta],
|
||||||
ids=[memory_id],
|
ids=[memory_id],
|
||||||
)
|
)
|
||||||
return memory_id
|
return memory_id
|
||||||
@@ -65,20 +103,36 @@ class MemoryStore:
|
|||||||
query_text: str,
|
query_text: str,
|
||||||
top_n: int = 5,
|
top_n: int = 5,
|
||||||
) -> list[dict]:
|
) -> list[dict]:
|
||||||
"""Return the *top_n* most semantically similar memories as text.
|
"""Return semantically similar memories visible to *agent_name*.
|
||||||
|
|
||||||
Each result dict contains:
|
Visibility rules
|
||||||
- ``text`` – the stored memory text
|
----------------
|
||||||
- ``metadata`` – any additional metadata (original ``text``
|
- All non‑private memories (``private=False``) are returned
|
||||||
field in metadata is stripped)
|
regardless of which agent created them.
|
||||||
- ``distance`` – cosine distance from the query
|
- 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)
|
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(
|
results = collection.query(
|
||||||
query_embeddings=[query_embedding],
|
query_embeddings=[query_embedding],
|
||||||
n_results=top_n,
|
n_results=top_n,
|
||||||
|
where=where_filter,
|
||||||
)
|
)
|
||||||
|
|
||||||
formatted: list[dict] = []
|
formatted: list[dict] = []
|
||||||
@@ -94,8 +148,6 @@ class MemoryStore:
|
|||||||
if results.get("metadatas")
|
if results.get("metadatas")
|
||||||
else {}
|
else {}
|
||||||
)
|
)
|
||||||
# Don't return the raw document text in metadata (it's
|
|
||||||
# already the top-level "text" field).
|
|
||||||
meta.pop("text", None)
|
meta.pop("text", None)
|
||||||
|
|
||||||
formatted.append({
|
formatted.append({
|
||||||
|
|||||||
Reference in New Issue
Block a user