diff --git a/lab_api.py b/lab_api.py index b5e1e8c..f581a0f 100644 --- a/lab_api.py +++ b/lab_api.py @@ -50,7 +50,7 @@ _allowed = os.environ.get("ALLOWED_ORIGINS", _DEFAULT_ORIGINS) 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 async def _lifespan(application: FastAPI): # ──── STARTUP ──── @@ -62,13 +62,21 @@ async def _lifespan(application: FastAPI): if not COMMANDER_META: load_commander_meta() + loop = asyncio.get_event_loop() + # RAG Phase 1: kick off Scryfall bulk download in background thread. # run_in_executor returns a Future; we do NOT await it so startup # completes immediately and the download runs concurrently. from services.scryfall_bulk import ensure_bulk_db - loop = asyncio.get_event_loop() 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 # ──── SHUTDOWN ──── @@ -76,7 +84,7 @@ async def _lifespan(application: FastAPI): await async_close_db() -# ── App ──────────────────────────────────────────────────────────────── +# ── App ──────────────────────────────────────────────────────────────────── app = FastAPI(title="Commander AI Lab API", version="3.0.0", lifespan=_lifespan) app.add_middleware( CORSMiddleware, @@ -100,7 +108,7 @@ app.include_router(ws_game_router) app.include_router(ml_router) -# ── API endpoints ──────────────────────────────────────────────────────────── +# ── API endpoints ──────────────────────────────────────────────────────────────── # NOTE: all @app.get routes MUST be registered before app.mount() calls. # Once StaticFiles is mounted at "/", it registers a wildcard catch-all that # shadows any routes added after it. @@ -111,8 +119,9 @@ async def health_check(): return {"status": "ok"} # ── Static UI (MUST come after all @app.get routes) ──────────────────────────── -_legacy_ui_dir = Path(__file__).parent / "ui" -_spa_dir = Path(__file__).parent / "frontend" / "commander-ai-lab-ui" / "dist" +from pathlib import Path as _Path +_legacy_ui_dir = _Path(__file__).parent / "ui" +_spa_dir = _Path(__file__).parent / "frontend" / "commander-ai-lab-ui" / "dist" if _legacy_ui_dir.exists(): 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")) -# ── Argument parsing + auto-detection helpers ───────────────────────────────── +# ── Argument parsing + auto-detection helpers ─────────────────────────────────── def _parse_args(): p = argparse.ArgumentParser(description="Commander AI Lab API Server") p.add_argument("--forge-jar", default=os.environ.get("FORGE_JAR", "")) @@ -298,15 +307,15 @@ def main(): init_coach_service() # 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 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 ensure_rag_store() + import uvicorn uvicorn.run(app, host="0.0.0.0", port=CFG.port, log_level="info") diff --git a/pyproject.toml b/pyproject.toml index 4daf5aa..9b0166d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,28 @@ dev = [ "httpx>=0.24.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] where = ["src"] diff --git a/routes/rag.py b/routes/rag.py index d3839ff..b2a37e3 100644 --- a/routes/rag.py +++ b/routes/rag.py @@ -77,15 +77,25 @@ async def rag_search(body: RAGSearchRequest): Returns ranked card results with relevance scores. """ try: - from services.rag_store import search_cards - results = search_cards( + from services.rag_store import query_cards + raw_results = query_cards( query=body.query, + n_results=body.top_k, 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( query=body.query, - results=[RAGCardResult(**r) for r in results], + results=results, total=len(results), ) except FileNotFoundError as exc: @@ -104,8 +114,8 @@ async def rag_status_detail(): Detailed RAG index status including card count and age. """ try: - from services.rag_store import get_index_status - return {"status": "ok", **get_index_status()} + from services.rag_store import get_rag_stats + return {"status": "ok", **get_rag_stats()} except Exception as exc: log.exception("RAG status check failed") return {"status": "error", "error": str(exc)} diff --git a/services/rag_store.py b/services/rag_store.py index a989eff..d839df2 100644 --- a/services/rag_store.py +++ b/services/rag_store.py @@ -12,6 +12,7 @@ oracle text into the LLM prompt before calling the coach. Public API: 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] get_rag_stats() -> dict @@ -30,27 +31,34 @@ from typing import Optional logger = logging.getLogger("commander_ai_lab.rag_store") -# ── Config ────────────────────────────────────────────────────────────────── +# ── Config ──────────────────────────────────────────────────────────────── 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") EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "nomic-embed-text") COLLECTION_NAME = "mtg_oracle" -BATCH_SIZE = 200 +DEFAULT_BATCH_SIZE = 500 _rag_lock = threading.Lock() _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: """Build the text document to embed for a single card. - Format: Name | Type | Oracle Text - This mirrors the structure used by the existing coach embeddings - (coach/embeddings.py) so vector similarity is consistent across - both systems. + Format: Name | Mana Cost | Type | Oracle Text + Matches the structure used by coach/embeddings._build_card_text() so + vector similarity is consistent across both embedding systems. """ parts = [card.get("name", "")] + mana_cost = card.get("mana_cost", "") + if mana_cost: + parts.append(mana_cost) type_line = card.get("type_line", "") if type_line: parts.append(type_line) @@ -82,16 +90,20 @@ def _build_metadata(card: dict) -> dict: 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: import chromadb except ImportError: raise ImportError( "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) - 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(): @@ -101,7 +113,7 @@ def _get_ollama_embedding_fn(): except ImportError: raise ImportError( "chromadb >= 0.4.0 with Ollama support is required. " - "Install with: pip install chromadb" + "Install with: pip install 'commander-ai-lab[rag]'" ) return OllamaEmbeddingFunction( url=f"{OLLAMA_BASE_URL}/api/embeddings", @@ -110,34 +122,67 @@ def _get_ollama_embedding_fn(): 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() embed_fn = _get_ollama_embedding_fn() - return client.get_or_create_collection( + _chroma_collection = client.get_or_create_collection( name=COLLECTION_NAME, embedding_function=embed_fn, 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: """Check if the ChromaDB collection is reasonably up-to-date. - We compare the stored card count against the bulk DB count. If they - differ by more than 5%, we consider the collection stale and rebuild. + Uses sidecar build_meta.json for the age check (survives restarts), + and compares card count within a 5% tolerance. """ chroma_count = collection.count() if chroma_count == 0: return False - try: - meta = collection.metadata or {} - built_at = float(meta.get("built_at", 0)) - age_days = (time.time() - built_at) / 86400 + + # Age check via sidecar file (not collection.metadata which doesn't persist) + meta = _read_build_meta() + built_at = meta.get("built_at", 0) + if built_at: + age_days = (time.time() - float(built_at)) / 86400 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 - except (ValueError, TypeError): - pass + ratio = chroma_count / max(bulk_card_count, 1) if ratio < 0.95 or ratio > 1.05: logger.info( @@ -148,12 +193,19 @@ def _collection_is_current(collection, bulk_card_count: int) -> bool: 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. - 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 @@ -164,28 +216,33 @@ def _build_collection_from_bulk() -> int: "WHERE oracle_text != '' " "ORDER BY edhrec_rank ASC" ) + all_rows = cursor.fetchall() + total_rows = len(all_rows) client = _get_chromadb_client() embed_fn = _get_ollama_embedding_fn() + + # Drop and recreate collection to ensure a clean build try: client.delete_collection(COLLECTION_NAME) except Exception: pass + _invalidate_collection_singleton() + collection = client.create_collection( name=COLLECTION_NAME, embedding_function=embed_fn, - metadata={ - "hnsw:space": "cosine", - "built_at": str(time.time()), - }, + metadata={"hnsw:space": "cosine"}, ) inserted = 0 + failed_cards = 0 + failed_batches = 0 batch_ids: list[str] = [] batch_docs: list[str] = [] batch_metas: list[dict] = [] - for row in cursor: + for row in all_rows: card = dict(row) doc = _build_document(card) if not doc.strip(): @@ -199,7 +256,7 @@ def _build_collection_from_bulk() -> int: batch_docs.append(doc) batch_metas.append(_build_metadata(card)) - if len(batch_ids) >= BATCH_SIZE: + if len(batch_ids) >= batch_size: try: collection.add( ids=batch_ids, @@ -208,15 +265,19 @@ def _build_collection_from_bulk() -> int: ) inserted += len(batch_ids) 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: + failed_cards += len(batch_ids) + failed_batches += 1 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_docs.clear() batch_metas.clear() + # Final partial batch if batch_ids: try: collection.add( @@ -226,13 +287,71 @@ def _build_collection_from_bulk() -> int: ) inserted += len(batch_ids) except Exception as exc: + failed_cards += len(batch_ids) + failed_batches += 1 logger.warning("RAG final batch insert failed: %s", exc) - logger.info("RAG vector store complete: %d cards indexed.", inserted) - return inserted + logger.info( + "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: """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 is a subset (e.g. ["W", "U", "B"]). type_filter: Optional type_line substring filter (e.g. "Creature"). + Applied as a post-query Python filter for ChromaDB compatibility. Returns: 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) return [] - where_filter = None - if type_filter: - where_filter = {"type_line": {"$contains": type_filter}} - 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( query_texts=[query], - n_results=min(fetch_n, 50), - where=where_filter, + n_results=min(fetch_n, 200), include=["documents", "metadatas", "distances"], ) except Exception as exc: @@ -338,15 +454,23 @@ def query_cards( meta = metas[i] if i < len(metas) else {} doc = docs[i] if i < len(docs) else "" + # Color identity post-filter (Python-side for ChromaDB compat) if 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) if not card_ci.issubset(deck_ci): continue - parts = doc.split(" | ", 2) - oracle = parts[2] if len(parts) > 2 else parts[-1] if parts else "" + # Type post-filter + 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({ "name": meta.get("name", ""), @@ -377,15 +501,18 @@ def get_rag_stats() -> dict: try: collection = _get_collection() 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: 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 diff --git a/tests/test_rag_store.py b/tests/test_rag_store.py new file mode 100644 index 0000000..9889517 --- /dev/null +++ b/tests/test_rag_store.py @@ -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