Breaking change: per-agent collections replaced by a single 'memories' collection with metadata-based access control. - All memories live in one collection - 'agent' metadata tracks the creator - 'private' (bool, default False) controls visibility - Query returns: shared (private=False) + own private (agent=asker) - Save endpoint accepts optional 'private' field - Version bumped to 0.2.0
124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
"""Vector Memory Server — REST API.
|
|
|
|
Endpoints
|
|
---------
|
|
- ``POST /api/{agent_name}/save`` — Save a memory (shared by default,
|
|
or private if ``private: true``)
|
|
- ``POST /api/{agent_name}/query`` — Query memories visible to the agent
|
|
(all shared + its own private ones)
|
|
- ``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. "
|
|
"Shared-by-default, private-opt-in. "
|
|
"ChromaDB + all-MiniLM-L6-v2.",
|
|
version="0.2.0",
|
|
)
|
|
|
|
store = MemoryStore(Config.DATA_DIR)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request / Response models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class SaveRequest(BaseModel):
|
|
text: str
|
|
private: bool = False
|
|
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 *agent_name*.
|
|
|
|
By default the memory is **shared** (visible to all agents).
|
|
Set ``private: true`` to restrict visibility to *agent_name* only.
|
|
"""
|
|
try:
|
|
metadata: dict = {
|
|
"private": request.private,
|
|
}
|
|
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 memories visible to *agent_name*.
|
|
|
|
Returns all **shared** memories from any agent, plus **private**
|
|
memories owned by *agent_name*.
|
|
"""
|
|
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,
|
|
)
|