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:
30
.gitignore
vendored
Normal file
30
.gitignore
vendored
Normal file
@@ -0,0 +1,30 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
|
||||
# Virtual environment
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
|
||||
# IDE / editors
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# ChromaDB data (persisted vectors — sacred, not in repo)
|
||||
*.sqlite3
|
||||
chroma_data/
|
||||
|
||||
# Environment overrides
|
||||
.env
|
||||
.env.local
|
||||
17
config.py
Normal file
17
config.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""Configuration for the Vector Memory Server."""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
class Config:
|
||||
DATA_DIR = os.environ.get(
|
||||
"MEMORY_SERVER_DATA",
|
||||
"/home/admin/agent-dir/vector_memory"
|
||||
)
|
||||
EMBEDDING_MODEL = os.environ.get(
|
||||
"MEMORY_SERVER_MODEL",
|
||||
"all-MiniLM-L6-v2"
|
||||
)
|
||||
DEFAULT_TOP_N = int(os.environ.get("MEMORY_SERVER_TOP_N", "5"))
|
||||
HOST = os.environ.get("MEMORY_SERVER_HOST", "127.0.0.1")
|
||||
PORT = int(os.environ.get("MEMORY_SERVER_PORT", "8000"))
|
||||
29
embedder.py
Normal file
29
embedder.py
Normal file
@@ -0,0 +1,29 @@
|
||||
"""Embedding model wrapper — singleton pattern.
|
||||
|
||||
Uses sentence-transformers with all-MiniLM-L6-v2 (384-dim, ~80MB).
|
||||
Loaded once, reused across all requests.
|
||||
"""
|
||||
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from config import Config
|
||||
|
||||
|
||||
class Embedder:
|
||||
"""Thread-safe singleton wrapper for the embedding model."""
|
||||
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> "Embedder":
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.model = SentenceTransformer(Config.EMBEDDING_MODEL)
|
||||
|
||||
def embed(self, text: str) -> list[float]:
|
||||
return self.model.encode(text).tolist()
|
||||
|
||||
def embed_batch(self, texts: list[str]) -> list[list[float]]:
|
||||
return self.model.encode(texts).tolist()
|
||||
109
main.py
Normal file
109
main.py
Normal file
@@ -0,0 +1,109 @@
|
||||
"""Vector Memory Server — REST API.
|
||||
|
||||
Endpoints
|
||||
---------
|
||||
- ``POST /api/{agent_name}/save`` — Save a memory
|
||||
- ``POST /api/{agent_name}/query`` — Query memories by semantic similarity
|
||||
- ``GET /health`` — Health check
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from config import Config
|
||||
from store import MemoryStore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# App
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
store = MemoryStore(Config.DATA_DIR)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request / Response models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SaveRequest(BaseModel):
|
||||
text: str
|
||||
type: Optional[str] = "fact"
|
||||
tags: Optional[list[str]] = None
|
||||
date: Optional[str] = None
|
||||
|
||||
|
||||
class SaveResponse(BaseModel):
|
||||
success: bool
|
||||
id: str
|
||||
|
||||
|
||||
class QueryRequest(BaseModel):
|
||||
query: str
|
||||
top_n: Optional[int] = Config.DEFAULT_TOP_N
|
||||
|
||||
|
||||
class QueryResponse(BaseModel):
|
||||
results: list[dict]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@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."""
|
||||
try:
|
||||
metadata: dict = {}
|
||||
if request.type:
|
||||
metadata["type"] = request.type
|
||||
if request.tags:
|
||||
metadata["tags"] = ",".join(request.tags)
|
||||
if request.date:
|
||||
metadata["date"] = request.date
|
||||
|
||||
memory_id = store.save(agent_name, request.text, metadata)
|
||||
return SaveResponse(success=True, id=memory_id)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
|
||||
|
||||
@app.post("/api/{agent_name}/query", response_model=QueryResponse)
|
||||
async def query_memory(
|
||||
agent_name: str, request: QueryRequest
|
||||
) -> QueryResponse:
|
||||
"""Retrieve semantically similar memories for the given agent."""
|
||||
try:
|
||||
results = store.query(agent_name, request.query, request.top_n)
|
||||
return QueryResponse(results=results)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict:
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entrypoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(
|
||||
"main:app",
|
||||
host=Config.HOST,
|
||||
port=Config.PORT,
|
||||
reload=False,
|
||||
)
|
||||
9
requirements.txt
Normal file
9
requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
# Vector Memory Server
|
||||
# Install: pip install -r requirements.txt
|
||||
|
||||
fastapi>=0.115.0
|
||||
uvicorn[standard]>=0.34.0
|
||||
chromadb>=0.6.0
|
||||
sentence-transformers>=3.4.0
|
||||
numpy>=2.0.0
|
||||
pydantic>=2.0.0
|
||||
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