From 82328b0a451bac5208df6855c289dc7012bb2daf Mon Sep 17 00:00:00 2001 From: Lucy Date: Thu, 25 Jun 2026 14:50:53 +0200 Subject: [PATCH] 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 --- .gitignore | 30 +++++++++++++ config.py | 17 ++++++++ embedder.py | 29 +++++++++++++ main.py | 109 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 9 ++++ store.py | 107 ++++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 301 insertions(+) create mode 100644 .gitignore create mode 100644 config.py create mode 100644 embedder.py create mode 100644 main.py create mode 100644 requirements.txt create mode 100644 store.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b56e7d6 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/config.py b/config.py new file mode 100644 index 0000000..daf8105 --- /dev/null +++ b/config.py @@ -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")) diff --git a/embedder.py b/embedder.py new file mode 100644 index 0000000..0149dcd --- /dev/null +++ b/embedder.py @@ -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() diff --git a/main.py b/main.py new file mode 100644 index 0000000..5d368b9 --- /dev/null +++ b/main.py @@ -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, + ) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..41a5a9e --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/store.py b/store.py new file mode 100644 index 0000000..13ce082 --- /dev/null +++ b/store.py @@ -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