KoalaTrapLord
4b0cd34eda
* feat(rag): Phase 2 — ChromaDB vector index over bulk card DB * feat(rag): add routes/rag.py with build, search, status endpoints * Update lab_api.pyfeat(rag): register rag_router in lab_api.py, remove old inline rag_status * Update requirements.txtfeat(rag): add chromadb>=0.5.0 to requirements.txtfeat(rag): add chromadb to requirements.txt
112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
"""
|
|
routes/rag.py
|
|
=============
|
|
RAG (Retrieval-Augmented Generation) endpoints for semantic card search.
|
|
|
|
Endpoints:
|
|
POST /api/rag/build — build / rebuild the ChromaDB vector index
|
|
POST /api/rag/search — semantic search over indexed cards
|
|
GET /api/rag/status — index health & stats
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
from fastapi import APIRouter, HTTPException
|
|
from pydantic import BaseModel, Field
|
|
|
|
log = logging.getLogger("commander_ai_lab.routes.rag")
|
|
|
|
router = APIRouter(tags=["rag"])
|
|
|
|
|
|
# ── Request / Response models ────────────────────────────────────────────────
|
|
|
|
class RAGSearchRequest(BaseModel):
|
|
query: str = Field(..., description="Natural-language card search query")
|
|
color_identity: Optional[List[str]] = Field(
|
|
None,
|
|
description="Optional color identity filter, e.g. ['W','U','B']",
|
|
)
|
|
top_k: int = Field(10, ge=1, le=50, description="Number of results to return")
|
|
|
|
|
|
class RAGCardResult(BaseModel):
|
|
name: str
|
|
type_line: str = ""
|
|
oracle_text: str = ""
|
|
color_identity: str = ""
|
|
score: float = 0.0
|
|
|
|
|
|
class RAGSearchResponse(BaseModel):
|
|
query: str
|
|
results: List[RAGCardResult]
|
|
total: int
|
|
|
|
|
|
class RAGBuildRequest(BaseModel):
|
|
force: bool = Field(False, description="Force full rebuild even if index exists")
|
|
batch_size: int = Field(500, ge=50, le=5000)
|
|
|
|
|
|
# ── Endpoints ────────────────────────────────────────────────────────────────
|
|
|
|
@router.post("/api/rag/build")
|
|
async def rag_build(body: RAGBuildRequest = RAGBuildRequest()):
|
|
"""
|
|
Build or rebuild the ChromaDB vector index from the Scryfall bulk DB.
|
|
Runs synchronously (may take 1-5 min on first build).
|
|
"""
|
|
try:
|
|
from services.rag_store import build_index
|
|
result = build_index(force=body.force, batch_size=body.batch_size)
|
|
return {"status": "ok", **result}
|
|
except FileNotFoundError as exc:
|
|
raise HTTPException(status_code=404, detail=str(exc))
|
|
except Exception as exc:
|
|
log.exception("RAG build failed")
|
|
raise HTTPException(status_code=500, detail=str(exc))
|
|
|
|
|
|
@router.post("/api/rag/search")
|
|
async def rag_search(body: RAGSearchRequest):
|
|
"""
|
|
Semantic search over the ChromaDB card index.
|
|
Returns ranked card results with relevance scores.
|
|
"""
|
|
try:
|
|
from services.rag_store import search_cards
|
|
results = search_cards(
|
|
query=body.query,
|
|
color_identity=body.color_identity,
|
|
top_k=body.top_k,
|
|
)
|
|
return RAGSearchResponse(
|
|
query=body.query,
|
|
results=[RAGCardResult(**r) for r in results],
|
|
total=len(results),
|
|
)
|
|
except FileNotFoundError as exc:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"Index not built yet. Call POST /api/rag/build first. {exc}",
|
|
)
|
|
except Exception as exc:
|
|
log.exception("RAG search failed")
|
|
raise HTTPException(status_code=500, detail=str(exc))
|
|
|
|
|
|
@router.get("/api/rag/status")
|
|
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()}
|
|
except Exception as exc:
|
|
log.exception("RAG status check failed")
|
|
return {"status": "error", "error": str(exc)}
|