""" JarvisChat — Score-based RAG vector eviction with hysteresis. """ import asyncio import logging from datetime import datetime, timezone, timedelta import httpx from config import ( QDRANT_URL, RAG_COLLECTION, 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") eviction_lock = asyncio.Lock() EVICTION_LOG: list[dict] = [] async def _update_retrieval_count(point_id: str, current_count: int = 0): 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 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: 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 ) 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