kumardatascience/Multi-Source-RAG-AI-System-with-Query-Routing

1"""2Day 4 — Document indexing.3Reads files from data/documents, chunks them, embeds them, stores in ChromaDB.4Run this as a script, not imported by main.py yet.5"""6 7from pathlib import Path8import chromadb9from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageContext, Settings10from llama_index.core.node_parser import SentenceSplitter11from llama_index.embeddings.huggingface import HuggingFaceEmbedding12from llama_index.vector_stores.chroma import ChromaVectorStore13 14# --- Configuration ---15PROJECT_ROOT = Path(__file__).resolve().parents[2]16DOCUMENTS_DIR = PROJECT_ROOT / "data" / "documents"17CHROMA_DIR = PROJECT_ROOT / "chroma_db"18COLLECTION_NAME = "themis_docs"19 20# Free, local embedding model. Runs on your Mac, no API key needed.21EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"22 23 24def build_index():25    """Read docs → chunk them → embed them → store in ChromaDB."""26 27    # 1. Tell LlamaIndex to use our local embedding model (no OpenAI)28    Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)29    Settings.node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)30 31    # 2. Load all files in data/documents (PDFs, txt, md, docx, etc.)32    print(f"📂 Reading documents from {DOCUMENTS_DIR}")33    reader = SimpleDirectoryReader(input_dir=str(DOCUMENTS_DIR))34    documents = reader.load_data()35    print(f"   Loaded {len(documents)} document(s)")36 37    # 3. Connect to ChromaDB (creates the folder if missing)38    print(f"💾 Connecting to ChromaDB at {CHROMA_DIR}")39    chroma_client = chromadb.PersistentClient(path=str(CHROMA_DIR))40    chroma_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)41 42    # 4. Wrap ChromaDB so LlamaIndex can talk to it43    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)44    storage_context = StorageContext.from_defaults(vector_store=vector_store)45 46    # 5. The magic step: chunk → embed → store. One line.47    print("🧠 Chunking, embedding, and storing...")48    VectorStoreIndex.from_documents(documents, storage_context=storage_context)49 50    # 6. Report51    count = chroma_collection.count()52    print(f"✅ Done! {count} chunks stored in collection '{COLLECTION_NAME}'.")53 54 55def get_retriever(top_k: int = 3):56    """57    Open the existing ChromaDB and return a retriever object.58    top_k = how many chunks to fetch per query.59    """60    # Use the same embedding model as indexing (must match!)61    Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)62 63    # Open the existing ChromaDB (don't re-index)64    chroma_client = chromadb.PersistentClient(path=str(CHROMA_DIR))65    chroma_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)66 67    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)68 69    # Wrap the existing store as a searchable index70    index = VectorStoreIndex.from_vector_store(vector_store=vector_store)71 72    return index.as_retriever(similarity_top_k=top_k)73 74 75def retrieve_context(question: str, top_k: int = 3) -> str:76    """77    Given a user's question, return the most relevant chunks joined as a single string.78    """79    retriever = get_retriever(top_k=top_k)80    nodes = retriever.retrieve(question)81 82    if not nodes:83        return ""84 85    # Combine chunks with a separator86    return "\n\n---\n\n".join(node.get_content() for node in nodes)    87 88# ============================================================89# Session-scoped indexing for user-uploaded files90# ============================================================91 92from llama_index.core import Document93from llama_index.core.readers import SimpleDirectoryReader as _Reader94 95 96# In-memory store of session collections: {session_id: chroma_collection}97_session_collections: dict = {}98 99 100def index_file_for_session(file_path: str, session_id: str) -> int:101    """102    Read a single file, chunk it, embed it, and store in a session-scoped103    ChromaDB collection. Returns the number of chunks stored.104    """105    # Configure LlamaIndex (same as build_index)106    Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)107    Settings.node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)108 109    # Read just this one file110    documents = _Reader(input_files=[file_path]).load_data()111 112    # Create an ephemeral (in-memory) ChromaDB client for this session113    if session_id not in _session_collections:114        ephemeral_client = chromadb.EphemeralClient()115        collection_name = f"session_{session_id[:20]}"116        _session_collections[session_id] = ephemeral_client.get_or_create_collection(117            collection_name118        )119 120    chroma_collection = _session_collections[session_id]121    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)122    storage_context = StorageContext.from_defaults(vector_store=vector_store)123 124    VectorStoreIndex.from_documents(documents, storage_context=storage_context)125    return chroma_collection.count()126 127 128def retrieve_context_with_session(129    question: str, session_id: str | None = None, top_k: int = 3130) -> str:131    """132    Retrieve chunks from the session collection (if any) AND the main collection,133    combined into a single context string.134    """135    context_parts = []136 137    # 1. Search the session collection (uploaded files) if it exists138    if session_id and session_id in _session_collections:139        Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)140        session_collection = _session_collections[session_id]141        session_vector_store = ChromaVectorStore(chroma_collection=session_collection)142        session_index = VectorStoreIndex.from_vector_store(143            vector_store=session_vector_store144        )145        session_retriever = session_index.as_retriever(similarity_top_k=top_k)146        session_nodes = session_retriever.retrieve(question)147        if session_nodes:148            context_parts.append(149                "--- FROM UPLOADED FILE ---\n"150                + "\n\n".join(n.get_content() for n in session_nodes)151            )152 153    # 2. Search the main pre-indexed collection154    main_context = retrieve_context(question, top_k=top_k)155    if main_context:156        context_parts.append(main_context)157 158    return "\n\n".join(context_parts)159 160 161def clear_session(session_id: str) -> None:162    """Wipe a session's uploaded-file collection when the chat ends."""163    if session_id in _session_collections:164        del _session_collections[session_id]165 166 167if __name__ == "__main__":168    build_index()