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
2026-07-06 08:00:26 -07:00
parent 8072fb3dd0
commit bb16cd6927
7 changed files with 227 additions and 210 deletions
+2
View File
@@ -50,6 +50,8 @@ BODY_LIMIT_PROFILE_BYTES = 256 * 1024
UPLOAD_DIR = "/tmp/jarvischat_uploads"
MAX_UPLOAD_BYTES = 20 * 1024 * 1024
SUPPORTED_UPLOAD_TYPES = {"text/plain", "text/markdown", "application/pdf", "application/json", "text/x-python", "text/html", "image/png", "image/jpeg", "image/gif", "image/svg+xml", "image/webp"}
QDRANT_URL = "http://192.168.50.108:6333"
RAG_COLLECTION = "jarvis_rag"
UPLOAD_CONTEXT_EXPIRY_HOURS = 1
BODY_LIMIT_UPLOAD_BYTES = MAX_UPLOAD_BYTES
+208
View File
@@ -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
+7 -201
View File
@@ -3,29 +3,25 @@ JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly.
"""
import asyncio
import logging
from datetime import datetime, timezone, timedelta
import httpx
from eviction import _update_retrieval_count
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,
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,
)
from config import MAX_SKILL_PROMPT_CHARS, QDRANT_URL, RAG_COLLECTION
log = logging.getLogger("jarvischat")
QDRANT_URL = "http://192.168.50.108:6333"
EMBED_URL = "http://192.168.50.210:11434"
EMBED_MODEL = "mxbai-embed-large"
RAG_COLLECTION = "jarvis_rag"
RAG_SCORE_THRESHOLD = 0.25
eviction_lock = asyncio.Lock()
EVICTION_LOG: list[dict] = []
# Re-export eviction symbols for backward compatibility
from eviction import ( # noqa: E402
maybe_evict, get_rag_operational_stats, EVICTION_LOG,
get_collection_count, get_collection_stats, evict_batch,
)
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list:
@@ -45,25 +41,6 @@ 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:
@@ -94,177 +71,6 @@ async def query_rag(query: str, limit: int = 3) -> list:
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()}
+2 -1
View File
@@ -7,7 +7,8 @@ from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import JSONResponse
from config import COMPLETIONS_API_KEY
from rag import chunk_text, maybe_evict, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
from eviction import maybe_evict
from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
log = logging.getLogger("jarvischat")
router = APIRouter()
+2 -4
View File
@@ -5,10 +5,8 @@ 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,
)
from eviction import get_rag_operational_stats, EVICTION_LOG
from rag import QDRANT_URL, RAG_COLLECTION
log = logging.getLogger("jarvischat")
router = APIRouter()
+2 -1
View File
@@ -10,7 +10,8 @@ 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, maybe_evict, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
from eviction import maybe_evict
from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
log = logging.getLogger("jarvischat")
router = APIRouter()
+4 -3
View File
@@ -7,6 +7,7 @@ import httpx
from fastapi.testclient import TestClient
import app
import config
import db
import rag
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
@@ -257,13 +258,13 @@ def test_maybe_evict_at_high_water(monkeypatch):
def test_maybe_evict_zero_config_disabled(monkeypatch):
"""RAG_MAX_VECTORS <= 0 should disable eviction."""
orig = rag.RAG_MAX_VECTORS
orig = config.RAG_MAX_VECTORS
try:
rag.RAG_MAX_VECTORS = 0
config.RAG_MAX_VECTORS = 0
evicted = asyncio.run(rag.maybe_evict())
assert evicted == 0
finally:
rag.RAG_MAX_VECTORS = orig
config.RAG_MAX_VECTORS = orig
# ---------- get_rag_operational_stats ----------