From bb16cd6927556692bf7fe70dca1097ef143924e2 Mon Sep 17 00:00:00 2001 From: gramps Date: Mon, 6 Jul 2026 08:00:26 -0700 Subject: [PATCH] =?UTF-8?q?refactor:=20extract=20eviction=20engine=20into?= =?UTF-8?q?=20eviction.py=20(rag.py=20303=E2=86=92109=20lines)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- config.py | 2 + eviction.py | 208 +++++++++++++++++++++++++++++++++++ rag.py | 208 ++--------------------------------- routers/ingest.py | 3 +- routers/rag_admin.py | 6 +- routers/upload.py | 3 +- tests/test_rag_management.py | 7 +- 7 files changed, 227 insertions(+), 210 deletions(-) create mode 100644 eviction.py diff --git a/config.py b/config.py index d1c277d..d5ebc94 100644 --- a/config.py +++ b/config.py @@ -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 diff --git a/eviction.py b/eviction.py new file mode 100644 index 0000000..a68b27e --- /dev/null +++ b/eviction.py @@ -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 diff --git a/rag.py b/rag.py index 3d743f4..8cbf80e 100644 --- a/rag.py +++ b/rag.py @@ -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()} diff --git a/routers/ingest.py b/routers/ingest.py index a27504c..f612f31 100644 --- a/routers/ingest.py +++ b/routers/ingest.py @@ -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() diff --git a/routers/rag_admin.py b/routers/rag_admin.py index b62b304..40a9c76 100644 --- a/routers/rag_admin.py +++ b/routers/rag_admin.py @@ -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() diff --git a/routers/upload.py b/routers/upload.py index 68f75d8..40bc552 100644 --- a/routers/upload.py +++ b/routers/upload.py @@ -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() diff --git a/tests/test_rag_management.py b/tests/test_rag_management.py index b0051ec..406325c 100644 --- a/tests/test_rag_management.py +++ b/tests/test_rag_management.py @@ -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 ----------