diff --git a/eviction.py b/eviction.py index a68b27e..7056db9 100644 --- a/eviction.py +++ b/eviction.py @@ -199,10 +199,70 @@ async def get_rag_operational_stats() -> dict: if datetime.fromisoformat(e["timestamp"]) > cutoff_30m ) + pinned_count = 0 + avg_retrieval_count = 0.0 + at_risk_count = 0 + + try: + async with httpx.AsyncClient() as client: + pinned_filter = { + "should": [ + {"match": {"key": "source", "value": src}} + for src in RAG_PINNED_SOURCES + ] + } + pinned_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll", + json={"filter": pinned_filter, "limit": 10000, "with_payload": True, "with_vector": False}, + timeout=10.0, + ) + if pinned_resp.status_code == 200: + pinned_count = len(pinned_resp.json().get("result", {}).get("points", [])) + + nonpinned_filter = { + "must_not": [ + {"match": {"key": "source", "value": src}} + for src in RAG_PINNED_SOURCES + ] + } + np_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll", + json={"filter": nonpinned_filter, "limit": 10000, "with_payload": True, "with_vector": False}, + timeout=10.0, + ) + if np_resp.status_code == 200: + points = np_resp.json().get("result", {}).get("points", []) + if points: + retrievals = [] + scored = [] + for p in points: + payload = p.get("payload", {}) + rc = payload.get("retrieval_count", 0) or 0 + retrievals.append(rc) + date_str = payload.get("ingest_date") or payload.get("upload_date", "") + if date_str: + age_hours = (now - datetime.fromisoformat(date_str)).total_seconds() / 3600 + else: + age_hours = 999999 + score = rc * RAG_ACCESS_WEIGHT + age_hours * RAG_AGE_WEIGHT + last_accessed = payload.get("last_accessed", date_str) + scored.append((score, last_accessed)) + + avg_retrieval_count = round(sum(retrievals) / len(retrievals), 2) + + scored.sort(key=lambda x: (x[0], x[1])) + at_risk_threshold = max(1, len(scored) // 10) + at_risk_count = at_risk_threshold + except Exception as e: + log.warning(f"RAG operational stats scroll error: {e}") + stats.update({ "grace_hours": RAG_GRACE_HOURS, "eviction_counts_last_1m": eviction_1m, "eviction_counts_last_5m": eviction_5m, "eviction_counts_last_30m": eviction_30m, + "pinned_count": pinned_count, + "avg_retrieval_count": avg_retrieval_count, + "at_risk_count": at_risk_count, }) return stats diff --git a/routers/rag_admin.py b/routers/rag_admin.py index 40a9c76..53736f7 100644 --- a/routers/rag_admin.py +++ b/routers/rag_admin.py @@ -2,7 +2,7 @@ import logging import httpx -from fastapi import APIRouter, Request +from fastapi import APIRouter, HTTPException, Request from fastapi.responses import JSONResponse from eviction import get_rag_operational_stats, EVICTION_LOG @@ -13,7 +13,9 @@ router = APIRouter() @router.get("/api/rag/stats") -async def rag_stats(): +async def rag_stats(request: Request): + if getattr(request.state, "session_role", "none") != "admin": + raise HTTPException(status_code=403, detail="Admin PIN required for this action") stats = await get_rag_operational_stats() stats["eviction_log_size"] = len(EVICTION_LOG) return stats diff --git a/tests/test_rag_management.py b/tests/test_rag_management.py index 406325c..9ab0dd3 100644 --- a/tests/test_rag_management.py +++ b/tests/test_rag_management.py @@ -291,14 +291,32 @@ def test_rag_operational_stats_shape(monkeypatch): def test_rag_stats_endpoint(tmp_path, monkeypatch): monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) with make_client(tmp_path) as client: - resp = client.get("/api/rag/stats", headers=_guest_headers(client)) + resp = client.get("/api/rag/stats", headers=_admin_headers(client)) assert resp.status_code == 200 data = resp.json() - assert "vector_count" in data - assert "max_vectors" in data - assert "grace_hours" in data - assert "eviction_log_size" in data assert data["vector_count"] == 123 + assert data["max_vectors"] == 50000 + assert "high_water_mark" in data + assert "low_water_mark" in data + assert data["high_water_pct"] == 80 + assert data["low_water_pct"] == 20 + assert "percent_full" in data + assert data["pinned_sources"] == ["upload", "profile"] + assert data["grace_hours"] == 1 + assert "eviction_counts_last_1m" in data + assert "eviction_counts_last_5m" in data + assert "eviction_counts_last_30m" in data + assert "pinned_count" in data + assert "avg_retrieval_count" in data + assert "at_risk_count" in data + assert "eviction_log_size" in data + + +def test_rag_stats_requires_admin(tmp_path, monkeypatch): + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + with make_client(tmp_path) as client: + resp = client.get("/api/rag/stats", headers=_guest_headers(client)) + assert resp.status_code == 403 # ---------- POST /api/rag/flush ----------