""" rag.py — Plexi RAG Engine ========================= Handles everything related to the LlamaIndex vector index: - Downloading the pre-built index from GitHub - Loading HuggingFace sentence-transformer embeddings - Embedding queries and retrieving top-k chunks scoped by semester + subject - Extracting text from PDFs for full-context fallback - Formatting retrieved chunks for the LLM system prompt """ import io import os import shutil # API-BUG-7: temp-dir cleanup after index is in memory import tempfile from pathlib import Path from typing import TypedDict # API-BUG-4: explicit typed return from retrieve_chunks import requests # --------------------------------------------------------------------------- # Optional LlamaIndex — graceful degradation if not installed # --------------------------------------------------------------------------- try: from llama_index.core import Settings, StorageContext, load_index_from_storage from llama_index.embeddings.huggingface import HuggingFaceEmbedding LLAMA_INDEX_AVAILABLE = True except ImportError: LLAMA_INDEX_AVAILABLE = False # API-BUG-5: MetadataFilters let the vector store do scope-filtering internally, # avoiding the over-fetch window that could miss relevant chunks. try: from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters METADATA_FILTERS_AVAILABLE = True except ImportError: try: from llama_index.core.vector_stores import MetadataFilter, MetadataFilters METADATA_FILTERS_AVAILABLE = True except ImportError: METADATA_FILTERS_AVAILABLE = False try: import PyPDF2 PYPDF2_AVAILABLE = True except ImportError: PYPDF2_AVAILABLE = False # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- MATERIALS_REPO = os.getenv("MATERIALS_REPO", "KunalGupta25/plexi-materials") MANIFEST_BRANCH = os.getenv("MANIFEST_BRANCH", "main") EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2" INDEX_FILES = [ "default__vector_store.json", "docstore.json", "graph_store.json", "image__vector_store.json", "index_store.json", ] DEFAULT_TOP_K = 5 # --------------------------------------------------------------------------- # Typed chunk return (API-BUG-4) # --------------------------------------------------------------------------- class ChunkDict(TypedDict): text: str score: float | None filename: str | None subject: str | None # --------------------------------------------------------------------------- # Index loading (called once at FastAPI startup via asyncio.to_thread) # --------------------------------------------------------------------------- def load_index(): """ Download the pre-built LlamaIndex from the materials repo and return a VectorStoreIndex ready for querying. Returns (index, error_msg). index is None if loading failed. API-BUG-7: The temp directory is always removed in the finally block. After load_index_from_storage() returns, all data is in memory — the files on disk are no longer needed. """ if not LLAMA_INDEX_AVAILABLE: return None, "llama-index-core is not installed." index_base_url = ( f"https://raw.githubusercontent.com/{MATERIALS_REPO}/{MANIFEST_BRANCH}/index" ) index_dir = tempfile.mkdtemp(prefix="plexi_index_") try: # Download each index shard from GitHub for filename in INDEX_FILES: url = f"{index_base_url}/{filename}" try: resp = requests.get(url, timeout=30) resp.raise_for_status() with open(os.path.join(index_dir, filename), "wb") as fh: fh.write(resp.content) except Exception as err: return None, f"Failed to download index file '{filename}': {err}" # Build the in-memory index from the downloaded shards try: embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL_ID) Settings.embed_model = embed_model Settings.llm = None storage_ctx = StorageContext.from_defaults(persist_dir=index_dir) index = load_index_from_storage(storage_ctx) return index, None except Exception as err: return None, f"Failed to load index from storage: {err}" finally: # API-BUG-7: always wipe the temp dir — data is now in RAM. shutil.rmtree(index_dir, ignore_errors=True) # --------------------------------------------------------------------------- # Retrieval # --------------------------------------------------------------------------- def _matches_scope(node, semester: str, subject: str) -> bool: """Return True when a retrieved node belongs to the active semester + subject.""" metadata = getattr(node.node, "metadata", {}) or {} return ( metadata.get("semester") == semester and metadata.get("subject") == subject ) def retrieve_chunks( index, query: str, semester: str, subject: str, top_k: int = DEFAULT_TOP_K, ) -> list[ChunkDict]: """ Embed the query, retrieve top-k chunks from the index scoped to the given semester + subject. API-BUG-4: Returns list[ChunkDict] (TypedDict) instead of list[dict] so type errors surface at the source rather than silently at Pydantic serialization time. API-BUG-5: Primary path uses MetadataFilters so the vector store does the scope-gating internally — no risk of the over-fetch window failing to reach chunks that belong to the active subject. Falls back to the generous over-fetch + manual filter approach when MetadataFilters are unavailable (e.g., older llama-index-core builds). """ if index is None: return [] try: if METADATA_FILTERS_AVAILABLE: # Primary: vector store filters by metadata at query time. filters = MetadataFilters( filters=[ MetadataFilter(key="semester", value=semester), MetadataFilter(key="subject", value=subject), ] ) retriever = index.as_retriever(similarity_top_k=top_k, filters=filters) nodes = retriever.retrieve(query) else: # Fallback: over-fetch (generous 10× window) + manual scope filter. retriever = index.as_retriever(similarity_top_k=max(top_k * 10, 50)) nodes = retriever.retrieve(query) nodes = [n for n in nodes if _matches_scope(n, semester, subject)] return [ ChunkDict( text=node.node.get_content(), score=round(float(node.score), 4) if node.score is not None else None, filename=(getattr(node.node, "metadata", {}) or {}).get("filename"), subject=(getattr(node.node, "metadata", {}) or {}).get("subject"), ) for node in nodes[:top_k] ] except Exception as err: print(f"Retrieval error: {err}") return [] # --------------------------------------------------------------------------- # Context formatting (for system prompt injection) # --------------------------------------------------------------------------- def format_context(chunks: list[ChunkDict]) -> str: """Format retrieved chunks as a numbered block for the LLM system prompt.""" if not chunks: return "(No relevant context retrieved for this query.)" parts = [] for i, chunk in enumerate(chunks, start=1): score_info = f" [relevance: {chunk['score']}]" if chunk.get("score") else "" source = chunk.get("filename") or chunk.get("subject") or "Unknown source" parts.append( f"--- Chunk {i} | {source}{score_info} ---\n{chunk['text']}\n" ) return "\n".join(parts) # --------------------------------------------------------------------------- # PDF text extraction # NOTE: read_pdf_text() is currently unused by the API routes — retained for # future use (CQ-4). load_embed_model() removed — was dead code (API-BUG-6). # --------------------------------------------------------------------------- def read_pdf_text(pdf_bytes: bytes) -> str: """Extract plain text from PDF bytes. Returns empty string on failure.""" if not PYPDF2_AVAILABLE: return "" text_parts = [] try: reader = PyPDF2.PdfReader(io.BytesIO(pdf_bytes)) for page in reader.pages: try: page_text = page.extract_text() if page_text: # Sanitise surrogate pairs that can appear in some PDFs filtered = page_text.encode("utf-16", "surrogatepass").decode( "utf-16", "ignore" ) text_parts.append(filtered) except Exception: pass except Exception: return pdf_bytes.decode("utf-8", errors="ignore") if pdf_bytes else "" return "\n".join(text_parts)