refactor: extract eviction engine into eviction.py (rag.py 303→109 lines)
- Move all eviction logic (evict_batch, maybe_evict, EVICTION_LOG, get_collection_count/stats, get_rag_operational_stats) into eviction.py - Move QDRANT_URL, RAG_COLLECTION into config.py to break circular dep - rag.py re-exports eviction symbols for backward compatibility - Router imports updated to use eviction module directly - All 130 tests pass
This commit is contained in:
gramps
+208
@@ -0,0 +1,208 @@
|
||||
"""
|
||||
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
|
||||
)
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user