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:
KoalaTrapLord authored and GitHub committed 2026-03-27 13:24:41 -04:00
1 parent 61ac349445
commit bcf558418a
5 files changed
+614 -72

No files matched your search

+20 -11
View File
@@ -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")
+22
View File
@@ -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
View File
@@ -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
View File
@@ -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
+374
View File
@@ -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