""" Basic RAG — Baseline Implementation Vector similarity search → top-K chunks → LLM prompt This is the BASELINE we're beating. No graph. No structure. Just cosine similarity over document chunks. """ import os import time import logging from pathlib import Path from dataclasses import dataclass import numpy as np import faiss from dotenv import load_dotenv from ..llm.gemini_client import gemini_generate # Load .env from project root load_dotenv(Path(__file__).parent.parent.parent / ".env", override=True) logger = logging.getLogger(__name__) TOP_K = int((os.getenv("TOP_K_BASIC_RAG", "5") or "5").strip()) EMBEDDING_MODEL = (os.getenv("EMBEDDING_MODEL", "text-embedding-3-small") or "").strip() EMBEDDING_DIM = 384 # sentence-transformers all-MiniLM-L6-v2 _sentence_transformer = None USE_SENTENCE_TRANSFORMERS = (os.getenv("USE_SENTENCE_TRANSFORMERS", "true") or "true").lower().strip() == "true" # Gemini calls go through shared gemini_client.gemini_generate() def _get_sentence_transformer(): """Lazily load sentence-transformer for local embeddings.""" global _sentence_transformer if _sentence_transformer is None: try: from sentence_transformers import SentenceTransformer _sentence_transformer = SentenceTransformer("all-MiniLM-L6-v2") logger.info("✅ Loaded sentence-transformers for local embeddings (384 dims)") except ImportError: logger.warning("sentence-transformers not installed. Install with: pip install sentence-transformers") _sentence_transformer = False return _sentence_transformer if _sentence_transformer else None @dataclass class RAGResult: answer: str retrieved_chunks: list[str] prompt_tokens: int completion_tokens: int total_tokens: int latency_ms: float method: str = "basic_rag" class BasicRAG: """ Standard RAG pipeline: 1. Embed query 2. Cosine similarity search over FAISS index 3. Concatenate top-K chunks 4. Send to LLM This is the BASELINE. Expensive. Indiscriminate. No structural understanding of relationships. """ def __init__(self): self.index = faiss.IndexFlatIP(EMBEDDING_DIM) # Inner product = cosine similarity self.chunks: list[str] = [] self.chunk_metadata: list[dict] = [] # Try to load pre-built FAISS index from disk # Use absolute path: work up from this file → backend → graphrag-hackathon → data faiss_path = Path(__file__).resolve().parent.parent.parent / "data" / "faiss_index.pkl" logger.info(f"Attempting to load FAISS index from: {faiss_path}") logger.info(f"Path exists: {faiss_path.exists()}") if faiss_path.exists(): try: import pickle with open(faiss_path, "rb") as f: data = pickle.load(f) self.index = data["index"] self.chunks = data["chunks"] self.chunk_metadata = data.get("metadata", []) logger.info(f"✅ Loaded FAISS index from {faiss_path} ({len(self.chunks)} chunks, {self.index.ntotal} vectors)") except Exception as e: logger.error(f"❌ Failed to load FAISS index: {e}", exc_info=True) # Fall through to build from sample docs self._build_index_from_sample_docs() else: logger.warning(f"⚠️ FAISS index not found at {faiss_path}. Building from sample documents...") self._build_index_from_sample_docs() def _build_index_from_sample_docs(self): """Dynamically build FAISS index from sample documents with semantic chunking""" sample_docs_dir = Path(__file__).resolve().parent.parent.parent / "data" / "sample_docs" if not sample_docs_dir.exists(): logger.warning(f"Sample docs directory not found at {sample_docs_dir}") return # Load all .md and .txt files doc_files = list(sample_docs_dir.glob("*.md")) + list(sample_docs_dir.glob("*.txt")) if not doc_files: logger.warning(f"No documents found in {sample_docs_dir}") return logger.info(f"Loading {len(doc_files)} documents from {sample_docs_dir}...") for doc_file in doc_files: try: with open(doc_file, "r", encoding="utf-8", errors="ignore") as f: content = f.read() # Semantic chunking: split by sections and paragraphs sections = content.split("##") for section_idx, section in enumerate(sections): # Split each section into paragraphs paragraphs = section.split("\n\n") # Group smaller paragraphs together, keep larger ones separate current_chunk = "" for para in paragraphs: para = para.strip() if not para: continue if len(current_chunk) + len(para) < 1000: # Combine paragraphs up to 1000 chars if current_chunk: current_chunk += "\n\n" current_chunk += para else: if current_chunk: self.chunks.append(current_chunk) self.chunk_metadata.append({ "source": doc_file.name, "section": section_idx, }) current_chunk = para if current_chunk: self.chunks.append(current_chunk) self.chunk_metadata.append({ "source": doc_file.name, "section": section_idx, }) logger.info(f" ✅ Loaded {doc_file.name} ({len(content)} chars → {len([c for c in self.chunks if c.split('###')[0].startswith(doc_file.name.split('.')[0][:20])])} semantic chunks)") except Exception as e: logger.warning(f"Failed to load {doc_file.name}: {e}") # Embed all chunks if self.chunks: try: logger.info(f"Embedding {len(self.chunks)} chunks...") embeddings = self._embed_batch(self.chunks) vectors = np.array(embeddings, dtype="float32") faiss.normalize_L2(vectors) self.index.add(vectors) logger.info(f"✅ Built FAISS index with {self.index.ntotal} vectors from {len(self.chunks)} chunks") except Exception as e: logger.error(f"Failed to embed chunks: {e}") self.chunks = [] # Clear on failure else: logger.warning("No chunks extracted from documents") def add_documents(self, chunks: list[str], metadata: list[dict] = None): """Add document chunks to the FAISS index.""" if not chunks: return logger.info(f"Embedding {len(chunks)} chunks for FAISS index...") embeddings = self._embed_batch(chunks) vectors = np.array(embeddings, dtype="float32") # Normalize for cosine similarity faiss.normalize_L2(vectors) self.index.add(vectors) self.chunks.extend(chunks) self.chunk_metadata.extend(metadata or [{} for _ in chunks]) logger.info(f"FAISS index now has {self.index.ntotal} vectors") def query(self, question: str, top_k: int = TOP_K) -> RAGResult: """ Full RAG pipeline for a single question. Returns answer + full token accounting. """ t0 = time.time() # 1. Embed the query try: query_embedding = self._embed(question) if not query_embedding or len(query_embedding) == 0: raise ValueError("Empty embedding returned") query_vec = np.array([query_embedding], dtype="float32") faiss.normalize_L2(query_vec) except Exception as e: logger.error(f"Embedding failed: {e}. Using fallback retrieval.") query_vec = None # 2. Retrieve top-K similar chunks retrieved = [] if query_vec is not None and self.index.ntotal > 0: try: distances, indices = self.index.search(query_vec, min(top_k, self.index.ntotal)) retrieved = [self.chunks[i] for i in indices[0] if i < len(self.chunks)] logger.info(f"✅ Retrieved {len(retrieved)} chunks from FAISS (ntotal: {self.index.ntotal})") except Exception as e: logger.error(f"FAISS search failed: {e}. Falling back to first chunks.") retrieved = self.chunks[:top_k] elif self.chunks: # Fallback: return first K chunks if embedding/search fails logger.warning(f"⚠️ Embedding or index unavailable. Using first {top_k} chunks as fallback.") retrieved = self.chunks[:top_k] else: retrieved = ["[No documents available. Please ingest documents first.]"] # 3. Build context (this is where tokens pile up) context = "\n\n---\n\n".join(retrieved) # 4. Build prompt system_prompt = ( "You are an expert assistant. Use the provided context when available. " "If context is insufficient, leverage your knowledge to provide a helpful answer. " "Always prioritize context when it's available." ) user_prompt = f"""Context from documents: {context} Question: {question} Provide a comprehensive answer using the context if available, or general knowledge:""" # 5. Call Gemini via shared client (accurate token counts) result = gemini_generate( system_prompt=system_prompt, user_prompt=user_prompt, temperature=0.1, max_tokens=1000, ) latency_ms = (time.time() - t0) * 1000 return RAGResult( answer=result["answer"], retrieved_chunks=retrieved, prompt_tokens=result["prompt_tokens"], completion_tokens=result["completion_tokens"], total_tokens=result["total_tokens"], latency_ms=latency_ms, method="basic_rag", ) def _embed(self, text: str) -> list[float]: """Embed a single text string using OpenAI, then sentence-transformers, then zeros.""" # Try OpenAI first try: from openai import OpenAI openai_client = OpenAI() response = openai_client.embeddings.create( model=EMBEDDING_MODEL, input=text[:8191], ) return response.data[0].embedding except Exception as e: logger.debug(f"OpenAI embedding failed: {e}") # Try sentence-transformers if USE_SENTENCE_TRANSFORMERS: try: model = _get_sentence_transformer() if model: embedding = model.encode(text[:8191]) return embedding.tolist() except Exception as e: logger.debug(f"Sentence-transformers embedding failed: {e}") # Fallback to zeros logger.warning("Falling back to zero embeddings") return [0.0] * 384 def _embed_batch(self, texts: list[str]) -> list[list[float]]: """Embed multiple texts efficiently using OpenAI, then sentence-transformers, then zeros.""" # Try OpenAI first (batch of 100) try: from openai import OpenAI openai_client = OpenAI() all_embeddings = [] batch_size = 100 for i in range(0, len(texts), batch_size): batch = texts[i:i + batch_size] response = openai_client.embeddings.create( model=EMBEDDING_MODEL, input=[t[:8191] for t in batch], ) all_embeddings.extend([item.embedding for item in response.data]) return all_embeddings except Exception as e: logger.debug(f"OpenAI batch embedding failed: {e}") # Try sentence-transformers (local, no API call) if USE_SENTENCE_TRANSFORMERS: try: model = _get_sentence_transformer() if model: embeddings = model.encode([t[:8191] for t in texts]) return embeddings.tolist() except Exception as e: logger.debug(f"Sentence-transformers batch embedding failed: {e}") # Fallback to zeros logger.warning(f"Falling back to zero embeddings for {len(texts)} texts") return [[0.0] * 384 for _ in texts] def count_tokens(self, text: str) -> int: from ..llm.gemini_client import count_tokens_gemini return count_tokens_gemini(text) def get_index_stats(self) -> dict: return { "total_chunks": len(self.chunks), "index_size": self.index.ntotal, "embedding_dim": EMBEDDING_DIM, }