ingest.py

import os
import argparse
import nbformat
from tqdm import tqdm
from langchain_community.document_loaders import PyPDFLoader, TextLoader, UnstructuredMarkdownLoader
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_core.documents import Document

# Supported file extensions
SUPPORTED_EXTENSIONS = {'.pdf', '.md', '.txt', '.ipynb'}

def load_notebook(file_path):
    """Custom loader for Jupyter Notebooks."""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            nb = nbformat.read(f, as_version=4)
        
        text = ""
        for cell in nb.cells:
            if cell.cell_type == 'markdown':
                text += cell.source + "\n\n"
            elif cell.cell_type == 'code':
                text += f"```python\n{cell.source}\n```\n\n"
        
        return [Document(page_content=text, metadata={"source": file_path})]
    except Exception as e:
        print(f"Error loading notebook {file_path}: {e}")
        return []

def load_documents(source_dir):
    """Loads documents from the specified directory recursively."""
    documents = []
    print(f"Scanning directory: {source_dir}")

    for root, _, files in os.walk(source_dir):
        for file in files:
            file_path = os.path.join(root, file)
            ext = os.path.splitext(file)[1].lower()

            if ext not in SUPPORTED_EXTENSIONS:
                continue

            try:
                if ext == '.pdf':
                    loader = PyPDFLoader(file_path)
                    documents.extend(loader.load())
                elif ext == '.md':
                    # Use TextLoader for markdown to keep it simple and dependency-free
                    documents.extend(TextLoader(file_path, encoding='utf-8').load())
                elif ext == '.txt':
                    loader = TextLoader(file_path, encoding='utf-8')
                    documents.extend(loader.load())
                elif ext == '.ipynb':
                    documents.extend(load_notebook(file_path))
                
                print(f"Loaded: {file_path}")
            except Exception as e:
                print(f"Failed to load {file_path}: {e}")

    return documents

def ingest_documents(source_dir, persist_dir="./chroma_db"):
    """Ingests documents into ChromaDB."""
    
    # 1. Load Documents
    raw_documents = load_documents(source_dir)
    if not raw_documents:
        print("No supported documents found.")
        return

    print(f"Total documents loaded: {len(raw_documents)}")

    # 2. Split Text
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
    )
    chunks = text_splitter.split_documents(raw_documents)
    print(f"Total chunks created: {len(chunks)}")

    # 3. Create Embeddings & Store in ChromaDB
    print("Creating embeddings and storing in ChromaDB... (This may take a while)")
    embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
    
    # Check if DB exists to append or create new
    if os.path.exists(persist_dir):
        print(f"Appending to existing database at {persist_dir}")
    else:
        print(f"Creating new database at {persist_dir}")

    db = Chroma.from_documents(
        documents=chunks,
        embedding=embedding_function,
        persist_directory=persist_dir
    )
    
    print(f"Success! Knowledge base created/updated at: {persist_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Ingest Data Engineering docs into local RAG.")
    parser.add_argument("--source_dir", type=str, required=True, help="Path to the directory containing documents.")
    parser.add_argument("--persist_dir", type=str, default="./chroma_db", help="Path to store the vector database.")
    
    args = parser.parse_args()
    
    if not os.path.exists(args.source_dir):
        print(f"Error: Directory '{args.source_dir}' not found.")
    else:
        ingest_documents(args.source_dir, args.persist_dir)