Bunnnyyy2005/Smart_engineering_RAG

1import os
2from langchain_community.document_loaders import PyPDFDirectoryLoader
3from langchain_text_splitters import RecursiveCharacterTextSplitter
4from langchain_huggingface import HuggingFaceEmbeddings
5from langchain_community.vectorstores import Chroma
6
7# Set up our directories
8DATA_DIR = "./data"
9CHROMA_DB_DIR = "./chroma_db"
10
11# Using a very fast, free, and standard embedding model
12embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
13
14def build_vector_database():
15    print("Loading engineering documents from data/ folder...")
16    loader = PyPDFDirectoryLoader(DATA_DIR)
17    documents = loader.load()
18    
19    if not documents:
20        print("❌ No PDFs found in the data/ folder! Please add some manuals.")
21        return None
22
23    print(f"✅ Found {len(documents)} pages. Splitting text into chunks...")
24    # Chunking: 1000 characters per piece, with 200 overlap so we don't cut sentences in half
25    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
26    chunks = text_splitter.split_documents(documents)
27
28    print("🧠 Building ChromaDB vector store. This might take a minute on the first run...")
29    vector_db = Chroma.from_documents(
30        documents=chunks, 
31        embedding=embeddings, 
32        persist_directory=CHROMA_DB_DIR
33    )
34    print("✅ Vector database built successfully! 🚀")
35    return vector_db
36
37def get_retriever():
38    """This function will be used by our LangChain Agent later to search the database."""
39    vector_db = Chroma(persist_directory=CHROMA_DB_DIR, embedding_function=embeddings)
40    # Return the top 3 most relevant chunks for any given question
41    return vector_db.as_retriever(search_kwargs={"k": 3})
42
43# When we run this file directly, it will build the database.
44if __name__ == "__main__":
45    build_vector_database()