Lucy
82328b0a45
REST API for per-agent semantic memory retrieval.
- FastAPI server with /api/{agent}/save and /api/{agent}/query
- ChromaDB for persistent vector storage
- all-MiniLM-L6-v2 via sentence-transformers for embeddings
- Per-agent collections for clean separation
- Config through env vars or config.py
- .venv ready with all dependencies
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
"""ChromaDB-backed memory store with per-agent collections."""
|
||
|
||
import uuid
|
||
|
||
import chromadb
|
||
from chromadb.config import Settings
|
||
|
||
from embedder import Embedder
|
||
|
||
|
||
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.
|
||
"""
|
||
|
||
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 _get_collection(self, agent_name: str):
|
||
"""Get or create a ChromaDB collection for *agent_name*."""
|
||
return self.client.get_or_create_collection(
|
||
name=f"agent_{agent_name}",
|
||
metadata={"agent": agent_name},
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public API
|
||
# ------------------------------------------------------------------
|
||
|
||
def save(
|
||
self,
|
||
agent_name: str,
|
||
text: str,
|
||
metadata: dict | None = None,
|
||
) -> str:
|
||
"""Embed *text* and store it in the agent's collection.
|
||
|
||
Returns the auto-generated memory ID.
|
||
"""
|
||
collection = self._get_collection(agent_name)
|
||
memory_id = str(uuid.uuid4())
|
||
embedding = self.embedder.embed(text)
|
||
|
||
collection.add(
|
||
embeddings=[embedding],
|
||
documents=[text],
|
||
metadatas=[metadata or {}],
|
||
ids=[memory_id],
|
||
)
|
||
return memory_id
|
||
|
||
def query(
|
||
self,
|
||
agent_name: str,
|
||
query_text: str,
|
||
top_n: int = 5,
|
||
) -> list[dict]:
|
||
"""Return the *top_n* most semantically similar memories as text.
|
||
|
||
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
|
||
"""
|
||
collection = self._get_collection(agent_name)
|
||
query_embedding = self.embedder.embed(query_text)
|
||
|
||
results = collection.query(
|
||
query_embeddings=[query_embedding],
|
||
n_results=top_n,
|
||
)
|
||
|
||
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 {}
|
||
)
|
||
# Don't return the raw document text in metadata (it's
|
||
# already the top-level "text" field).
|
||
meta.pop("text", None)
|
||
|
||
formatted.append({
|
||
"text": doc,
|
||
"metadata": meta,
|
||
"distance": distance,
|
||
})
|
||
|
||
return formatted
|