387 lines
14 KiB
Python
387 lines
14 KiB
Python
"""
|
|
RAG Permanence v2 — Cloudflare R2 cold storage via REST API.
|
|
Hot (Redis) → Warm (local 7-day cache) → Cold (R2 permanent).
|
|
|
|
R2 free tier: 10GB storage, 10M Class A ops/month, zero egress.
|
|
Uses CF REST API with Bearer token — no S3 credentials needed.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# CONFIG
|
|
# ═══════════════════════════════════════════════════════════════
|
|
R2_ACCOUNT_ID = os.getenv("R2_ACCOUNT_ID", "8f9bd9165c1250b426c66dc1967deefd")
|
|
R2_BUCKET = os.getenv("R2_BUCKET", "rag-backup")
|
|
R2_API_TOKEN = os.getenv("R2_API_TOKEN", "")
|
|
R2_BASE = f"https://api.cloudflare.com/client/v4/accounts/{R2_ACCOUNT_ID}/r2/buckets/{R2_BUCKET}"
|
|
|
|
LOCAL_CACHE = Path(os.getenv("RAG_LOCAL_CACHE", "/data/rag-storage"))
|
|
LOCAL_RETENTION_DAYS = 7
|
|
EMBEDDING_VERSION = os.getenv("EMBEDDING_MODEL_VERSION", "v1")
|
|
|
|
LOCAL_CACHE.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def _headers():
|
|
return {"Authorization": f"Bearer {R2_API_TOKEN}", "Content-Type": "application/json"}
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# R2 REST API OPERATIONS
|
|
# ═══════════════════════════════════════════════════════════════
|
|
async def _r2_put(key: str, data: bytes, content_type: str = "application/json") -> bool:
|
|
"""Upload an object to R2 via REST API."""
|
|
if not R2_API_TOKEN:
|
|
logger.error("R2_API_TOKEN not set")
|
|
return False
|
|
url = f"{R2_BASE}/objects/{key}"
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
try:
|
|
resp = await client.put(
|
|
url,
|
|
content=data,
|
|
headers={"Authorization": f"Bearer {R2_API_TOKEN}", "Content-Type": content_type},
|
|
)
|
|
if resp.status_code == 200:
|
|
return True
|
|
logger.error(f"R2 PUT {key}: {resp.status_code} {resp.text[:200]}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"R2 PUT {key}: {e}")
|
|
return False
|
|
|
|
|
|
async def _r2_get(key: str) -> bytes | None:
|
|
"""Download an object from R2 via REST API."""
|
|
if not R2_API_TOKEN:
|
|
return None
|
|
url = f"{R2_BASE}/objects/{key}"
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
try:
|
|
resp = await client.get(url, headers=_headers())
|
|
if resp.status_code == 200:
|
|
return resp.content
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"R2 GET {key}: {e}")
|
|
return None
|
|
|
|
|
|
async def _r2_list(prefix: str = "") -> list:
|
|
"""List objects in R2 bucket."""
|
|
if not R2_API_TOKEN:
|
|
return []
|
|
url = f"{R2_BASE}/objects"
|
|
if prefix:
|
|
url += f"?prefix={prefix}"
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
try:
|
|
resp = await client.get(url, headers=_headers())
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
return data.get("result", [])
|
|
return []
|
|
except Exception as e:
|
|
logger.error(f"R2 LIST: {e}")
|
|
return []
|
|
|
|
|
|
async def _r2_delete(key: str) -> bool:
|
|
"""Delete an object from R2."""
|
|
if not R2_API_TOKEN:
|
|
return False
|
|
url = f"{R2_BASE}/objects/{key}"
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
try:
|
|
resp = await client.delete(url, headers=_headers())
|
|
return resp.status_code in (200, 204, 404)
|
|
except Exception as e:
|
|
logger.error(f"R2 DELETE {key}: {e}")
|
|
return False
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# RAG PERSISTENCE OPERATIONS
|
|
# ═══════════════════════════════════════════════════════════════
|
|
async def snapshot_all():
|
|
"""Snapshot all RAG collections to R2. Returns per-collection status."""
|
|
results = {}
|
|
ts = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
|
|
|
|
# Get collections from rag_service (Redis-backed)
|
|
try:
|
|
from app.rag_service import COLLECTIONS, _get_redis
|
|
|
|
redis = await _get_redis()
|
|
has_redis = True
|
|
except Exception:
|
|
COLLECTIONS = []
|
|
has_redis = False
|
|
|
|
for name in COLLECTIONS:
|
|
try:
|
|
# Get document IDs from Redis
|
|
if not has_redis:
|
|
results[name] = {"collection": name, "status": "skipped", "count": 0}
|
|
continue
|
|
|
|
doc_ids = await redis.smembers(f"rag:idx:{name}")
|
|
count = len(doc_ids) if doc_ids else 0
|
|
if count == 0:
|
|
results[name] = {"collection": name, "status": "empty", "count": 0}
|
|
continue
|
|
|
|
# Collect all documents for this collection
|
|
# Redis stores docs as strings at rag:{collection}:{id}, NOT as hashes
|
|
items = []
|
|
for doc_id in doc_ids:
|
|
try:
|
|
raw = await redis.get(f"rag:{name}:{doc_id}")
|
|
if raw:
|
|
doc = json.loads(raw) if isinstance(raw, str) else raw
|
|
items.append(doc)
|
|
except Exception:
|
|
pass
|
|
|
|
# Upload in chunks if payload is large (R2 REST limit ~5GB but we chunk at 5MB)
|
|
MAX_CHUNK_SIZE = 5 * 1024 * 1024 # 5MB per chunk
|
|
all_items = items
|
|
chunk_idx = 0
|
|
total_uploaded = 0
|
|
|
|
while all_items:
|
|
# Build chunk — add items until we exceed size limit
|
|
chunk_items = []
|
|
chunk_size = 0
|
|
while all_items and chunk_size < MAX_CHUNK_SIZE:
|
|
item = all_items.pop(0)
|
|
item_json = json.dumps(item, default=str)
|
|
if chunk_size + len(item_json.encode()) > MAX_CHUNK_SIZE and chunk_items:
|
|
# Put it back and upload this chunk
|
|
all_items.insert(0, item)
|
|
break
|
|
chunk_items.append(item)
|
|
chunk_size += len(item_json.encode())
|
|
|
|
payload = json.dumps(
|
|
{
|
|
"collection": name,
|
|
"version": EMBEDDING_VERSION,
|
|
"timestamp": ts,
|
|
"chunk": chunk_idx,
|
|
"count": len(chunk_items),
|
|
"items": chunk_items,
|
|
},
|
|
default=str,
|
|
).encode()
|
|
|
|
key = f"rag_snapshots/{name}/{ts}_chunk{chunk_idx}.json"
|
|
ok = await _r2_put(key, payload)
|
|
total_uploaded += len(chunk_items)
|
|
|
|
if not ok and chunk_idx == 0:
|
|
# First chunk failed — report failure
|
|
results[name] = {
|
|
"collection": name,
|
|
"status": "failed",
|
|
"count": count,
|
|
"key": key,
|
|
"size_bytes": len(payload),
|
|
"chunks_uploaded": chunk_idx,
|
|
}
|
|
break
|
|
chunk_idx += 1
|
|
|
|
# Save latest pointer (only if at least one chunk succeeded)
|
|
if chunk_idx > 0 or count == 0:
|
|
await _r2_put(
|
|
f"rag_snapshots/{name}/latest.json",
|
|
json.dumps(
|
|
{
|
|
"key": f"rag_snapshots/{name}/{ts}",
|
|
"timestamp": ts,
|
|
"count": count,
|
|
"version": EMBEDDING_VERSION,
|
|
"chunks": chunk_idx,
|
|
}
|
|
).encode(),
|
|
)
|
|
|
|
results[name] = {
|
|
"collection": name,
|
|
"status": "uploaded" if count > 0 else "empty",
|
|
"count": count,
|
|
"key": f"rag_snapshots/{name}/{ts}",
|
|
"size_bytes": sum(len(json.dumps(i, default=str).encode()) for i in items),
|
|
"chunks": chunk_idx,
|
|
}
|
|
except Exception as e:
|
|
results[name] = {"collection": name, "error": str(e)[:200]}
|
|
|
|
# Save local metadata
|
|
meta = {
|
|
"last_snapshot": ts,
|
|
"results": {k: v.get("status", "error") for k, v in results.items()},
|
|
}
|
|
(LOCAL_CACHE / "snapshot_meta.json").write_text(json.dumps(meta, indent=2))
|
|
|
|
return results
|
|
|
|
|
|
async def restore_all():
|
|
"""Restore all RAG collections from latest R2 snapshots."""
|
|
# Get collections and Redis connection
|
|
try:
|
|
from app.rag_service import COLLECTIONS, _get_redis
|
|
|
|
redis = await _get_redis()
|
|
except Exception:
|
|
return {"error": "Cannot connect to RAG service"}
|
|
|
|
results = {}
|
|
|
|
for name in COLLECTIONS:
|
|
try:
|
|
# Get latest pointer
|
|
latest_raw = await _r2_get(f"rag_snapshots/{name}/latest.json")
|
|
if not latest_raw:
|
|
results[name] = {"collection": name, "status": "no_snapshot"}
|
|
continue
|
|
|
|
latest = json.loads(latest_raw)
|
|
key = latest["key"]
|
|
|
|
# Download the snapshot
|
|
data_raw = await _r2_get(key)
|
|
if not data_raw:
|
|
results[name] = {"collection": name, "status": "download_failed"}
|
|
continue
|
|
|
|
snapshot = json.loads(data_raw)
|
|
items = snapshot.get("items", [])
|
|
snapshot_ts = latest.get("timestamp", "")
|
|
|
|
# Handle chunked snapshots — download all chunks
|
|
total_chunks = latest.get("chunks", 1)
|
|
if total_chunks > 1:
|
|
for ci in range(1, total_chunks):
|
|
chunk_key = f"rag_snapshots/{name}/{snapshot_ts}_chunk{ci}.json"
|
|
chunk_raw = await _r2_get(chunk_key)
|
|
if chunk_raw:
|
|
chunk_data = json.loads(chunk_raw)
|
|
items.extend(chunk_data.get("items", []))
|
|
|
|
# Restore to Redis: set each doc as string and add to index set
|
|
restored = 0
|
|
for item in items:
|
|
try:
|
|
doc_id = item.get("id", "")
|
|
if not doc_id:
|
|
continue
|
|
# Store as JSON string (matches rag_service format)
|
|
await redis.set(f"rag:{name}:{doc_id}", json.dumps(item, default=str))
|
|
await redis.sadd(f"rag:idx:{name}", doc_id)
|
|
restored += 1
|
|
except Exception:
|
|
pass
|
|
|
|
results[name] = {
|
|
"collection": name,
|
|
"status": "restored",
|
|
"count": restored,
|
|
"total_in_snapshot": len(items),
|
|
"snapshot_key": key,
|
|
}
|
|
except Exception as e:
|
|
results[name] = {"collection": name, "error": str(e)[:200]}
|
|
|
|
return results
|
|
|
|
|
|
async def nightly_cycle():
|
|
"""Full nightly cycle: snapshot to R2, clean local cache."""
|
|
snap = await snapshot_all()
|
|
ok = sum(1 for v in snap.values() if v.get("status") in ("uploaded", "empty"))
|
|
total = len(snap)
|
|
|
|
# Clean local cache older than retention period
|
|
cleaned = 0
|
|
cutoff = datetime.now(UTC) - timedelta(days=LOCAL_RETENTION_DAYS)
|
|
for f in LOCAL_CACHE.glob("*.json"):
|
|
try:
|
|
if datetime.fromtimestamp(f.stat().st_mtime, tz=UTC) < cutoff:
|
|
f.unlink()
|
|
cleaned += 1
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
"snapshot_results": {k: v.get("status", "error") for k, v in snap.items()},
|
|
"collections_ok": ok,
|
|
"collections_total": total,
|
|
"local_cache_cleaned": cleaned,
|
|
}
|
|
|
|
|
|
async def r2_stats():
|
|
"""Get R2 usage + local cache stats."""
|
|
# Get collections and Redis connection
|
|
try:
|
|
from app.rag_service import COLLECTIONS, _get_redis
|
|
|
|
redis = await _get_redis()
|
|
has_redis = True
|
|
except Exception:
|
|
COLLECTIONS = []
|
|
has_redis = False
|
|
|
|
# R2 bucket stats
|
|
objects = await _r2_list("rag_snapshots/")
|
|
total_size = 0
|
|
snapshot_count = 0
|
|
for obj in objects:
|
|
size = int(obj.get("size", 0))
|
|
total_size += size
|
|
snapshot_count += 1
|
|
|
|
# Local cache stats
|
|
local_files = list(LOCAL_CACHE.glob("*.json"))
|
|
local_size = sum(f.stat().st_size for f in local_files) if local_files else 0
|
|
|
|
# Collection stats (from Redis)
|
|
coll_info = {}
|
|
if has_redis:
|
|
for name in COLLECTIONS:
|
|
try:
|
|
count = await redis.scard(f"rag:idx:{name}")
|
|
coll_info[name] = count
|
|
except Exception:
|
|
coll_info[name] = 0
|
|
|
|
return {
|
|
"r2": {
|
|
"bucket": R2_BUCKET,
|
|
"endpoint": R2_BASE,
|
|
"snapshot_count": snapshot_count,
|
|
"total_size_bytes": total_size,
|
|
"total_size_mb": round(total_size / 1024 / 1024, 2),
|
|
"has_token": bool(R2_API_TOKEN),
|
|
},
|
|
"local": {
|
|
"path": str(LOCAL_CACHE),
|
|
"file_count": len(local_files),
|
|
"size_bytes": local_size,
|
|
"retention_days": LOCAL_RETENTION_DAYS,
|
|
},
|
|
"collections": coll_info,
|
|
"model_version": EMBEDDING_VERSION,
|
|
}
|