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
110 lines
3.0 KiB
Python
110 lines
3.0 KiB
Python
"""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,
|
|
)
|