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