app.py

# Modules for environment variables
import os
from dotenv import load_dotenv

# Module for Interface (UI)
import gradio as gr

# Modules for Sourcing Data, transforming, embedding, and storing
from langchain_community.document_loaders import (
    PyPDFDirectoryLoader,
    PyPDFLoader,
    Docx2txtLoader,
    TextLoader,
    CSVLoader,
    UnstructuredPowerPointLoader,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

# Modules for Document chain, retriever, and retrieval chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain 
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_classic.chains import create_retrieval_chain

# Get the env variables
load_dotenv()
os.environ['GROQ_API_KEY'] = os.getenv('GROQ_API_KEY')

# Load embeddings model
embeddings = HuggingFaceEmbeddings(
  model_name = 'sentence-transformers/all-MiniLM-L6-v2',
  model_kwargs = {"device": "cpu"},
  encode_kwargs = {"batch_size": 8}
)

# Load the source data and chunk it
KNOWLEDGE_DIR = os.path.join(os.path.dirname(__file__), "knowledge")
loader = PyPDFDirectoryLoader(KNOWLEDGE_DIR)
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=200)
documents = text_splitter.split_documents(docs)

# Embed the data and store it in the database if it is not already stored, if it is stored, then just load it from local index file
if os.path.exists("faiss_index"):
  db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
else:
  db = FAISS.from_documents(documents, embeddings)
  db.save_local("faiss_index")

# Initialize the LLM model and generate a rag prompt
llm = ChatGroq(
  model_name="llama-3.3-70b-versatile"
)
rag_prompt = ChatPromptTemplate.from_template(
  """
    Answer the questions based only on the provided context. Don't hallucinate. Don't reply to anything else that is not in the context. Make sure you provide the most accurate answer and I will give you praise when you will provide the best accurate answer which the user finds helpful.
    <context> {context} </context>
    Question: {input}
  """
)

# Create retriever, document chain, and retrieval chain for retrieving the embeddings from the vector store
retriever = db.as_retriever()
document_chain = create_stuff_documents_chain(llm, rag_prompt)
retrieval_chain = create_retrieval_chain(retriever, document_chain)

# For multiple file types uploaded by the user
def load_file_by_extension(file_path):
  ext = os.path.splitext(file_path)[-1].lower()
  try:
    if ext == ".pdf":
      return PyPDFLoader(file_path).load()
    elif ext == ".docx":
      return Docx2txtLoader(file_path).load()
    elif ext == ".txt":
      return TextLoader(file_path).load()
    elif ext == ".csv":
      return CSVLoader(file_path).load()
    elif ext == ".pptx":
      return UnstructuredPowerPointLoader(file_path).load()
    else:
            print(f"Unsupported file type: {ext}")
            return []
  except Exception as e:
    print(f"Error loading {file_path}: {e}")
    return []

# Function for gradio chat interface
def response(message, history):
  user_text = message["text"]
  user_files = message["files"]

  if user_files:
    new_docs = []
    for file in user_files:
      file_path = file["path"] if isinstance(file, dict) else file
      loaded = load_file_by_extension(file_path)
      new_docs.extend(loaded)
    
    split_new_docs = text_splitter.split_documents(new_docs)
    db.add_documents(split_new_docs)
  

  response = retrieval_chain.invoke({"input": user_text})
  answer = response["answer"]

  return answer

# Gradio Initialization
with gr.Blocks(title='RAG Knowledge Based Chatbot') as demo:
  gr.Markdown('RAG Chatbot')
  gr.Markdown('Ask anything about your uploaded documents and about mr chips questions and answers')

  chatbot = gr.ChatInterface(
    fn=response,
    multimodal=True,
    chatbot=gr.Chatbot(height=420, show_label=False),
    textbox=gr.MultimodalTextbox(
      placeholder='Ask anything about your uploaded files',
      file_types=[".pdf", ".txt", ".docx", ".csv", ".xlsx", "image"],
      container=False,
      scale=8
    ),
    examples=[
      "What is Logistics Regression?",
      "What services did Brookfield render during the war?",
      "When and how did Katherine die?"
    ]
  )

demo.launch()