chatbot.py

import os
import uuid
import numpy as np
from pathlib import Path
from typing import List, Dict, Any
from dotenv import load_dotenv
from langchain_community.document_loaders import ( # loading
    PyPDFLoader,
    PyMuPDFLoader,
    TextLoader,
    CSVLoader,
    JSONLoader,
    UnstructuredHTMLLoader,
    UnstructuredWordDocumentLoader,
    UnstructuredPowerPointLoader,
    UnstructuredExcelLoader
)
from langchain_text_splitters import RecursiveCharacterTextSplitter # chunking
from sentence_transformers import SentenceTransformer # embeddings
import chromadb # vector store
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq


# Load and parse PDFs

def load_any_document(file_path: str):
    file_ext = Path(file_path).suffix.lower()

    if file_ext in [".pdf"]:
        try:
            return PyMuPDFLoader(file_path).load()
        except:
            return PyPDFLoader(file_path).load()

    elif file_ext in [".txt"]:
        return TextLoader(file_path).load()

    elif file_ext in [".csv"]:
        return CSVLoader(file_path).load()

    elif file_ext in [".json"]:
        return JSONLoader(file_path).load()

    elif file_ext in [".docx"]:
        return UnstructuredWordDocumentLoader(file_path).load()

    elif file_ext in [".pptx"]:
        return UnstructuredPowerPointLoader(file_path).load()

    elif file_ext in [".xlsx", ".xls"]:
        return UnstructuredExcelLoader(file_path).load()

    elif file_ext in [".html", ".htm"]:
        return UnstructuredHTMLLoader(file_path).load()

    else:
        print(f"❌ Unsupported file type: {file_ext}")
        return []

def process_all_files(directory: str):
    all_docs = []
    data_dir = Path(directory)

    supported_exts = [
        "*.pdf", "*.txt", "*.csv", "*.json",
        "*.docx", "*.pptx", "*.xlsx", "*.html"
    ]

    files = []
    for pattern in supported_exts:
        files.extend(list(data_dir.glob(f"**/{pattern}")))

    print(f"Found {len(files)} documents to process.")

    for file_path in files:
        print(f"\n📄 Loading {file_path.name}")
        try:
            docs = load_any_document(str(file_path))
            for d in docs:
                d.metadata["source_file"] = file_path.name
            all_docs.extend(docs)

        except Exception as e:
            print(f"❌ Error loading {file_path.name}: {e}")

    print(f"\nTotal loaded documents: {len(all_docs)}")
    return all_docs


# 2️ Split text into smaller chunks

def split_documents(documents, chunk_size=1000, chunk_overlap=200):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", " ", ""],
    )
    split_docs = text_splitter.split_documents(documents)
    print(f"Split {len(documents)} documents into {len(split_docs)} chunks")

    if split_docs:
        print("\nExample chunk:")
        print(f"Content: {split_docs[0].page_content[:200]}...")
        print(f"Metadata: {split_docs[0].metadata}")

    return split_docs


# Embedding Manager — SentenceTransformer


class EmbeddingManager:

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self.model = None
        self._load_model()

    def _load_model(self):
        try:
            print(f"Loading embedding model: {self.model_name}")
            self.model = SentenceTransformer(self.model_name)
            print(
                f"Model loaded successfully. Embedding dimension: {self.model.get_sentence_embedding_dimension()}"
            )
        except Exception as e:
            print(f"Error loading model {self.model_name}: {e}")
            raise

    def generate_embeddings(self, texts: List[str]) -> np.ndarray:
        if not self.model:
            raise ValueError("Model not loaded")

        print(f"Generating embeddings for {len(texts)} texts...")
        embeddings = self.model.encode(texts, show_progress_bar=True)
        print(f"Generated embeddings with shape: {embeddings.shape}")
        return embeddings


#  Vector Store — ChromaDB (vector storage)


class VectorStore:
    def __init__(self, collection_name: str = "pdf_documents", persist_directory: str = "data/vector_store"):
        self.collection_name = collection_name
        self.persist_directory = persist_directory
        self.client = None
        self.collection = None
        self._initialize_store()

    def _initialize_store(self):
        """Initialize ChromaDB client and collection"""
        try:
            os.makedirs(self.persist_directory, exist_ok=True)
            self.client = chromadb.PersistentClient(path=self.persist_directory)
            self.collection = self.client.get_or_create_collection(
                name=self.collection_name,
                metadata={"description": "PDF document embeddings for RAG"},
            )
            print(f"Vector store initialized. Collection: {self.collection_name}")
            print(f"Existing documents in collection: {self.collection.count()}")
        except Exception as e:
            print(f"Error initializing vector store: {e}")
            raise

    def add_documents(self, documents: List[Any], embeddings: np.ndarray):
        """Add documents and their embeddings to vector store"""
        if len(documents) != len(embeddings):
            raise ValueError("Number of documents must match number of embeddings")

        print(f"Adding {len(documents)} documents to vector store...")

        ids, metadatas, documents_text, embeddings_list = [], [], [], []

        for i, (doc, embedding) in enumerate(zip(documents, embeddings)):
            doc_id = f"doc_{uuid.uuid4().hex[:8]}_{i}"
            ids.append(doc_id)

            metadata = dict(doc.metadata)
            metadata["doc_index"] = i
            metadata["content_length"] = len(doc.page_content)
            metadatas.append(metadata)

            documents_text.append(doc.page_content)
            embeddings_list.append(embedding.tolist())

        try:
            self.collection.add(
                ids=ids,
                embeddings=embeddings_list,
                metadatas=metadatas,
                documents=documents_text,
            )
            print(f"✅ Added {len(documents)} documents to vector store")
            print(f"📚 Total documents in collection: {self.collection.count()}")
        except Exception as e:
            print(f"Error adding documents to vector store: {e}")
            raise
    
    # to add doucments from outside
    def add_single_document(self, doc, embedding):
        """Add one new document chunk + embedding to the vector store"""
        try:
            doc_id = f"doc_{uuid.uuid4().hex[:8]}"
            metadata = dict(doc.metadata)
            metadata["content_length"] = len(doc.page_content)

            self.collection.add(
                ids=[doc_id],
                embeddings=[embedding.tolist()],
                metadatas=[metadata],
                documents=[doc.page_content],
            )

            print(f"📌 Added new document chunk: {doc_id}")
        except Exception as e:
            print(f"Error adding single document to vector store: {e}")
            raise


