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:
2026-06-25 14:50:53 +02:00
commit 82328b0a45
6 changed files with 301 additions and 0 deletions

107
store.py Normal file
View 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