import os import csv import glob from pypdf import PdfReader from llama_index.core import VectorStoreIndex, Settings from llama_index.core.node_parser import SentenceSplitter from llama_index.core.schema import Document from llama_index.embeddings.huggingface import HuggingFaceEmbedding # ------------------------------------------------------------------------- # LlamaIndex configuration # - Embeddings: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 (Supports 50+ languages) # - LLM: disabled — we manage the LLM ourselves in src/llm.py # ------------------------------------------------------------------------- Settings.embed_model = HuggingFaceEmbedding( model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" ) Settings.llm = None _index: VectorStoreIndex | None = None # ── File loaders ───────────────────────────────────────────────────────── def _load_csv(filepath: str) -> list[Document]: """Loads a CSV with 'Question' (and optionally 'Answer') columns.""" docs = [] try: with open(filepath, newline="", encoding="utf-8") as fh: reader = csv.DictReader(fh) if reader.fieldnames is None or "Question" not in reader.fieldnames: print(f"[RAG] Skipping {filepath}: no 'Question' column.") return docs has_answer = "Answer" in (reader.fieldnames or []) for row in reader: question = row.get("Question", "").strip() answer = row.get("Answer", "").strip() if has_answer else "" if question: text = f"Q: {question}\nA: {answer}" if answer else f"Q: {question}" docs.append(Document(text=text, metadata={"source": filepath})) except Exception as exc: print(f"[RAG] Error reading CSV {filepath}: {exc}") return docs def _load_pdf(filepath: str) -> list[Document]: """Extracts text from each page of a PDF as separate documents.""" docs = [] try: reader = PdfReader(filepath) basename = os.path.basename(filepath) for i, page in enumerate(reader.pages): text = page.extract_text() if text and text.strip(): docs.append(Document( text=text.strip(), metadata={"source": basename, "page": i + 1}, )) print(f"[RAG] Loaded {len(docs)} pages from {basename}") except Exception as exc: print(f"[RAG] Error reading PDF {filepath}: {exc}") return docs def _load_text(filepath: str) -> list[Document]: """Loads a plain text file as a single document.""" docs = [] try: with open(filepath, "r", encoding="utf-8") as fh: text = fh.read().strip() if text: docs.append(Document(text=text, metadata={"source": filepath})) except Exception as exc: print(f"[RAG] Error reading text file {filepath}: {exc}") return docs # ── Vector store init ──────────────────────────────────────────────────── def init_vector_store(data_path: str = "data") -> None: """ Scans the data directory and loads ALL supported files: - .csv → FAQ-style (Question/Answer columns) - .pdf → RBI guidelines, circulars, etc. - .txt → Plain text documents Builds an in-memory vector index from all loaded content. """ global _index # Support both file path and directory path for backward compatibility if os.path.isfile(data_path): data_dir = os.path.dirname(data_path) else: data_dir = data_path if not os.path.isdir(data_dir): print(f"[RAG] Warning: data directory not found at '{data_dir}'. Vector store empty.") return documents: list[Document] = [] # Auto-discover all supported files for pattern, loader in [("*.csv", _load_csv), ("*.pdf", _load_pdf), ("*.txt", _load_text)]: for filepath in sorted(glob.glob(os.path.join(data_dir, pattern))): loaded = loader(filepath) documents.extend(loaded) if not documents: print("[RAG] Warning: No documents loaded from any file.") return try: splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50) nodes = splitter.get_nodes_from_documents(documents) _index = VectorStoreIndex(nodes) print(f"[RAG] Loaded {len(documents)} documents → {len(nodes)} chunks into vector store.") except Exception as exc: print(f"[RAG] Error building vector index: {exc}") # ── Retrieval ──────────────────────────────────────────────────────────── def retrieve_context(query: str, k: int = 5) -> str: """ Returns the top-k most relevant passages for a query string. Falls back to an informative message if the index is not ready. """ if not _index: return "Knowledge base not available." if not query or not query.strip(): return "Knowledge base not available." try: retriever = _index.as_retriever(similarity_top_k=k) results = retriever.retrieve(query) if not results: return "No relevant information found in the knowledge base." return "\n\n---\n\n".join(node.get_content() for node in results) except Exception as exc: print(f"[RAG] Error during retrieval: {exc}") return "Knowledge base lookup failed." def retrieve_context_multi(english_query: str, hindi_text: str, k: int = 15) -> str: """ Searches using BOTH the English query and the original Hindi text, deduplicates, and returns the top-k most relevant passages. """ if not _index: return "Knowledge base not available." seen_content: set[str] = set() all_results = [] try: retriever = _index.as_retriever(similarity_top_k=k) # Search with English query if english_query and english_query.strip(): for node in retriever.retrieve(english_query): content = node.get_content() if content not in seen_content: seen_content.add(content) all_results.append((node.get_score(), content)) # Search with Hindi text (backup) if hindi_text and hindi_text.strip(): for node in retriever.retrieve(hindi_text): content = node.get_content() if content not in seen_content: seen_content.add(content) all_results.append((node.get_score(), content)) if not all_results: return "No relevant information found in the knowledge base." all_results.sort(key=lambda x: x[0], reverse=True) top_results = [content for _, content in all_results[:k]] print(f"[RAG] Retrieved {len(top_results)} unique passages from dual query") return "\n\n---\n\n".join(top_results) except Exception as exc: print(f"[RAG] Error during retrieval: {exc}") return "Knowledge base lookup failed."