#  RAG Retriever — Query & Retrieve

class RAGRetriever:

    def __init__(self, vector_store: VectorStore, embedding_manager: EmbeddingManager):
        self.vector_store = vector_store
        self.embedding_manager = embedding_manager

    def retrieve(
        self, query: str, top_k: int = 5, score_threshold: float = 0.0) -> List[Dict[str, Any]]:
        print(f"\n🔍 Retrieving documents for query: '{query}'")
        print(f"Top K: {top_k}, Score threshold: {score_threshold}")

        query_embedding = self.embedding_manager.generate_embeddings([query])[0]

        try:
            results = self.vector_store.collection.query(
                query_embeddings=[query_embedding.tolist()],
                n_results=top_k,
            )

            retrieved_docs = []
            if results["documents"] and results["documents"][0]:
                documents = results["documents"][0]
                metadatas = results["metadatas"][0]
                distances = results["distances"][0]
                ids = results["ids"][0]

                for i, (doc_id, document, metadata, distance) in enumerate(
                    zip(ids, documents, metadatas, distances)
                ):
                    similarity_score = 1 - distance
                    if similarity_score >= score_threshold:
                        retrieved_docs.append(
                            {
                                "id": doc_id,
                                "content": document,
                                "metadata": metadata,
                                "similarity_score": similarity_score,
                                "distance": distance,
                                "rank": i + 1,
                            }
                        )

                print(f"Retrieved {len(retrieved_docs)} documents (after filtering)")
            else:
                print("No documents found")

            return retrieved_docs

        except Exception as e:
            print(f"Error during retrieval: {e}")
            return []



#  Run a Sample Query (without LLM integration)

'''
rag_retriever = RAGRetriever(vectorstore, embedding_manager)
results = rag_retriever.retrieve("LLM Fine-tuning", top_k=3) # type in any query u want

print("\n================= QUERY RESULTS =================")
for r in results:
    print(f"\nRank: {r['rank']}, Score: {r['similarity_score']:.4f}")
    print(f"Source: {r['metadata'].get('source_file', 'Unknown')}")
    print(f"Content snippet: {r['content'][:300]}...")
    
'''

# Function to build RAG pipeline

def build_rag_pipeline():
    # 1. load documents
    all_documents = process_all_files("data")

    # 2. split chunks
    chunks = split_documents(all_documents)

    # 3. embeddings
    embedding_manager = EmbeddingManager()
    texts = [doc.page_content for doc in chunks]
    embeddings = embedding_manager.generate_embeddings(texts)

    # 4. vector store
    vectorstore = VectorStore()
    vectorstore.add_documents(chunks, embeddings)

    # 5. load retriever
    embeddings2 = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    chroma_vs = Chroma(
        persist_directory="data/vector_store",
        embedding_function=embeddings2,
        collection_name="pdf_documents"
    )
    retriever = chroma_vs.as_retriever(search_kwargs={"k": 3})

    # 6. LLM
    load_dotenv()
    groq_key = os.getenv("GROQ_API_KEY")
    llm = ChatGroq(model="llama-3.1-8b-instant", groq_api_key=groq_key)

    # 7. template + chain
    template = """
    You are an AI assistant that answers questions using the context provided.

    You must:
    - Use ONLY the provided context (do not hallucinate).
    - Provide a clear and helpful answer.
    - Include a "Sources:" section at the end.
    - For each source, cite it like this:
    [source_file | similarity_score]
    and include 1-2 lines of the relevant extracted text.

    Context (retrieved documents):
    {context}

    Question:
    {question}

    Answer format example:
    The number of factors of a number is given by increasing each exponent by 1 and multiplying.

    Sources:
    1) Aptitude.pdf | 0.92
    “Number of Factors = (a+1)(b+1)(c+1)”
    """

    prompt = ChatPromptTemplate.from_template(template)

    rag_chain = (
        {"context": retriever | (lambda docs: "\n\n".join(d.page_content for d in docs)),
         "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    return rag_chain, vectorstore, embedding_manager, llm, prompt


#adding documnets from outside
def load_any_document(file_path: str):
    ext = Path(file_path).suffix.lower()

    if ext == ".pdf":
        try:
            return PyMuPDFLoader(file_path).load()
        except:
            return PyPDFLoader(file_path).load()

    if ext == ".txt":
        return TextLoader(file_path).load()

    if ext == ".docx":
        return UnstructuredWordDocumentLoader(file_path).load()

    if ext == ".csv":
        return CSVLoader(file_path).load()

    if ext == ".json":
        return JSONLoader(file_path).load()

    return []

def chunk_single_doc(doc):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    return splitter.split_documents([doc])


def embed_chunks(chunks, embedding_manager):
    texts = [chunk.page_content for chunk in chunks]
    return embedding_manager.generate_embeddings(texts)