fez5587/commander-ai-lab
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rag): batch_size param, failure tracking, sidecar build_meta, mana_cost in doc (closes #124, #126, #127, #128)

Browse Source
This commit is contained in:
KoalaTrapLord committed 2026-03-27 12:13:36 -04:00
1 parent 1d756bca49
commit 96e51f59ae
1 file changed
+182 -55
Download Patch File Download Diff File Expand all files Collapse all files
services/rag_store.py
+182 -55
View File
@@ -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