feat: Roadmap K — RAG corpus management with score-based eviction (v0.13.0)
- 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
This commit is contained in:
@@ -37,6 +37,7 @@ import routers.completions as completions
|
|||||||
import routers.upload as upload
|
import routers.upload as upload
|
||||||
import routers.ingest as ingest
|
import routers.ingest as ingest
|
||||||
import routers.hardware as hardware
|
import routers.hardware as hardware
|
||||||
|
import routers.rag_admin as rag_admin
|
||||||
|
|
||||||
# --- Logging ---
|
# --- Logging ---
|
||||||
log = logging.getLogger("jarvischat")
|
log = logging.getLogger("jarvischat")
|
||||||
@@ -146,6 +147,7 @@ for router_module in [
|
|||||||
auth_router, conversations.router, memories.router, models.router,
|
auth_router, conversations.router, memories.router, models.router,
|
||||||
presets.router, profile.router, settings.router, skills.router,
|
presets.router, profile.router, settings.router, skills.router,
|
||||||
chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router,
|
chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router,
|
||||||
|
rag_admin.router,
|
||||||
]:
|
]:
|
||||||
app.include_router(router_module)
|
app.include_router(router_module)
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ import logging
|
|||||||
|
|
||||||
log = logging.getLogger("jarvischat")
|
log = logging.getLogger("jarvischat")
|
||||||
|
|
||||||
VERSION = "v0.11.0"
|
VERSION = "v0.13.0"
|
||||||
OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434")
|
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")
|
LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081")
|
||||||
SEARXNG_BASE = "http://localhost:8888"
|
SEARXNG_BASE = "http://localhost:8888"
|
||||||
@@ -53,6 +53,16 @@ SUPPORTED_UPLOAD_TYPES = {"text/plain", "text/markdown", "application/pdf", "app
|
|||||||
UPLOAD_CONTEXT_EXPIRY_HOURS = 1
|
UPLOAD_CONTEXT_EXPIRY_HOURS = 1
|
||||||
BODY_LIMIT_UPLOAD_BYTES = MAX_UPLOAD_BYTES
|
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_CHAT_MESSAGE_CHARS = 8000
|
||||||
MAX_SEARCH_QUERY_CHARS = 500
|
MAX_SEARCH_QUERY_CHARS = 500
|
||||||
MAX_PROFILE_CHARS = 32000
|
MAX_PROFILE_CHARS = 32000
|
||||||
|
|||||||
@@ -1,13 +1,20 @@
|
|||||||
"""
|
"""
|
||||||
JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly.
|
JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly.
|
||||||
"""
|
"""
|
||||||
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
|
from datetime import datetime, timezone, timedelta
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
from db import get_db, get_setting, list_skills_with_state, format_active_skills_prompt
|
from db import get_db, get_setting, list_skills_with_state, format_active_skills_prompt
|
||||||
from memory import search_memories
|
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")
|
log = logging.getLogger("jarvischat")
|
||||||
|
|
||||||
@@ -17,10 +24,12 @@ EMBED_MODEL = "mxbai-embed-large"
|
|||||||
RAG_COLLECTION = "jarvis_rag"
|
RAG_COLLECTION = "jarvis_rag"
|
||||||
RAG_SCORE_THRESHOLD = 0.25
|
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:
|
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list:
|
||||||
words = text.split()
|
words = text.split()
|
||||||
# Approximate token count as len(words) * 1.3
|
|
||||||
target_words = int(chunk_size / 1.3)
|
target_words = int(chunk_size / 1.3)
|
||||||
overlap_words = int(overlap / 1.3)
|
overlap_words = int(overlap / 1.3)
|
||||||
if not words:
|
if not words:
|
||||||
@@ -36,6 +45,25 @@ def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list:
|
|||||||
return chunks
|
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:
|
async def query_rag(query: str, limit: int = 3) -> list:
|
||||||
try:
|
try:
|
||||||
async with httpx.AsyncClient() as client:
|
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:
|
if search_resp.status_code != 200:
|
||||||
return []
|
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:
|
except Exception as e:
|
||||||
log.warning(f"RAG query error: {e}")
|
log.warning(f"RAG query error: {e}")
|
||||||
return []
|
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:
|
async def build_system_prompt(db, extra_prompt: str = "", user_message: str = "") -> str:
|
||||||
parts = []
|
parts = []
|
||||||
settings = {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM settings").fetchall()}
|
settings = {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM settings").fetchall()}
|
||||||
|
|||||||
+6
-1
@@ -7,7 +7,7 @@ from fastapi import APIRouter, HTTPException, Request
|
|||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
from config import COMPLETIONS_API_KEY
|
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")
|
log = logging.getLogger("jarvischat")
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
@@ -61,4 +61,9 @@ async def ingest_content(request: Request):
|
|||||||
else:
|
else:
|
||||||
log.warning(f"Ingest Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}")
|
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}"}
|
return {"chunks_ingested": ingested, "source": source, "message": f"Ingested {ingested} chunks from {source}"}
|
||||||
|
|||||||
@@ -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}"})
|
||||||
+5
-1
@@ -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 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 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")
|
log = logging.getLogger("jarvischat")
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
@@ -87,6 +87,10 @@ async def upload_file(
|
|||||||
else:
|
else:
|
||||||
log.warning(f"Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}")
|
log.warning(f"Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}")
|
||||||
result["chunks_ingested"] = ingested
|
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"):
|
if mode in ("context", "both"):
|
||||||
expires = (datetime.now(timezone.utc) + timedelta(hours=UPLOAD_CONTEXT_EXPIRY_HOURS)).isoformat()
|
expires = (datetime.now(timezone.utc) + timedelta(hours=UPLOAD_CONTEXT_EXPIRY_HOURS)).isoformat()
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user