Initial commit: standalone vector memory server
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
This commit is contained in:
107
store.py
Normal file
107
store.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user