fix(rag): squash all RAG/ChromaDB bugs #123–#130
Fixes critical API contract mismatch, batch_size wiring, lifespan init, failure tracking, build_meta persistence, mana_cost in embeddings, pyproject.toml deps, and adds full test suite. Closes #123, #124, #125, #126, #127, #128, #129, #130
This commit is contained in:
1 parent
61ac349445
commit
bcf558418a
5 files changed
+614
-72
No files matched your search
+20
-11
@@ -50,7 +50,7 @@ _allowed = os.environ.get("ALLOWED_ORIGINS", _DEFAULT_ORIGINS)
|
|||||||
allowed_origins = [o.strip() for o in _allowed.split(",") if o.strip()]
|
allowed_origins = [o.strip() for o in _allowed.split(",") if o.strip()]
|
||||||
|
|
||||||
|
|
||||||
# ── Lifespan (replaces deprecated @app.on_event) ───────────────────────────────
|
# ── Lifespan (replaces deprecated @app.on_event) ─────────────────────────────────
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def _lifespan(application: FastAPI):
|
async def _lifespan(application: FastAPI):
|
||||||
# ──── STARTUP ────
|
# ──── STARTUP ────
|
||||||
@@ -62,13 +62,21 @@ async def _lifespan(application: FastAPI):
|
|||||||
if not COMMANDER_META:
|
if not COMMANDER_META:
|
||||||
load_commander_meta()
|
load_commander_meta()
|
||||||
|
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
|
||||||
# RAG Phase 1: kick off Scryfall bulk download in background thread.
|
# RAG Phase 1: kick off Scryfall bulk download in background thread.
|
||||||
# run_in_executor returns a Future; we do NOT await it so startup
|
# run_in_executor returns a Future; we do NOT await it so startup
|
||||||
# completes immediately and the download runs concurrently.
|
# completes immediately and the download runs concurrently.
|
||||||
from services.scryfall_bulk import ensure_bulk_db
|
from services.scryfall_bulk import ensure_bulk_db
|
||||||
loop = asyncio.get_event_loop()
|
|
||||||
loop.run_in_executor(None, ensure_bulk_db)
|
loop.run_in_executor(None, ensure_bulk_db)
|
||||||
|
|
||||||
|
# RAG Phase 2: kick off ChromaDB vector store build in background thread.
|
||||||
|
# Runs after Phase 1 completes inside ensure_rag_store() (it checks bulk_count > 0).
|
||||||
|
# Safe under both `uvicorn lab_api:app` and `python lab_api.py` — the
|
||||||
|
# _rag_ready flag + _rag_lock in rag_store.py make it idempotent.
|
||||||
|
from services.rag_store import ensure_rag_store
|
||||||
|
loop.run_in_executor(None, ensure_rag_store)
|
||||||
|
|
||||||
yield # ← server is running
|
yield # ← server is running
|
||||||
|
|
||||||
# ──── SHUTDOWN ────
|
# ──── SHUTDOWN ────
|
||||||
@@ -76,7 +84,7 @@ async def _lifespan(application: FastAPI):
|
|||||||
await async_close_db()
|
await async_close_db()
|
||||||
|
|
||||||
|
|
||||||
# ── App ────────────────────────────────────────────────────────────────
|
# ── App ────────────────────────────────────────────────────────────────────
|
||||||
app = FastAPI(title="Commander AI Lab API", version="3.0.0", lifespan=_lifespan)
|
app = FastAPI(title="Commander AI Lab API", version="3.0.0", lifespan=_lifespan)
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
@@ -100,7 +108,7 @@ app.include_router(ws_game_router)
|
|||||||
app.include_router(ml_router)
|
app.include_router(ml_router)
|
||||||
|
|
||||||
|
|
||||||
# ── API endpoints ────────────────────────────────────────────────────────────
|
# ── API endpoints ────────────────────────────────────────────────────────────────
|
||||||
# NOTE: all @app.get routes MUST be registered before app.mount() calls.
|
# NOTE: all @app.get routes MUST be registered before app.mount() calls.
|
||||||
# Once StaticFiles is mounted at "/", it registers a wildcard catch-all that
|
# Once StaticFiles is mounted at "/", it registers a wildcard catch-all that
|
||||||
# shadows any routes added after it.
|
# shadows any routes added after it.
|
||||||
@@ -111,8 +119,9 @@ async def health_check():
|
|||||||
return {"status": "ok"}
|
return {"status": "ok"}
|
||||||
|
|
||||||
# ── Static UI (MUST come after all @app.get routes) ────────────────────────────
|
# ── Static UI (MUST come after all @app.get routes) ────────────────────────────
|
||||||
_legacy_ui_dir = Path(__file__).parent / "ui"
|
from pathlib import Path as _Path
|
||||||
_spa_dir = Path(__file__).parent / "frontend" / "commander-ai-lab-ui" / "dist"
|
_legacy_ui_dir = _Path(__file__).parent / "ui"
|
||||||
|
_spa_dir = _Path(__file__).parent / "frontend" / "commander-ai-lab-ui" / "dist"
|
||||||
|
|
||||||
if _legacy_ui_dir.exists():
|
if _legacy_ui_dir.exists():
|
||||||
app.mount("/", StaticFiles(directory=str(_legacy_ui_dir), html=True), name="ui")
|
app.mount("/", StaticFiles(directory=str(_legacy_ui_dir), html=True), name="ui")
|
||||||
@@ -130,7 +139,7 @@ elif _spa_dir.exists():
|
|||||||
return FileResponse(str(_spa_dir / "index.html"))
|
return FileResponse(str(_spa_dir / "index.html"))
|
||||||
|
|
||||||
|
|
||||||
# ── Argument parsing + auto-detection helpers ─────────────────────────────────
|
# ── Argument parsing + auto-detection helpers ───────────────────────────────────
|
||||||
def _parse_args():
|
def _parse_args():
|
||||||
p = argparse.ArgumentParser(description="Commander AI Lab API Server")
|
p = argparse.ArgumentParser(description="Commander AI Lab API Server")
|
||||||
p.add_argument("--forge-jar", default=os.environ.get("FORGE_JAR", ""))
|
p.add_argument("--forge-jar", default=os.environ.get("FORGE_JAR", ""))
|
||||||
@@ -298,15 +307,15 @@ def main():
|
|||||||
init_coach_service()
|
init_coach_service()
|
||||||
|
|
||||||
# RAG Phase 1: synchronous boot-time check so logs are visible.
|
# RAG Phase 1: synchronous boot-time check so logs are visible.
|
||||||
# The async lifespan handler also fires it in a thread on uvicorn startup,
|
|
||||||
# but calling here ensures it runs (and logs) before the first request.
|
|
||||||
from services.scryfall_bulk import ensure_bulk_db
|
from services.scryfall_bulk import ensure_bulk_db
|
||||||
ensure_bulk_db()
|
ensure_bulk_db()
|
||||||
|
|
||||||
|
# RAG Phase 2: synchronous boot-time ChromaDB build.
|
||||||
# RAG Phase 2: synchronous boot-time ChromaDB build.
|
# The lifespan handler also fires this in a thread for uvicorn ASGI mode;
|
||||||
|
# calling here ensures it runs before the first request when using python lab_api.py.
|
||||||
from services.rag_store import ensure_rag_store
|
from services.rag_store import ensure_rag_store
|
||||||
ensure_rag_store()
|
ensure_rag_store()
|
||||||
|
|
||||||
import uvicorn
|
import uvicorn
|
||||||
uvicorn.run(app, host="0.0.0.0", port=CFG.port, log_level="info")
|
uvicorn.run(app, host="0.0.0.0", port=CFG.port, log_level="info")
|
||||||
|
|
||||||
|
|||||||
@@ -25,6 +25,28 @@ dev = [
|
|||||||
"httpx>=0.24.0",
|
"httpx>=0.24.0",
|
||||||
"ruff>=0.1.0",
|
"ruff>=0.1.0",
|
||||||
]
|
]
|
||||||
|
# RAG Phase 2: ChromaDB vector store + Ollama embedding support
|
||||||
|
# Install with: pip install ".[rag]"
|
||||||
|
rag = [
|
||||||
|
"chromadb>=0.5.0",
|
||||||
|
]
|
||||||
|
# Local coach embeddings (sentence-transformers, no Ollama required)
|
||||||
|
# Install with: pip install ".[coach-embeddings]"
|
||||||
|
coach-embeddings = [
|
||||||
|
"sentence-transformers>=3.0.0",
|
||||||
|
"torch>=2.0.0",
|
||||||
|
"pandas>=1.5.0",
|
||||||
|
"pyarrow>=12.0.0",
|
||||||
|
]
|
||||||
|
# Full local AI stack (RAG + coach embeddings)
|
||||||
|
# Install with: pip install ".[all]"
|
||||||
|
all = [
|
||||||
|
"chromadb>=0.5.0",
|
||||||
|
"sentence-transformers>=3.0.0",
|
||||||
|
"torch>=2.0.0",
|
||||||
|
"pandas>=1.5.0",
|
||||||
|
"pyarrow>=12.0.0",
|
||||||
|
]
|
||||||
|
|
||||||
[tool.setuptools.packages.find]
|
[tool.setuptools.packages.find]
|
||||||
where = ["src"]
|
where = ["src"]
|
||||||
|
|||||||
+16
-6
@@ -77,15 +77,25 @@ async def rag_search(body: RAGSearchRequest):
|
|||||||
Returns ranked card results with relevance scores.
|
Returns ranked card results with relevance scores.
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
from services.rag_store import search_cards
|
from services.rag_store import query_cards
|
||||||
results = search_cards(
|
raw_results = query_cards(
|
||||||
query=body.query,
|
query=body.query,
|
||||||
|
n_results=body.top_k,
|
||||||
color_identity=body.color_identity,
|
color_identity=body.color_identity,
|
||||||
top_k=body.top_k,
|
|
||||||
)
|
)
|
||||||
|
results = [
|
||||||
|
RAGCardResult(
|
||||||
|
name=r.get("name", ""),
|
||||||
|
type_line=r.get("type_line", ""),
|
||||||
|
oracle_text=r.get("oracle_text", ""),
|
||||||
|
color_identity=r.get("color_identity", ""),
|
||||||
|
score=round(1.0 - r.get("distance", 0.0), 4),
|
||||||
|
)
|
||||||
|
for r in raw_results
|
||||||
|
]
|
||||||
return RAGSearchResponse(
|
return RAGSearchResponse(
|
||||||
query=body.query,
|
query=body.query,
|
||||||
results=[RAGCardResult(**r) for r in results],
|
results=results,
|
||||||
total=len(results),
|
total=len(results),
|
||||||
)
|
)
|
||||||
except FileNotFoundError as exc:
|
except FileNotFoundError as exc:
|
||||||
@@ -104,8 +114,8 @@ async def rag_status_detail():
|
|||||||
Detailed RAG index status including card count and age.
|
Detailed RAG index status including card count and age.
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
from services.rag_store import get_index_status
|
from services.rag_store import get_rag_stats
|
||||||
return {"status": "ok", **get_index_status()}
|
return {"status": "ok", **get_rag_stats()}
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
log.exception("RAG status check failed")
|
log.exception("RAG status check failed")
|
||||||
return {"status": "error", "error": str(exc)}
|
return {"status": "error", "error": str(exc)}
|
||||||
+182
-55
@@ -12,6 +12,7 @@ oracle text into the LLM prompt before calling the coach.
|
|||||||
|
|
||||||
Public API:
|
Public API:
|
||||||
ensure_rag_store() -> None (build if missing / stale)
|
ensure_rag_store() -> None (build if missing / stale)
|
||||||
|
build_index(force, batch_size) -> dict (explicit build, returns stats)
|
||||||
query_cards(text, n, colors) -> list[dict]
|
query_cards(text, n, colors) -> list[dict]
|
||||||
get_rag_stats() -> dict
|
get_rag_stats() -> dict
|
||||||
|
|
||||||
@@ -30,27 +31,34 @@ from typing import Optional
|
|||||||
|
|
||||||
logger = logging.getLogger("commander_ai_lab.rag_store")
|
logger = logging.getLogger("commander_ai_lab.rag_store")
|
||||||
|
|
||||||
# ── Config ──────────────────────────────────────────────────────────────────
|
# ── Config ────────────────────────────────────────────────────────────────
|
||||||
RAG_CHROMA_DIR = Path(__file__).parent.parent / "data" / "rag_chroma"
|
RAG_CHROMA_DIR = Path(__file__).parent.parent / "data" / "rag_chroma"
|
||||||
|
BUILD_META_PATH = RAG_CHROMA_DIR / "build_meta.json"
|
||||||
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
|
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||||
EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "nomic-embed-text")
|
EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "nomic-embed-text")
|
||||||
COLLECTION_NAME = "mtg_oracle"
|
COLLECTION_NAME = "mtg_oracle"
|
||||||
BATCH_SIZE = 200
|
DEFAULT_BATCH_SIZE = 500
|
||||||
_rag_lock = threading.Lock()
|
_rag_lock = threading.Lock()
|
||||||
_rag_ready = False
|
_rag_ready = False
|
||||||
|
|
||||||
|
# Module-level singletons to avoid re-creating client/collection on every call
|
||||||
|
_chroma_client = None
|
||||||
|
_chroma_collection = None
|
||||||
|
|
||||||
# ── Helpers ──────────────────────────────────────────────────────────────────
|
|
||||||
|
# ── Helpers ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
def _build_document(card: dict) -> str:
|
def _build_document(card: dict) -> str:
|
||||||
"""Build the text document to embed for a single card.
|
"""Build the text document to embed for a single card.
|
||||||
|
|
||||||
Format: Name | Type | Oracle Text
|
Format: Name | Mana Cost | Type | Oracle Text
|
||||||
This mirrors the structure used by the existing coach embeddings
|
Matches the structure used by coach/embeddings._build_card_text() so
|
||||||
(coach/embeddings.py) so vector similarity is consistent across
|
vector similarity is consistent across both embedding systems.
|
||||||
both systems.
|
|
||||||
"""
|
"""
|
||||||
parts = [card.get("name", "")]
|
parts = [card.get("name", "")]
|
||||||
|
mana_cost = card.get("mana_cost", "")
|
||||||
|
if mana_cost:
|
||||||
|
parts.append(mana_cost)
|
||||||
type_line = card.get("type_line", "")
|
type_line = card.get("type_line", "")
|
||||||
if type_line:
|
if type_line:
|
||||||
parts.append(type_line)
|
parts.append(type_line)
|
||||||
@@ -82,16 +90,20 @@ def _build_metadata(card: dict) -> dict:
|
|||||||
|
|
||||||
|
|
||||||
def _get_chromadb_client():
|
def _get_chromadb_client():
|
||||||
"""Return a persistent ChromaDB client, creating the data dir if needed."""
|
"""Return the module-level persistent ChromaDB client (singleton)."""
|
||||||
|
global _chroma_client
|
||||||
|
if _chroma_client is not None:
|
||||||
|
return _chroma_client
|
||||||
try:
|
try:
|
||||||
import chromadb
|
import chromadb
|
||||||
except ImportError:
|
except ImportError:
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"chromadb is required for RAG Phase 2. "
|
"chromadb is required for RAG Phase 2. "
|
||||||
"Install with: pip install chromadb"
|
"Install with: pip install 'commander-ai-lab[rag]'"
|
||||||
)
|
)
|
||||||
RAG_CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
RAG_CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
return chromadb.PersistentClient(path=str(RAG_CHROMA_DIR))
|
_chroma_client = chromadb.PersistentClient(path=str(RAG_CHROMA_DIR))
|
||||||
|
return _chroma_client
|
||||||
|
|
||||||
|
|
||||||
def _get_ollama_embedding_fn():
|
def _get_ollama_embedding_fn():
|
||||||
@@ -101,7 +113,7 @@ def _get_ollama_embedding_fn():
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"chromadb >= 0.4.0 with Ollama support is required. "
|
"chromadb >= 0.4.0 with Ollama support is required. "
|
||||||
"Install with: pip install chromadb"
|
"Install with: pip install 'commander-ai-lab[rag]'"
|
||||||
)
|
)
|
||||||
return OllamaEmbeddingFunction(
|
return OllamaEmbeddingFunction(
|
||||||
url=f"{OLLAMA_BASE_URL}/api/embeddings",
|
url=f"{OLLAMA_BASE_URL}/api/embeddings",
|
||||||
@@ -110,34 +122,67 @@ def _get_ollama_embedding_fn():
|
|||||||
|
|
||||||
|
|
||||||
def _get_collection():
|
def _get_collection():
|
||||||
"""Return the ChromaDB collection, creating it if needed."""
|
"""Return the ChromaDB collection singleton, creating it if needed."""
|
||||||
|
global _chroma_collection
|
||||||
|
if _chroma_collection is not None:
|
||||||
|
return _chroma_collection
|
||||||
client = _get_chromadb_client()
|
client = _get_chromadb_client()
|
||||||
embed_fn = _get_ollama_embedding_fn()
|
embed_fn = _get_ollama_embedding_fn()
|
||||||
return client.get_or_create_collection(
|
_chroma_collection = client.get_or_create_collection(
|
||||||
name=COLLECTION_NAME,
|
name=COLLECTION_NAME,
|
||||||
embedding_function=embed_fn,
|
embedding_function=embed_fn,
|
||||||
metadata={"hnsw:space": "cosine"},
|
metadata={"hnsw:space": "cosine"},
|
||||||
)
|
)
|
||||||
|
return _chroma_collection
|
||||||
|
|
||||||
|
|
||||||
|
def _invalidate_collection_singleton() -> None:
|
||||||
|
"""Reset the cached collection after a rebuild so the next call re-opens it."""
|
||||||
|
global _chroma_collection
|
||||||
|
_chroma_collection = None
|
||||||
|
|
||||||
|
|
||||||
|
def _read_build_meta() -> dict:
|
||||||
|
"""Read the sidecar build_meta.json, returning {} if missing or corrupt."""
|
||||||
|
try:
|
||||||
|
if BUILD_META_PATH.exists():
|
||||||
|
return json.loads(BUILD_META_PATH.read_text(encoding="utf-8"))
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("Could not read build_meta.json: %s", exc)
|
||||||
|
return {}
|
||||||
|
|
||||||
|
|
||||||
|
def _write_build_meta(card_count: int) -> None:
|
||||||
|
"""Persist build timestamp and card count to sidecar JSON file."""
|
||||||
|
try:
|
||||||
|
RAG_CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
|
BUILD_META_PATH.write_text(
|
||||||
|
json.dumps({"built_at": time.time(), "card_count": card_count}),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("Could not write build_meta.json: %s", exc)
|
||||||
|
|
||||||
|
|
||||||
def _collection_is_current(collection, bulk_card_count: int) -> bool:
|
def _collection_is_current(collection, bulk_card_count: int) -> bool:
|
||||||
"""Check if the ChromaDB collection is reasonably up-to-date.
|
"""Check if the ChromaDB collection is reasonably up-to-date.
|
||||||
|
|
||||||
We compare the stored card count against the bulk DB count. If they
|
Uses sidecar build_meta.json for the age check (survives restarts),
|
||||||
differ by more than 5%, we consider the collection stale and rebuild.
|
and compares card count within a 5% tolerance.
|
||||||
"""
|
"""
|
||||||
chroma_count = collection.count()
|
chroma_count = collection.count()
|
||||||
if chroma_count == 0:
|
if chroma_count == 0:
|
||||||
return False
|
return False
|
||||||
try:
|
|
||||||
meta = collection.metadata or {}
|
# Age check via sidecar file (not collection.metadata which doesn't persist)
|
||||||
built_at = float(meta.get("built_at", 0))
|
meta = _read_build_meta()
|
||||||
age_days = (time.time() - built_at) / 86400
|
built_at = meta.get("built_at", 0)
|
||||||
|
if built_at:
|
||||||
|
age_days = (time.time() - float(built_at)) / 86400
|
||||||
if age_days > 14:
|
if age_days > 14:
|
||||||
logger.info("RAG collection older than 14 days — rebuilding.")
|
logger.info("RAG collection older than 14 days (%0.1f) — rebuilding.", age_days)
|
||||||
return False
|
return False
|
||||||
except (ValueError, TypeError):
|
|
||||||
pass
|
|
||||||
ratio = chroma_count / max(bulk_card_count, 1)
|
ratio = chroma_count / max(bulk_card_count, 1)
|
||||||
if ratio < 0.95 or ratio > 1.05:
|
if ratio < 0.95 or ratio > 1.05:
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -148,12 +193,19 @@ def _collection_is_current(collection, bulk_card_count: int) -> bool:
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
# ── Build pipeline ───────────────────────────────────────────────────────────
|
# ── Build pipeline ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
def _build_collection_from_bulk() -> int:
|
def _build_collection_from_bulk(batch_size: int = DEFAULT_BATCH_SIZE) -> dict:
|
||||||
"""Read all cards from the bulk SQLite DB and insert into ChromaDB.
|
"""Read all cards from the bulk SQLite DB and insert into ChromaDB.
|
||||||
|
|
||||||
Returns the number of cards indexed.
|
Args:
|
||||||
|
batch_size: Number of cards to embed and insert per ChromaDB call.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict with keys: indexed (int), failed (int).
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If more than 1% of cards failed to insert.
|
||||||
"""
|
"""
|
||||||
from services.scryfall_bulk import get_bulk_db
|
from services.scryfall_bulk import get_bulk_db
|
||||||
|
|
||||||
@@ -164,28 +216,33 @@ def _build_collection_from_bulk() -> int:
|
|||||||
"WHERE oracle_text != '' "
|
"WHERE oracle_text != '' "
|
||||||
"ORDER BY edhrec_rank ASC"
|
"ORDER BY edhrec_rank ASC"
|
||||||
)
|
)
|
||||||
|
all_rows = cursor.fetchall()
|
||||||
|
total_rows = len(all_rows)
|
||||||
|
|
||||||
client = _get_chromadb_client()
|
client = _get_chromadb_client()
|
||||||
embed_fn = _get_ollama_embedding_fn()
|
embed_fn = _get_ollama_embedding_fn()
|
||||||
|
|
||||||
|
# Drop and recreate collection to ensure a clean build
|
||||||
try:
|
try:
|
||||||
client.delete_collection(COLLECTION_NAME)
|
client.delete_collection(COLLECTION_NAME)
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
_invalidate_collection_singleton()
|
||||||
|
|
||||||
collection = client.create_collection(
|
collection = client.create_collection(
|
||||||
name=COLLECTION_NAME,
|
name=COLLECTION_NAME,
|
||||||
embedding_function=embed_fn,
|
embedding_function=embed_fn,
|
||||||
metadata={
|
metadata={"hnsw:space": "cosine"},
|
||||||
"hnsw:space": "cosine",
|
|
||||||
"built_at": str(time.time()),
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
|
|
||||||
inserted = 0
|
inserted = 0
|
||||||
|
failed_cards = 0
|
||||||
|
failed_batches = 0
|
||||||
batch_ids: list[str] = []
|
batch_ids: list[str] = []
|
||||||
batch_docs: list[str] = []
|
batch_docs: list[str] = []
|
||||||
batch_metas: list[dict] = []
|
batch_metas: list[dict] = []
|
||||||
|
|
||||||
for row in cursor:
|
for row in all_rows:
|
||||||
card = dict(row)
|
card = dict(row)
|
||||||
doc = _build_document(card)
|
doc = _build_document(card)
|
||||||
if not doc.strip():
|
if not doc.strip():
|
||||||
@@ -199,7 +256,7 @@ def _build_collection_from_bulk() -> int:
|
|||||||
batch_docs.append(doc)
|
batch_docs.append(doc)
|
||||||
batch_metas.append(_build_metadata(card))
|
batch_metas.append(_build_metadata(card))
|
||||||
|
|
||||||
if len(batch_ids) >= BATCH_SIZE:
|
if len(batch_ids) >= batch_size:
|
||||||
try:
|
try:
|
||||||
collection.add(
|
collection.add(
|
||||||
ids=batch_ids,
|
ids=batch_ids,
|
||||||
@@ -208,15 +265,19 @@ def _build_collection_from_bulk() -> int:
|
|||||||
)
|
)
|
||||||
inserted += len(batch_ids)
|
inserted += len(batch_ids)
|
||||||
if inserted % 2000 == 0:
|
if inserted % 2000 == 0:
|
||||||
logger.info(" RAG indexed %d cards...", inserted)
|
logger.info(" RAG indexed %d / %d cards...", inserted, total_rows)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
|
failed_cards += len(batch_ids)
|
||||||
|
failed_batches += 1
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"RAG batch insert failed at %d: %s", inserted, exc
|
"RAG batch insert failed (batch #%d, %d cards lost): %s",
|
||||||
|
failed_batches, len(batch_ids), exc,
|
||||||
)
|
)
|
||||||
batch_ids.clear()
|
batch_ids.clear()
|
||||||
batch_docs.clear()
|
batch_docs.clear()
|
||||||
batch_metas.clear()
|
batch_metas.clear()
|
||||||
|
|
||||||
|
# Final partial batch
|
||||||
if batch_ids:
|
if batch_ids:
|
||||||
try:
|
try:
|
||||||
collection.add(
|
collection.add(
|
||||||
@@ -226,13 +287,71 @@ def _build_collection_from_bulk() -> int:
|
|||||||
)
|
)
|
||||||
inserted += len(batch_ids)
|
inserted += len(batch_ids)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
|
failed_cards += len(batch_ids)
|
||||||
|
failed_batches += 1
|
||||||
logger.warning("RAG final batch insert failed: %s", exc)
|
logger.warning("RAG final batch insert failed: %s", exc)
|
||||||
|
|
||||||
logger.info("RAG vector store complete: %d cards indexed.", inserted)
|
logger.info(
|
||||||
return inserted
|
"RAG vector store complete: %d indexed, %d failed (%d batches).",
|
||||||
|
inserted, failed_cards, failed_batches,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Abort if failure rate exceeds 1% — delete partial collection
|
||||||
|
if total_rows > 0 and (failed_cards / total_rows) > 0.01:
|
||||||
|
logger.error(
|
||||||
|
"RAG build failure rate %.1f%% exceeds 1%% threshold — deleting partial collection.",
|
||||||
|
100.0 * failed_cards / total_rows,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
client.delete_collection(COLLECTION_NAME)
|
||||||
|
_invalidate_collection_singleton()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
raise RuntimeError(
|
||||||
|
f"RAG build aborted: {failed_cards}/{total_rows} cards failed to insert. "
|
||||||
|
"Check Ollama connectivity and retry."
|
||||||
|
)
|
||||||
|
|
||||||
|
_write_build_meta(card_count=inserted)
|
||||||
|
# Invalidate singleton so next _get_collection() picks up the new collection
|
||||||
|
_invalidate_collection_singleton()
|
||||||
|
return {"indexed": inserted, "failed": failed_cards}
|
||||||
|
|
||||||
|
|
||||||
# ── Public API ───────────────────────────────────────────────────────────────
|
# ── Public API ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def build_index(force: bool = False, batch_size: int = DEFAULT_BATCH_SIZE) -> dict:
|
||||||
|
"""Explicitly build or rebuild the ChromaDB vector index.
|
||||||
|
|
||||||
|
This is the function called by POST /api/rag/build.
|
||||||
|
Unlike ensure_rag_store(), this always runs (no _rag_ready short-circuit)
|
||||||
|
and returns stats so the API can surface them to the caller.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
force: If True, rebuild even if the index appears current.
|
||||||
|
batch_size: Cards per ChromaDB insert batch (default 500).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict with keys: indexed (int), failed (int).
|
||||||
|
"""
|
||||||
|
global _rag_ready
|
||||||
|
with _rag_lock:
|
||||||
|
if not force:
|
||||||
|
try:
|
||||||
|
from services.scryfall_bulk import get_stats as bulk_stats
|
||||||
|
stats = bulk_stats()
|
||||||
|
bulk_count = stats.get("card_count", 0)
|
||||||
|
collection = _get_collection()
|
||||||
|
if _collection_is_current(collection, bulk_count):
|
||||||
|
logger.info("RAG index is current — skipping rebuild (use force=True to override).")
|
||||||
|
return {"indexed": collection.count(), "failed": 0, "skipped": True}
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
result = _build_collection_from_bulk(batch_size=batch_size)
|
||||||
|
_rag_ready = True
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
def ensure_rag_store() -> None:
|
def ensure_rag_store() -> None:
|
||||||
"""Build the ChromaDB vector store if it is missing or stale.
|
"""Build the ChromaDB vector store if it is missing or stale.
|
||||||
@@ -294,6 +413,7 @@ def query_cards(
|
|||||||
color_identity: If provided, only return cards whose color identity
|
color_identity: If provided, only return cards whose color identity
|
||||||
is a subset (e.g. ["W", "U", "B"]).
|
is a subset (e.g. ["W", "U", "B"]).
|
||||||
type_filter: Optional type_line substring filter (e.g. "Creature").
|
type_filter: Optional type_line substring filter (e.g. "Creature").
|
||||||
|
Applied as a post-query Python filter for ChromaDB compatibility.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of dicts with keys: name, type_line, oracle_text, cmc,
|
List of dicts with keys: name, type_line, oracle_text, cmc,
|
||||||
@@ -309,16 +429,12 @@ def query_cards(
|
|||||||
logger.error("Failed to get RAG collection: %s", exc)
|
logger.error("Failed to get RAG collection: %s", exc)
|
||||||
return []
|
return []
|
||||||
|
|
||||||
where_filter = None
|
|
||||||
if type_filter:
|
|
||||||
where_filter = {"type_line": {"$contains": type_filter}}
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
fetch_n = n_results * 3 if color_identity else n_results
|
# Fetch extra results to account for post-query color/type filtering
|
||||||
|
fetch_n = n_results * 4 if (color_identity or type_filter) else n_results
|
||||||
results = collection.query(
|
results = collection.query(
|
||||||
query_texts=[query],
|
query_texts=[query],
|
||||||
n_results=min(fetch_n, 50),
|
n_results=min(fetch_n, 200),
|
||||||
where=where_filter,
|
|
||||||
include=["documents", "metadatas", "distances"],
|
include=["documents", "metadatas", "distances"],
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -338,15 +454,23 @@ def query_cards(
|
|||||||
meta = metas[i] if i < len(metas) else {}
|
meta = metas[i] if i < len(metas) else {}
|
||||||
doc = docs[i] if i < len(docs) else ""
|
doc = docs[i] if i < len(docs) else ""
|
||||||
|
|
||||||
|
# Color identity post-filter (Python-side for ChromaDB compat)
|
||||||
if color_identity:
|
if color_identity:
|
||||||
card_ci_str = meta.get("color_identity", "")
|
card_ci_str = meta.get("color_identity", "")
|
||||||
card_ci = set(card_ci_str.split(",")) if card_ci_str else set()
|
card_ci = set(c for c in card_ci_str.split(",") if c) if card_ci_str else set()
|
||||||
deck_ci = set(color_identity)
|
deck_ci = set(color_identity)
|
||||||
if not card_ci.issubset(deck_ci):
|
if not card_ci.issubset(deck_ci):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
parts = doc.split(" | ", 2)
|
# Type post-filter
|
||||||
oracle = parts[2] if len(parts) > 2 else parts[-1] if parts else ""
|
if type_filter:
|
||||||
|
card_type = meta.get("type_line", "")
|
||||||
|
if type_filter.lower() not in card_type.lower():
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Document format: Name | Mana Cost | Type | Oracle Text
|
||||||
|
parts = doc.split(" | ", 3)
|
||||||
|
oracle = parts[3] if len(parts) > 3 else (parts[-1] if parts else "")
|
||||||
|
|
||||||
cards.append({
|
cards.append({
|
||||||
"name": meta.get("name", ""),
|
"name": meta.get("name", ""),
|
||||||
@@ -377,15 +501,18 @@ def get_rag_stats() -> dict:
|
|||||||
try:
|
try:
|
||||||
collection = _get_collection()
|
collection = _get_collection()
|
||||||
result["card_count"] = collection.count()
|
result["card_count"] = collection.count()
|
||||||
meta = collection.metadata or {}
|
|
||||||
built_at = meta.get("built_at")
|
|
||||||
if built_at:
|
|
||||||
result["built_at"] = built_at
|
|
||||||
try:
|
|
||||||
age = (time.time() - float(built_at)) / 86400
|
|
||||||
result["age_days"] = round(age, 1)
|
|
||||||
except (ValueError, TypeError):
|
|
||||||
pass
|
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
# Read age from sidecar file (persists across restarts)
|
||||||
|
meta = _read_build_meta()
|
||||||
|
built_at = meta.get("built_at")
|
||||||
|
if built_at:
|
||||||
|
result["built_at"] = built_at
|
||||||
|
try:
|
||||||
|
age = (time.time() - float(built_at)) / 86400
|
||||||
|
result["age_days"] = round(age, 1)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
pass
|
||||||
|
|
||||||
return result
|
return result
|
||||||
@@ -0,0 +1,374 @@
|
|||||||
|
"""
|
||||||
|
tests/test_rag_store.py
|
||||||
|
=======================
|
||||||
|
Unit tests for services/rag_store.py.
|
||||||
|
|
||||||
|
All tests use mocked ChromaDB and Ollama dependencies — no external services
|
||||||
|
or network access required.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import types
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import MagicMock, patch, mock_open
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Helpers to reset module-level singletons between tests
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def _reset_rag_module():
|
||||||
|
"""Reset all module-level state in rag_store between tests."""
|
||||||
|
import services.rag_store as rs
|
||||||
|
rs._rag_ready = False
|
||||||
|
rs._chroma_client = None
|
||||||
|
rs._chroma_collection = None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def reset_rag_state():
|
||||||
|
_reset_rag_module()
|
||||||
|
yield
|
||||||
|
_reset_rag_module()
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# _build_document
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildDocument:
|
||||||
|
def test_includes_all_four_fields(self):
|
||||||
|
from services.rag_store import _build_document
|
||||||
|
card = {
|
||||||
|
"name": "Rampant Growth",
|
||||||
|
"mana_cost": "{1}{G}",
|
||||||
|
"type_line": "Sorcery",
|
||||||
|
"oracle_text": "Search your library for a basic land card.",
|
||||||
|
}
|
||||||
|
doc = _build_document(card)
|
||||||
|
assert "Rampant Growth" in doc
|
||||||
|
assert "{1}{G}" in doc
|
||||||
|
assert "Sorcery" in doc
|
||||||
|
assert "Search your library" in doc
|
||||||
|
|
||||||
|
def test_mana_cost_present_in_document(self):
|
||||||
|
"""Regression: mana_cost was previously missing (bug #128)."""
|
||||||
|
from services.rag_store import _build_document
|
||||||
|
card = {"name": "Lightning Bolt", "mana_cost": "{R}", "type_line": "Instant", "oracle_text": "Deal 3 damage."}
|
||||||
|
doc = _build_document(card)
|
||||||
|
assert "{R}" in doc
|
||||||
|
|
||||||
|
def test_missing_mana_cost_skipped(self):
|
||||||
|
from services.rag_store import _build_document
|
||||||
|
card = {"name": "Plains", "mana_cost": "", "type_line": "Basic Land", "oracle_text": ""}
|
||||||
|
doc = _build_document(card)
|
||||||
|
# Should not produce " | | " style empty fields
|
||||||
|
assert "Plains" in doc
|
||||||
|
assert doc.count(" | | ") == 0
|
||||||
|
|
||||||
|
def test_empty_card_returns_name_only(self):
|
||||||
|
from services.rag_store import _build_document
|
||||||
|
card = {"name": "Unnamed"}
|
||||||
|
doc = _build_document(card)
|
||||||
|
assert doc == "Unnamed"
|
||||||
|
|
||||||
|
def test_pipe_delimiter_format(self):
|
||||||
|
from services.rag_store import _build_document
|
||||||
|
card = {"name": "A", "mana_cost": "B", "type_line": "C", "oracle_text": "D"}
|
||||||
|
assert _build_document(card) == "A | B | C | D"
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# _build_metadata
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildMetadata:
|
||||||
|
def test_color_identity_json_string_parsed(self):
|
||||||
|
from services.rag_store import _build_metadata
|
||||||
|
card = {"name": "X", "color_identity": '["W", "U"]', "cmc": 3}
|
||||||
|
meta = _build_metadata(card)
|
||||||
|
assert meta["color_identity"] == "U,W" # sorted
|
||||||
|
|
||||||
|
def test_color_identity_list_input(self):
|
||||||
|
from services.rag_store import _build_metadata
|
||||||
|
card = {"name": "X", "color_identity": ["B", "G"], "cmc": 2}
|
||||||
|
meta = _build_metadata(card)
|
||||||
|
assert meta["color_identity"] == "B,G"
|
||||||
|
|
||||||
|
def test_colorless_card(self):
|
||||||
|
from services.rag_store import _build_metadata
|
||||||
|
card = {"name": "Sol Ring", "color_identity": "[]", "cmc": 1}
|
||||||
|
meta = _build_metadata(card)
|
||||||
|
assert meta["color_identity"] == ""
|
||||||
|
|
||||||
|
def test_malformed_color_identity_defaults_empty(self):
|
||||||
|
from services.rag_store import _build_metadata
|
||||||
|
card = {"name": "X", "color_identity": "not-json", "cmc": 0}
|
||||||
|
meta = _build_metadata(card)
|
||||||
|
assert meta["color_identity"] == ""
|
||||||
|
|
||||||
|
def test_cmc_none_defaults_zero(self):
|
||||||
|
from services.rag_store import _build_metadata
|
||||||
|
card = {"name": "X", "cmc": None}
|
||||||
|
meta = _build_metadata(card)
|
||||||
|
assert meta["cmc"] == 0.0
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# _read_build_meta / _write_build_meta
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildMeta:
|
||||||
|
def test_write_then_read_roundtrip(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_dir = rs.RAG_CHROMA_DIR
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.RAG_CHROMA_DIR = tmp_path
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
rs._write_build_meta(card_count=12345)
|
||||||
|
meta = rs._read_build_meta()
|
||||||
|
assert meta["card_count"] == 12345
|
||||||
|
assert abs(meta["built_at"] - time.time()) < 5
|
||||||
|
finally:
|
||||||
|
rs.RAG_CHROMA_DIR = original_dir
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
|
||||||
|
def test_read_missing_returns_empty_dict(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "nonexistent.json"
|
||||||
|
assert rs._read_build_meta() == {}
|
||||||
|
finally:
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# _collection_is_current
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestCollectionIsCurrent:
|
||||||
|
def _make_collection(self, count: int) -> MagicMock:
|
||||||
|
col = MagicMock()
|
||||||
|
col.count.return_value = count
|
||||||
|
return col
|
||||||
|
|
||||||
|
def test_empty_collection_returns_false(self):
|
||||||
|
from services.rag_store import _collection_is_current
|
||||||
|
col = self._make_collection(0)
|
||||||
|
assert _collection_is_current(col, 1000) is False
|
||||||
|
|
||||||
|
def test_current_collection_returns_true(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
rs._write_build_meta(card_count=1000)
|
||||||
|
col = self._make_collection(1000)
|
||||||
|
assert rs._collection_is_current(col, 1000) is True
|
||||||
|
finally:
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
|
||||||
|
def test_stale_by_age_returns_false(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
# Write a built_at 15 days ago
|
||||||
|
rs.BUILD_META_PATH.write_text(
|
||||||
|
json.dumps({"built_at": time.time() - (15 * 86400), "card_count": 1000}),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
col = self._make_collection(1000)
|
||||||
|
assert rs._collection_is_current(col, 1000) is False
|
||||||
|
finally:
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
|
||||||
|
def test_count_mismatch_returns_false(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
rs._write_build_meta(card_count=500)
|
||||||
|
col = self._make_collection(500) # chroma has 500
|
||||||
|
assert rs._collection_is_current(col, 1000) is False # bulk has 1000
|
||||||
|
finally:
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# query_cards
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestQueryCards:
|
||||||
|
def test_returns_empty_when_not_ready(self):
|
||||||
|
from services.rag_store import query_cards
|
||||||
|
# _rag_ready is False (reset by fixture)
|
||||||
|
result = query_cards("draw cards")
|
||||||
|
assert result == []
|
||||||
|
|
||||||
|
def test_color_filter_removes_off_color_cards(self):
|
||||||
|
import services.rag_store as rs
|
||||||
|
rs._rag_ready = True
|
||||||
|
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.query.return_value = {
|
||||||
|
"ids": [["id1", "id2"]],
|
||||||
|
"documents": [["Wrath | {3}{W}{W} | Sorcery | Destroy all", "Dark Ritual | {B} | Instant | Add BBB"]],
|
||||||
|
"metadatas": [[
|
||||||
|
{"name": "Wrath", "type_line": "Sorcery", "color_identity": "W", "cmc": 5.0, "mana_cost": "{3}{W}{W}"},
|
||||||
|
{"name": "Dark Ritual", "type_line": "Instant", "color_identity": "B", "cmc": 1.0, "mana_cost": "{B}"},
|
||||||
|
]],
|
||||||
|
"distances": [[0.1, 0.2]],
|
||||||
|
}
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
results = rs.query_cards("mana acceleration", n_results=10, color_identity=["W"])
|
||||||
|
names = [r["name"] for r in results]
|
||||||
|
assert "Wrath" in names
|
||||||
|
assert "Dark Ritual" not in names
|
||||||
|
|
||||||
|
def test_distance_converted_to_score_in_route(self):
|
||||||
|
"""Verify distance field is present and in [0, 1] range."""
|
||||||
|
import services.rag_store as rs
|
||||||
|
rs._rag_ready = True
|
||||||
|
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.query.return_value = {
|
||||||
|
"ids": [["id1"]],
|
||||||
|
"documents": [["Sol Ring | {1} | Artifact | Add two colorless mana."]],
|
||||||
|
"metadatas": [[
|
||||||
|
{"name": "Sol Ring", "type_line": "Artifact", "color_identity": "", "cmc": 1.0, "mana_cost": "{1}"},
|
||||||
|
]],
|
||||||
|
"distances": [[0.15]],
|
||||||
|
}
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
results = rs.query_cards("mana rock", n_results=1)
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0]["distance"] == 0.15
|
||||||
|
assert results[0]["name"] == "Sol Ring"
|
||||||
|
|
||||||
|
def test_empty_chroma_result_returns_empty_list(self):
|
||||||
|
import services.rag_store as rs
|
||||||
|
rs._rag_ready = True
|
||||||
|
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.query.return_value = {"ids": [[]], "documents": [[]], "metadatas": [[]], "distances": [[]]}
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
assert rs.query_cards("nothing") == []
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# get_rag_stats
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestGetRagStats:
|
||||||
|
def test_returns_expected_keys(self):
|
||||||
|
import services.rag_store as rs
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.count.return_value = 42000
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
stats = rs.get_rag_stats()
|
||||||
|
for key in ["ready", "chroma_dir", "embedding_model", "ollama_url", "card_count", "built_at", "age_days"]:
|
||||||
|
assert key in stats, f"Missing key: {key}"
|
||||||
|
|
||||||
|
def test_card_count_from_collection(self):
|
||||||
|
import services.rag_store as rs
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.count.return_value = 55000
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
stats = rs.get_rag_stats()
|
||||||
|
assert stats["card_count"] == 55000
|
||||||
|
|
||||||
|
def test_ready_false_when_not_initialized(self):
|
||||||
|
from services.rag_store import get_rag_stats
|
||||||
|
stats = get_rag_stats()
|
||||||
|
assert stats["ready"] is False
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# build_index (public API for POST /api/rag/build)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildIndex:
|
||||||
|
def test_skips_rebuild_when_current(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
original_dir = rs.RAG_CHROMA_DIR
|
||||||
|
try:
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
rs.RAG_CHROMA_DIR = tmp_path
|
||||||
|
rs._write_build_meta(card_count=1000)
|
||||||
|
|
||||||
|
mock_col = MagicMock()
|
||||||
|
mock_col.count.return_value = 1000
|
||||||
|
rs._chroma_collection = mock_col
|
||||||
|
|
||||||
|
with patch("services.rag_store._get_collection", return_value=mock_col), \
|
||||||
|
patch("services.scryfall_bulk.get_stats", return_value={"card_count": 1000}):
|
||||||
|
result = rs.build_index(force=False)
|
||||||
|
|
||||||
|
assert result.get("skipped") is True
|
||||||
|
finally:
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
|
rs.RAG_CHROMA_DIR = original_dir
|
||||||
|
|
||||||
|
def test_force_triggers_rebuild(self, tmp_path):
|
||||||
|
import services.rag_store as rs
|
||||||
|
with patch("services.rag_store._build_collection_from_bulk") as mock_build:
|
||||||
|
mock_build.return_value = {"indexed": 500, "failed": 0}
|
||||||
|
result = rs.build_index(force=True, batch_size=250)
|
||||||
|
mock_build.assert_called_once_with(batch_size=250)
|
||||||
|
assert result["indexed"] == 500
|
||||||
|
assert rs._rag_ready is True
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Failure tracking in _build_collection_from_bulk
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildCollectionFailureTracking:
|
||||||
|
def test_high_failure_rate_raises_runtime_error(self, tmp_path):
|
||||||
|
"""If >1% of batches fail, a RuntimeError should be raised."""
|
||||||
|
import services.rag_store as rs
|
||||||
|
|
||||||
|
fake_rows = [
|
||||||
|
{"name": f"Card {i}", "oracle_text": "text", "type_line": "Instant",
|
||||||
|
"mana_cost": "{1}", "cmc": 1, "color_identity": "[]",
|
||||||
|
"rarity": "common", "scryfall_id": f"id-{i}"}
|
||||||
|
for i in range(200)
|
||||||
|
]
|
||||||
|
|
||||||
|
mock_conn = MagicMock()
|
||||||
|
mock_cursor = MagicMock()
|
||||||
|
mock_cursor.fetchall.return_value = [types.SimpleNamespace(**r) for r in fake_rows]
|
||||||
|
mock_conn.execute.return_value = mock_cursor
|
||||||
|
|
||||||
|
mock_collection = MagicMock()
|
||||||
|
mock_collection.add.side_effect = Exception("Ollama timeout")
|
||||||
|
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.create_collection.return_value = mock_collection
|
||||||
|
|
||||||
|
original_dir = rs.RAG_CHROMA_DIR
|
||||||
|
original_path = rs.BUILD_META_PATH
|
||||||
|
try:
|
||||||
|
rs.RAG_CHROMA_DIR = tmp_path
|
||||||
|
rs.BUILD_META_PATH = tmp_path / "build_meta.json"
|
||||||
|
with patch("services.scryfall_bulk.get_bulk_db", return_value=mock_conn), \
|
||||||
|
patch("services.rag_store._get_chromadb_client", return_value=mock_client), \
|
||||||
|
patch("services.rag_store._get_ollama_embedding_fn", return_value=MagicMock()), \
|
||||||
|
patch("services.rag_store._invalidate_collection_singleton"):
|
||||||
|
with pytest.raises(RuntimeError, match="RAG build aborted"):
|
||||||
|
rs._build_collection_from_bulk(batch_size=50)
|
||||||
|
finally:
|
||||||
|
rs.RAG_CHROMA_DIR = original_dir
|
||||||
|
rs.BUILD_META_PATH = original_path
|
||||||
Reference in new issue
Block a user