Files
pi-qmem/gateway/main.py
T

294 lines
11 KiB
Python

"""
Memory Gateway — memoria centralizzata condivisa per agenti AI.
Stack snello: FastAPI + Qdrant + Ollama (BGE-M3). Nessun LLM in scrittura.
Accesso: UNA o più API key condivise con accesso COMPLETO in lettura e
scrittura all'intera conoscenza. Nessun isolamento per agente: qualsiasi
agente (attuale o futuro) con la chiave può consultare e aggiungere
informazioni liberamente. L'agent_id è solo metadata di provenienza.
Endpoints:
POST /v1/memories → crea un record di memoria
POST /v1/memories:search → ricerca semantica con filtri
GET /v1/memories/{id} → recupera per UUID
DELETE /v1/memories/{id} → elimina per UUID
GET /v1/status → health + statistiche
"""
import asyncio
import json
import logging
import os
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Literal, Optional
import httpx
import uvicorn
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from pydantic import BaseModel, Field
from qdrant_client import QdrantClient
from qdrant_client.http import models as qm
# ---------------------------------------------------------------------------
# Configurazione (env)
# ---------------------------------------------------------------------------
QDRANT_URL = os.environ.get("QDRANT_URL", "http://127.0.0.1:6333")
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
EMBED_MODEL = os.environ.get("EMBED_MODEL", "bge-m3")
EMBED_DIM = int(os.environ.get("EMBED_DIM", "1024"))
COLLECTION = os.environ.get("COLLECTION", "memories")
# Chiavi condivise (separate da virgola): accesso completo in lettura/scrittura
API_KEYS: set[str] = {k.strip() for k in os.environ.get("API_KEYS", "").split(",") if k.strip()}
RATE_LIMIT_PER_MIN = int(os.environ.get("RATE_LIMIT_PER_MIN", "120"))
MAX_TEXT_LEN = int(os.environ.get("MAX_TEXT_LEN", "8000"))
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("memory-gateway")
app = FastAPI(title="Memory Gateway", version="2.0.0")
qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
# Rate limit in-memory: {key: [timestamps]}
_ratelimit: dict[str, list[float]] = {}
# ---------------------------------------------------------------------------
# Modelli
# ---------------------------------------------------------------------------
class MemoryIn(BaseModel):
text: str = Field(min_length=1, max_length=MAX_TEXT_LEN)
kind: Literal["decision", "fact", "episode", "preference"] = "fact"
agent_id: Optional[str] = Field(default=None, max_length=64, description="Solo provenienza, nessun isolamento")
project_id: Optional[str] = Field(default=None, max_length=64)
scope: Literal["agent", "project", "org"] = "agent"
source: Optional[str] = Field(default=None, max_length=256)
expires_at: Optional[str] = None # ISO 8601
supersedes_id: Optional[str] = None
class SearchIn(BaseModel):
query: str = Field(min_length=1, max_length=512)
kind: Optional[Literal["decision", "fact", "episode", "preference"]] = None
project_id: Optional[str] = None
scope: Optional[Literal["agent", "project", "org"]] = None
top_k: int = Field(default=5, ge=1, le=20)
# ---------------------------------------------------------------------------
# Auth: chiave condivisa → accesso completo (nessun isolamento)
# ---------------------------------------------------------------------------
def require_auth(x_api_key: str = Header(...)) -> str:
if x_api_key not in API_KEYS:
raise HTTPException(status_code=401, detail="API key non valida")
# rate limit per chiave
now = time.monotonic()
window = _ratelimit.setdefault(x_api_key, [])
window[:] = [t for t in window if now - t < 60]
if len(window) >= RATE_LIMIT_PER_MIN:
raise HTTPException(status_code=429, detail="Rate limit superato")
window.append(now)
return x_api_key
def _audit(key: str, action: str, **extra: Any) -> None:
"""Audit log in JSON lines (catturato da docker logs)."""
entry = {
"ts": datetime.now(timezone.utc).isoformat(),
"key": key[:8] + "...",
"action": action,
**extra,
}
log.info(json.dumps(entry, default=str))
def _now_iso() -> str:
return datetime.now(timezone.utc).isoformat()
def _parse_ts(value: Optional[str]) -> Optional[float]:
"""Converte ISO 8601 in timestamp Unix (per i range query Qdrant)."""
if not value:
return None
try:
return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
except ValueError:
return None
# ---------------------------------------------------------------------------
# Embedding via Ollama (BGE-M3)
# ---------------------------------------------------------------------------
async def embed(text: str) -> list[float]:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(
f"{OLLAMA_URL}/api/embed",
json={"model": EMBED_MODEL, "input": text},
)
r.raise_for_status()
return r.json()["embeddings"][0]
# ---------------------------------------------------------------------------
# Startup: crea collection e indici se non esistono
# ---------------------------------------------------------------------------
@app.on_event("startup")
def startup() -> None:
collections = qdrant.get_collections().collections
if not any(c.name == COLLECTION for c in collections):
qdrant.create_collection(
collection_name=COLLECTION,
vectors_config=qm.VectorParams(size=EMBED_DIM, distance=qm.Distance.COSINE),
)
for field in ("agent_id", "project_id", "scope", "kind"):
qdrant.create_payload_index(
collection_name=COLLECTION,
field_name=field,
field_schema=qm.PayloadSchemaType.KEYWORD,
)
log.info("collection %s creata con indici", COLLECTION)
else:
log.info("collection %s già esistente", COLLECTION)
# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------
@app.post("/v1/memories")
async def add_memory(body: MemoryIn, key: str = Depends(require_auth)) -> dict:
memory_id = str(uuid.uuid4())
vector = await embed(body.text)
payload: dict[str, Any] = {
"text": body.text,
"kind": body.kind,
"agent_id": body.agent_id or "shared",
"project_id": body.project_id,
"scope": body.scope,
"source": body.source,
"created_at": _now_iso(),
"expires_at": _parse_ts(body.expires_at),
"supersedes_id": body.supersedes_id,
"embedding_model": EMBED_MODEL,
}
qdrant.upsert(
collection_name=COLLECTION,
points=[qm.PointStruct(id=memory_id, vector=vector, payload=payload)],
)
_audit(key, "create", memory_id=memory_id, kind=body.kind, agent_id=payload["agent_id"])
return {"memory_id": memory_id, "created_at": payload["created_at"]}
@app.post("/v1/memories:search")
async def search_memories(body: SearchIn, key: str = Depends(require_auth)) -> dict:
vector = await embed(body.query)
must: list[Any] = []
if body.kind:
must.append(qm.FieldCondition(key="kind", match=qm.MatchValue(value=body.kind)))
if body.project_id:
must.append(qm.FieldCondition(key="project_id", match=qm.MatchValue(value=body.project_id)))
if body.scope:
must.append(qm.FieldCondition(key="scope", match=qm.MatchValue(value=body.scope)))
hits = qdrant.search(
collection_name=COLLECTION,
query_vector=vector,
query_filter=qm.Filter(must=must) if must else None,
limit=body.top_k,
with_payload=True,
)
results = [
{
"memory_id": h.id,
"score": round(h.score, 4),
"text": h.payload.get("text"),
"kind": h.payload.get("kind"),
"agent_id": h.payload.get("agent_id"),
"scope": h.payload.get("scope"),
"project_id": h.payload.get("project_id"),
"created_at": h.payload.get("created_at"),
"source": h.payload.get("source"),
}
for h in hits
]
_audit(key, "search", query=body.query[:80], top_k=body.top_k, hits=len(results))
return {"results": results}
@app.get("/v1/memories/{memory_id}")
async def get_memory(memory_id: str, key: str = Depends(require_auth)) -> dict:
point = qdrant.retrieve(
collection_name=COLLECTION, ids=[memory_id], with_payload=True
)
if not point:
raise HTTPException(status_code=404, detail="Memoria non trovata")
_audit(key, "get", memory_id=memory_id)
return {"memory_id": memory_id, **point[0].payload}
@app.delete("/v1/memories/{memory_id}")
async def delete_memory(memory_id: str, key: str = Depends(require_auth)) -> dict:
point = qdrant.retrieve(
collection_name=COLLECTION, ids=[memory_id], with_payload=True
)
if not point:
raise HTTPException(status_code=404, detail="Memoria non trovata")
qdrant.delete(collection_name=COLLECTION, points_selector=[memory_id])
_audit(key, "delete", memory_id=memory_id)
return {"deleted": memory_id}
@app.get("/v1/status")
async def status() -> dict:
info = qdrant.get_collection(COLLECTION)
return {
"status": "ok",
"collection": COLLECTION,
"points": info.points_count,
"embedding_model": EMBED_MODEL,
"embedding_dim": EMBED_DIM,
"access": "shared",
"api_keys": len(API_KEYS),
}
# ---------------------------------------------------------------------------
# Cleanup periodico: rimuove record scaduti (expires_at < now)
# ---------------------------------------------------------------------------
async def _cleanup_loop() -> None:
while True:
try:
now = time.time()
scroll = qdrant.scroll(
collection_name=COLLECTION,
scroll_filter=qm.Filter(
must=[
qm.FieldCondition(
key="expires_at",
range=qm.Range(lt=now),
)
]
),
limit=100,
with_payload=False,
)
ids = [p.id for p in scroll[0]]
if ids:
qdrant.delete(collection_name=COLLECTION, points_selector=ids)
log.info("cleanup: rimossi %d record scaduti", len(ids))
except Exception as e: # noqa: BLE001
log.warning("cleanup error: %s", e)
await asyncio.sleep(3600)
@app.on_event("startup")
async def start_cleanup() -> None:
asyncio.create_task(_cleanup_loop())
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8080)