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:
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,
|
||||
)
|
||||
Reference in New Issue
Block a user