From 8072fb3dd08292391c7e88760614932f88d5eafc Mon Sep 17 00:00:00 2001 From: gramps Date: Mon, 6 Jul 2026 07:56:09 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20Roadmap=20K=20=E2=80=94=20RAG=20corpus?= =?UTF-8?q?=20management=20with=20score-based=20eviction=20(v0.13.0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Config: RAG_MAX_VECTORS, high/low water marks, grace period, weights - rag.py: get_collection_count, evict_batch, maybe_evict (asyncio.Lock), get_rag_operational_stats, EVICTION_LOG, retrieval_count tracking - routers/rag_admin.py: GET /api/rag/stats, POST /api/rag/flush (admin) - Wire maybe_evict() into upload.py and ingest.py after Qdrant upsert - 16 tests: collection stats, eviction scoring, pinned/grace/batch guards, endpoint auth, race lock, flush, operational stats shape - Bump to v0.13.0 --- app.py | 2 + config.py | 12 +- rag.py | 211 +++++++++++++++++++- routers/ingest.py | 7 +- routers/rag_admin.py | 60 ++++++ routers/upload.py | 6 +- tests/test_rag_management.py | 373 +++++++++++++++++++++++++++++++++++ 7 files changed, 665 insertions(+), 6 deletions(-) create mode 100644 routers/rag_admin.py create mode 100644 tests/test_rag_management.py diff --git a/app.py b/app.py index 234464c..d4214c1 100644 --- a/app.py +++ b/app.py @@ -37,6 +37,7 @@ import routers.completions as completions import routers.upload as upload import routers.ingest as ingest import routers.hardware as hardware +import routers.rag_admin as rag_admin # --- Logging --- log = logging.getLogger("jarvischat") @@ -146,6 +147,7 @@ for router_module in [ auth_router, conversations.router, memories.router, models.router, presets.router, profile.router, settings.router, skills.router, chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router, + rag_admin.router, ]: app.include_router(router_module) diff --git a/config.py b/config.py index 50a6b17..d1c277d 100644 --- a/config.py +++ b/config.py @@ -9,7 +9,7 @@ import logging log = logging.getLogger("jarvischat") -VERSION = "v0.11.0" +VERSION = "v0.13.0" OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434") LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081") SEARXNG_BASE = "http://localhost:8888" @@ -53,6 +53,16 @@ SUPPORTED_UPLOAD_TYPES = {"text/plain", "text/markdown", "application/pdf", "app UPLOAD_CONTEXT_EXPIRY_HOURS = 1 BODY_LIMIT_UPLOAD_BYTES = MAX_UPLOAD_BYTES +# --- RAG eviction --- +RAG_MAX_VECTORS = 50000 +RAG_EVICTION_HIGH_WATER = 0.80 +RAG_EVICTION_LOW_WATER = 0.20 +RAG_EVICTION_BATCH = 1000 +RAG_PINNED_SOURCES = ["upload", "profile"] +RAG_GRACE_HOURS = 1 +RAG_ACCESS_WEIGHT = 1.0 +RAG_AGE_WEIGHT = 0.1 + MAX_CHAT_MESSAGE_CHARS = 8000 MAX_SEARCH_QUERY_CHARS = 500 MAX_PROFILE_CHARS = 32000 diff --git a/rag.py b/rag.py index 8c023bc..3d743f4 100644 --- a/rag.py +++ b/rag.py @@ -1,13 +1,20 @@ """ JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly. """ +import asyncio import logging +from datetime import datetime, timezone, timedelta import httpx from db import get_db, get_setting, list_skills_with_state, format_active_skills_prompt from memory import search_memories -from config import MAX_SKILL_PROMPT_CHARS +from config import ( + MAX_SKILL_PROMPT_CHARS, + RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, + RAG_EVICTION_BATCH, RAG_PINNED_SOURCES, RAG_GRACE_HOURS, + RAG_ACCESS_WEIGHT, RAG_AGE_WEIGHT, +) log = logging.getLogger("jarvischat") @@ -17,10 +24,12 @@ EMBED_MODEL = "mxbai-embed-large" RAG_COLLECTION = "jarvis_rag" RAG_SCORE_THRESHOLD = 0.25 +eviction_lock = asyncio.Lock() +EVICTION_LOG: list[dict] = [] + def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list: words = text.split() - # Approximate token count as len(words) * 1.3 target_words = int(chunk_size / 1.3) overlap_words = int(overlap / 1.3) if not words: @@ -36,6 +45,25 @@ def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list: return chunks +async def _update_retrieval_count(point_id: str, current_count: int = 0): + """Fire-and-forget increment of retrieval_count.""" + try: + async with httpx.AsyncClient() as client: + payload = { + "retrieval_count": current_count + 1, + "last_accessed": datetime.now(timezone.utc).isoformat(), + } + resp = await client.put( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/payload", + json={"points": [point_id], "payload": payload}, + timeout=5.0, + ) + if resp.status_code not in (200, 201): + log.warning(f"Failed to increment retrieval count for {point_id}: {resp.status_code}") + except Exception as e: + log.warning(f"Error incrementing retrieval count for {point_id}: {e}") + + async def query_rag(query: str, limit: int = 3) -> list: try: async with httpx.AsyncClient() as client: @@ -54,12 +82,189 @@ async def query_rag(query: str, limit: int = 3) -> list: ) if search_resp.status_code != 200: return [] - return search_resp.json().get("result", []) + results = search_resp.json().get("result", []) + for r in results: + pid = r.get("id") + if pid: + current = r.get("payload", {}).get("retrieval_count", 0) or 0 + asyncio.ensure_future(_update_retrieval_count(pid, current)) + return results except Exception as e: log.warning(f"RAG query error: {e}") return [] +async def get_collection_count() -> int: + try: + async with httpx.AsyncClient() as client: + resp = await client.get( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}", + timeout=10.0, + ) + if resp.status_code == 200: + return resp.json().get("result", {}).get("vectors_count", 0) + except Exception as e: + log.warning(f"get_collection_count error: {e}") + return 0 + + +async def get_collection_stats() -> dict: + count = await get_collection_count() + high_water_pct = int(RAG_EVICTION_HIGH_WATER * 100) + low_water_pct = int(RAG_EVICTION_LOW_WATER * 100) + percent_full = round((count / RAG_MAX_VECTORS) * 100, 1) if RAG_MAX_VECTORS > 0 else 0 + return { + "vector_count": count, + "max_vectors": RAG_MAX_VECTORS, + "high_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER), + "low_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER), + "high_water_pct": high_water_pct, + "low_water_pct": low_water_pct, + "percent_full": percent_full, + "pinned_sources": list(RAG_PINNED_SOURCES), + } + + +async def evict_batch(batch_size: int) -> int: + """Scroll non-pinned, out-of-grace-period vectors, compute scores, delete lowest-scoring.""" + filter_conditions = { + "must_not": [ + {"match": {"key": "source", "value": src}} + for src in RAG_PINNED_SOURCES + ] + } + try: + async with httpx.AsyncClient() as client: + scroll_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll", + json={ + "filter": filter_conditions, + "limit": min(batch_size * 10, 10000), + "with_payload": True, + "with_vector": False, + }, + timeout=30.0, + ) + if scroll_resp.status_code != 200: + log.warning(f"Eviction scroll failed: {scroll_resp.status_code}") + return 0 + + points = scroll_resp.json().get("result", {}).get("points", []) + if not points: + return 0 + + now = datetime.now(timezone.utc) + scored = [] + for p in points: + payload = p.get("payload", {}) + 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 + + if age_hours < RAG_GRACE_HOURS: + continue + + retrieval_count = payload.get("retrieval_count", 0) or 0 + score = retrieval_count * RAG_ACCESS_WEIGHT + age_hours * RAG_AGE_WEIGHT + last_accessed = payload.get("last_accessed", date_str) + scored.append((score, last_accessed, p["id"])) + + if not scored: + log.warning("No evictable vectors found (all pinned or newborn)") + return 0 + + scored.sort(key=lambda x: (x[0], x[1])) + to_delete = [p[2] for p in scored[:batch_size]] + if not to_delete: + return 0 + + delete_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/delete", + json={"points": to_delete}, + timeout=30.0, + ) + if delete_resp.status_code not in (200, 201): + log.warning(f"Eviction delete failed: {delete_resp.status_code}") + return 0 + + return len(to_delete) + except Exception as e: + log.warning(f"evict_batch error: {e}") + return 0 + + +async def maybe_evict() -> int: + if RAG_MAX_VECTORS <= 0: + return 0 + effective_batch = max(RAG_EVICTION_BATCH, 1) + + async with eviction_lock: + count = await get_collection_count() + threshold_high = int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER) + threshold_low = int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER) + + if count < threshold_high: + return 0 + + total_evicted = 0 + while count >= threshold_low: + if total_evicted > 0 and count < threshold_low: + break + deleted = await evict_batch(effective_batch) + if deleted == 0: + break + total_evicted += deleted + count -= deleted + if count < threshold_high and total_evicted > 0: + break + if count < threshold_low: + break + + if total_evicted > 0: + entry = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "count": total_evicted, + "remaining": count, + } + EVICTION_LOG.append(entry) + if len(EVICTION_LOG) > 1000: + EVICTION_LOG.pop(0) + log.info(f"Evicted {total_evicted} vectors ({count} remaining)") + + return total_evicted + + +async def get_rag_operational_stats() -> dict: + stats = await get_collection_stats() + now = datetime.now(timezone.utc) + cutoff_1m = now - timedelta(minutes=1) + cutoff_5m = now - timedelta(minutes=5) + cutoff_30m = now - timedelta(minutes=30) + + eviction_1m = sum( + e["count"] for e in EVICTION_LOG + if datetime.fromisoformat(e["timestamp"]) > cutoff_1m + ) + eviction_5m = sum( + e["count"] for e in EVICTION_LOG + if datetime.fromisoformat(e["timestamp"]) > cutoff_5m + ) + eviction_30m = sum( + e["count"] for e in EVICTION_LOG + if datetime.fromisoformat(e["timestamp"]) > cutoff_30m + ) + + 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, + }) + return stats + + async def build_system_prompt(db, extra_prompt: str = "", user_message: str = "") -> str: parts = [] settings = {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM settings").fetchall()} diff --git a/routers/ingest.py b/routers/ingest.py index 82aa5d5..a27504c 100644 --- a/routers/ingest.py +++ b/routers/ingest.py @@ -7,7 +7,7 @@ from fastapi import APIRouter, HTTPException, Request from fastapi.responses import JSONResponse from config import COMPLETIONS_API_KEY -from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION +from rag import chunk_text, maybe_evict, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION log = logging.getLogger("jarvischat") router = APIRouter() @@ -61,4 +61,9 @@ async def ingest_content(request: Request): else: log.warning(f"Ingest Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}") + if ingested > 0: + evicted = await maybe_evict() + if evicted: + log.info(f"Evicted {evicted} vectors after ingest") + return {"chunks_ingested": ingested, "source": source, "message": f"Ingested {ingested} chunks from {source}"} diff --git a/routers/rag_admin.py b/routers/rag_admin.py new file mode 100644 index 0000000..b62b304 --- /dev/null +++ b/routers/rag_admin.py @@ -0,0 +1,60 @@ +"""JarvisChat routers — RAG corpus management admin endpoints.""" +import logging + +import httpx +from fastapi import APIRouter, Request +from fastapi.responses import JSONResponse + +from rag import ( + get_rag_operational_stats, EVICTION_LOG, + QDRANT_URL, RAG_COLLECTION, +) + +log = logging.getLogger("jarvischat") +router = APIRouter() + + +@router.get("/api/rag/stats") +async def rag_stats(): + stats = await get_rag_operational_stats() + stats["eviction_log_size"] = len(EVICTION_LOG) + return stats + + +@router.post("/api/rag/flush") +async def rag_flush(request: Request): + try: + async with httpx.AsyncClient() as client: + scroll_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll", + json={"limit": 10000, "with_payload": False, "with_vector": False}, + timeout=30.0, + ) + if scroll_resp.status_code != 200: + return JSONResponse(status_code=502, content={"detail": f"Qdrant scroll failed: {scroll_resp.status_code}"}) + + all_points = scroll_resp.json().get("result", {}).get("points", []) + point_ids = [p["id"] for p in all_points] + + if not point_ids: + return {"deleted_count": 0, "collection": RAG_COLLECTION, "status": "flushed"} + + delete_resp = await client.post( + f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/delete", + json={"points": point_ids}, + timeout=30.0, + ) + if delete_resp.status_code not in (200, 201): + return JSONResponse(status_code=502, content={"detail": f"Qdrant delete failed: {delete_resp.status_code}"}) + + EVICTION_LOG.clear() + log.warning(f"RAG collection '{RAG_COLLECTION}' flushed ({len(point_ids)} points deleted)") + + return { + "deleted_count": len(point_ids), + "collection": RAG_COLLECTION, + "status": "flushed", + } + except Exception as e: + log.warning(f"RAG flush error: {e}") + return JSONResponse(status_code=502, content={"detail": f"RAG flush error: {e}"}) diff --git a/routers/upload.py b/routers/upload.py index 818a546..68f75d8 100644 --- a/routers/upload.py +++ b/routers/upload.py @@ -10,7 +10,7 @@ from fastapi.responses import JSONResponse from config import UPLOAD_DIR, MAX_UPLOAD_BYTES, SUPPORTED_UPLOAD_TYPES, UPLOAD_CONTEXT_EXPIRY_HOURS from db import get_db, insert_upload_context, list_upload_context_by_conversation, delete_upload_context_by_id -from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION +from rag import chunk_text, maybe_evict, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION log = logging.getLogger("jarvischat") router = APIRouter() @@ -87,6 +87,10 @@ async def upload_file( else: log.warning(f"Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}") result["chunks_ingested"] = ingested + if ingested > 0: + evicted = await maybe_evict() + if evicted: + log.info(f"Evicted {evicted} vectors after upload") if mode in ("context", "both"): expires = (datetime.now(timezone.utc) + timedelta(hours=UPLOAD_CONTEXT_EXPIRY_HOURS)).isoformat() diff --git a/tests/test_rag_management.py b/tests/test_rag_management.py new file mode 100644 index 0000000..b0051ec --- /dev/null +++ b/tests/test_rag_management.py @@ -0,0 +1,373 @@ +import asyncio +import os +from datetime import datetime, timezone, timedelta +from pathlib import Path + +import httpx +from fastapi.testclient import TestClient + +import app +import db +import rag +from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS + + +def make_client(tmp_path: Path) -> TestClient: + os.environ["JARVISCHAT_ADMIN_PIN"] = "1234" + db.DB_PATH = tmp_path / "jarvischat-rag-mgmt.db" + SESSIONS.clear() + PIN_ATTEMPTS.clear() + RATE_EVENTS.clear() + db.init_db() + return TestClient(app.app, raise_server_exceptions=False) + + +def _admin_headers(client: TestClient) -> dict: + login = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"}) + sid = login.json()["session_id"] + return {"X-Session-ID": sid, "Origin": "http://testserver"} + + +def _guest_headers(client: TestClient) -> dict: + sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"] + return {"X-Session-ID": sid, "Origin": "http://testserver"} + + +class FakeResponse: + def __init__(self, status, json_data=None): + self.status_code = status + self._json = json_data or {} + + def json(self): + return self._json + + +class FakeAsyncClient: + def __init__(self, *a, **kw): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + pass + + async def get(self, url, **kw): + if "/collections/jarvis_rag" in url: + return FakeResponse(200, {"result": {"vectors_count": 123}}) + return FakeResponse(200) + + async def post(self, url, **kw): + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": []}}) + if "/points/delete" in url: + return FakeResponse(200) + return FakeResponse(200) + + async def put(self, url, **kw): + return FakeResponse(200) + + +def _old_ts(hours_ago: float = 24) -> str: + return (datetime.now(timezone.utc) - timedelta(hours=hours_ago)).isoformat() + + +def _young_ts(hours_ago: float = 0.1) -> str: + return (datetime.now(timezone.utc) - timedelta(hours=hours_ago)).isoformat() + + +# ---------- get_collection_count ---------- + +def test_get_collection_count(monkeypatch): + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + count = asyncio.run(rag.get_collection_count()) + assert count == 123 + + +# ---------- get_collection_stats ---------- + +def test_get_collection_stats_shape(monkeypatch): + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + stats = asyncio.run(rag.get_collection_stats()) + assert stats["vector_count"] == 123 + assert stats["max_vectors"] == 50000 + assert stats["high_water_mark"] == 40000 + assert stats["low_water_mark"] == 10000 + assert stats["high_water_pct"] == 80 + assert stats["low_water_pct"] == 20 + assert 0 < stats["percent_full"] < 1 + assert "upload" in stats["pinned_sources"] + assert "profile" in stats["pinned_sources"] + + +# ---------- evict_batch ---------- + +def test_evict_batch_excludes_pinned_sources(monkeypatch): + """Pinned sources ('upload', 'profile') should be in the must_not scroll filter.""" + old = _old_ts(48) + # Only non-pinned points are returned (real Qdrant would honour must_not filter) + scroll_points = [ + {"id": "old-data", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}}, + ] + + class ScrollClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/scroll" in url: + must_not = kw.get("json", {}).get("filter", {}).get("must_not", []) + pinned_values = [m["match"]["value"] for m in must_not] + assert "upload" in pinned_values + assert "profile" in pinned_values + return FakeResponse(200, {"result": {"points": scroll_points}}) + if "/points/delete" in url: + return FakeResponse(200) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: ScrollClient()) + deleted = asyncio.run(rag.evict_batch(10)) + assert deleted == 1 + + +def test_evict_batch_respects_grace_period(monkeypatch): + """Vectors younger than RAG_GRACE_HOURS should be skipped.""" + old = _old_ts(48) + young = _young_ts(0.1) + scroll_points = [ + {"id": "mature", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}}, + {"id": "newborn", "payload": {"source": "terminal", "ingest_date": young, "retrieval_count": 0}}, + ] + + class GraceClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": scroll_points}}) + if "/points/delete" in url: + deleted = kw.get("json", {}).get("points", []) + assert "newborn" not in deleted + return FakeResponse(200) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: GraceClient()) + deleted = asyncio.run(rag.evict_batch(10)) + assert deleted == 1 + + +def test_evict_batch_respects_batch_size(monkeypatch): + """Only up to batch_size vectors should be deleted per call.""" + old = _old_ts(48) + points = [{"id": f"p{i}", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}} for i in range(50)] + + class BatchClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/delete" in url: + deleted = kw.get("json", {}).get("points", []) + assert len(deleted) == 10 + return FakeResponse(200) + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": points}}) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: BatchClient()) + deleted = asyncio.run(rag.evict_batch(10)) + assert deleted == 10 + + +def test_evict_batch_all_pinned_returns_zero(monkeypatch): + """If scroll returns nothing (all points filtered by must_not), evict_batch returns 0.""" + class EmptyClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": []}}) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: EmptyClient()) + deleted = asyncio.run(rag.evict_batch(10)) + assert deleted == 0 + + +def test_evict_batch_scores_lowest_first(monkeypatch): + """Vectors with lower scores should be evicted first.""" + old = _old_ts(48) + points = [ + {"id": "high-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 100}}, + {"id": "low-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}}, + {"id": "mid-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 50}}, + ] + + class ScoreClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/delete" in url: + deleted = kw.get("json", {}).get("points", []) + assert "low-score" in deleted + assert "high-score" not in deleted + assert "mid-score" not in deleted + return FakeResponse(200) + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": points}}) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: ScoreClient()) + deleted = asyncio.run(rag.evict_batch(1)) + assert deleted == 1 + + +# ---------- maybe_evict ---------- + +def test_maybe_evict_below_high_water(monkeypatch): + """When count is below high-water mark, eviction should not fire.""" + class LowCountClient(FakeAsyncClient): + async def get(self, url, **kw): + # 30000 < 40000 high water + return FakeResponse(200, {"result": {"vectors_count": 30000}}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: LowCountClient()) + rag.EVICTION_LOG.clear() + evicted = asyncio.run(rag.maybe_evict()) + assert evicted == 0 + assert len(rag.EVICTION_LOG) == 0 + + +def test_maybe_evict_at_high_water(monkeypatch): + """When count reaches high-water mark, eviction should fire.""" + class HighCountClient(FakeAsyncClient): + def __init__(self, *a, **kw): + super().__init__() + self.call_count = 0 + + async def get(self, url, **kw): + return FakeResponse(200, {"result": {"vectors_count": 45000}}) + + async def post(self, url, **kw): + if "/points/scroll" in url: + old = _old_ts(48) + points = [{"id": f"evict-me-{i}", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}} for i in range(100)] + return FakeResponse(200, {"result": {"points": points}}) + if "/points/delete" in url: + return FakeResponse(200) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: HighCountClient()) + rag.EVICTION_LOG.clear() + evicted = asyncio.run(rag.maybe_evict()) + assert evicted > 0 + assert len(rag.EVICTION_LOG) == 1 + entry = rag.EVICTION_LOG[0] + assert "timestamp" in entry + assert entry["count"] > 0 + + +def test_maybe_evict_zero_config_disabled(monkeypatch): + """RAG_MAX_VECTORS <= 0 should disable eviction.""" + orig = rag.RAG_MAX_VECTORS + try: + rag.RAG_MAX_VECTORS = 0 + evicted = asyncio.run(rag.maybe_evict()) + assert evicted == 0 + finally: + rag.RAG_MAX_VECTORS = orig + + +# ---------- get_rag_operational_stats ---------- + +def test_rag_operational_stats_shape(monkeypatch): + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + rag.EVICTION_LOG.clear() + rag.EVICTION_LOG.append({ + "timestamp": datetime.now(timezone.utc).isoformat(), + "count": 500, + "remaining": 40000, + }) + stats = asyncio.run(rag.get_rag_operational_stats()) + assert stats["vector_count"] == 123 + assert stats["grace_hours"] == 1 + assert "eviction_counts_last_1m" in stats + assert "eviction_counts_last_5m" in stats + assert "eviction_counts_last_30m" in stats + assert stats["eviction_counts_last_1m"] == 500 + + +# ---------- GET /api/rag/stats ---------- + +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)) + 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 + + +# ---------- POST /api/rag/flush ---------- + +def test_rag_flush_endpoint(tmp_path, monkeypatch): + class FlushClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": [{"id": "a"}, {"id": "b"}]}}) + if "/points/delete" in url: + return FakeResponse(200) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FlushClient()) + with make_client(tmp_path) as client: + resp = client.post("/api/rag/flush", headers=_admin_headers(client)) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "flushed" + assert data["deleted_count"] == 2 + assert data["collection"] == rag.RAG_COLLECTION + + +def test_rag_flush_requires_admin(tmp_path, monkeypatch): + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient()) + with make_client(tmp_path) as client: + resp = client.post("/api/rag/flush", headers=_guest_headers(client)) + assert resp.status_code == 403 + + +def test_rag_flush_empty_collection(tmp_path, monkeypatch): + class EmptyClient(FakeAsyncClient): + async def post(self, url, **kw): + if "/points/scroll" in url: + return FakeResponse(200, {"result": {"points": []}}) + return FakeResponse(200) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: EmptyClient()) + with make_client(tmp_path) as client: + resp = client.post("/api/rag/flush", headers=_admin_headers(client)) + assert resp.status_code == 200 + data = resp.json() + assert data["deleted_count"] == 0 + assert data["status"] == "flushed" + + +# ---------- Race lock ---------- + +def test_eviction_lock_prevents_concurrent_eviction(monkeypatch): + """Concurrent calls to maybe_evict should queue; only one evicts.""" + call_order = [] + + async def slow_get_collection_count(): + call_order.append("count") + return 45000 + + async def slow_evict_batch(bs): + call_order.append("evict") + await asyncio.sleep(0.05) + return 500 + + monkeypatch.setattr(rag, "get_collection_count", slow_get_collection_count) + monkeypatch.setattr(rag, "evict_batch", slow_evict_batch) + rag.EVICTION_LOG.clear() + + async def run_concurrent(): + r1, r2 = await asyncio.gather(rag.maybe_evict(), rag.maybe_evict()) + return r1, r2 + + r1, r2 = asyncio.run(run_concurrent()) + # First call evicted, second found count already below high water or lock serialized + assert r1 >= 0 + assert r2 >= 0