pi-qmem/gateway/main.py

"""
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 (con supersedes_id corregge un record esistente)
  POST   /v1/memories:search    → ricerca semantica con filtri (include_superseded per la lineage)
  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.1.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
    supersede_reason: Optional[str] = Field(default=None, max_length=512)


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
    include_superseded: bool = False
    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", "supersedes_id", "superseded_by"):
            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())

    # Supersede: il nuovo record corregge uno esistente, che resta in archivio marcato
    superseded_id: Optional[str] = None
    if body.supersedes_id:
        old = qdrant.retrieve(collection_name=COLLECTION, ids=[body.supersedes_id], with_payload=True)
        if not old:
            raise HTTPException(status_code=404, detail="Memoria da supersedere non trovata")
        if old[0].payload.get("superseded_by"):
            raise HTTPException(status_code=409, detail="La memoria è già stata superseduta: correggi la versione attiva")
        superseded_id = body.supersedes_id

    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": superseded_id,
        "supersede_reason": body.supersede_reason,
        "embedding_model": EMBED_MODEL,
    }
    qdrant.upsert(
        collection_name=COLLECTION,
        points=[qm.PointStruct(id=memory_id, vector=vector, payload=payload)],
    )

    if superseded_id:
        qdrant.set_payload(
            collection_name=COLLECTION,
            payload={
                "superseded_by": memory_id,
                "superseded_at": _now_iso(),
                "supersede_reason": body.supersede_reason,
            },
            points=[superseded_id],
        )
        _audit(key, "supersede", old_id=superseded_id, new_id=memory_id, kind=body.kind, agent_id=payload["agent_id"])
    else:
        _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"], "supersedes_id": superseded_id}


@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)))
    if not body.include_superseded:
        # default: esclude i record già corretti (superseded_by presente)
        must.append(qm.IsEmptyCondition(is_empty=qm.PayloadField(key="superseded_by")))

    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"),
            "supersedes_id": h.payload.get("supersedes_id"),
            "superseded_by": h.payload.get("superseded_by"),
            "supersede_reason": h.payload.get("supersede_reason"),
        }
        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)