app.py

import gradio as gr
import PyPDF2
import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
from langchain.vectorstores import Chroma
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# Initialize global variables for RAG components
rag_chain = None
vector_store = None
llm = None
embedding_model = None
embedding_function = None

# Set your Groq API key here - Read from environment variable
# os.environ["GROQ_API_KEY"] = "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" # Removed hardcoded key
groq_api_key = os.environ.get("GROQ_API_KEY")

def build_rag_from_pdf(pdf_path):
    """Builds the RAG chain from an uploaded PDF."""
    global rag_chain, vector_store, llm, embedding_model, embedding_function

    if not groq_api_key:
        error_message = "Groq API key not set. Please set the GROQ_API_KEY environment variable."
        print(error_message)
        return False, error_message

    pdf_text = ""
    try:
        with open(pdf_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            for page in reader.pages:
                pdf_text += page.extract_text()
        print("PDF text extracted successfully.")
    except FileNotFoundError:
        error_message = f"Error: The file '{pdf_path}' was not found."
        print(error_message)
        return False, error_message
    except Exception as e:
        error_message = f"An error occurred during PDF extraction: {e}"
        print(error_message)
        return False, error_message

    if not pdf_text:
        error_message = "Extracted text is empty."
        print(error_message)
        return False, error_message

    # Split text into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    text_chunks = text_splitter.split_text(pdf_text)
    print(f"Split PDF text into {len(text_chunks)} chunks.")

    if not text_chunks:
        error_message = "No text chunks created."
        print(error_message)
        return False, error_message

    # Create embeddings
    try:
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        embeddings = embedding_model.encode(text_chunks)
        print(f"Generated embeddings for {len(text_chunks)} text chunks.")
    except Exception as e:
        error_message = f"An error occurred during embedding generation: {e}"
        print(error_message)
        return False, error_message

    # Set up vector store
    try:
        embedding_function = SentenceTransformerEmbeddings(model_name='all-MiniLM-L6-v2')
        vector_store = Chroma.from_texts(texts=text_chunks, embedding=embedding_function)
        print("Chroma vector store created successfully.")
    except Exception as e:
        error_message = f"An error occurred during vector store creation: {e}"
        print(error_message)
        return False, error_message

    # Initialize LLM
    try:
        llm = ChatGroq(model_name="openai/gpt-oss-20b", groq_api_key=groq_api_key) # Pass API key
        print("Groq language model set up successfully.")
    except Exception as e:
        error_message = f"An error occurred during LLM setup: {e}"
        print(error_message)
        return False, error_message

    # Build RAG chain
    try:
        template = """Use the following pieces of context to answer queries about the provided document efficiently. If you don't know the answer, just say that you don't know, don't try to make up an answer.

        {context}

        Question: {question}
        Helpful Answer:"""
        QA_CHAIN_PROMPT = PromptTemplate.from_template(template)

        rag_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=vector_store.as_retriever(),
            chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
        )
        print("RAG chain built successfully.")
        return True, "Document processed. You can now ask questions."
    except Exception as e:
        error_message = f"An error occurred while building the RAG chain: {e}"
        print(error_message)
        rag_chain = None
        return False, error_message

def chat_with_rag(message, history):
    """Handles chat interactions using the RAG chain."""
    global rag_chain
    print(f"chat_with_rag received history type: {type(history)}, content: {history}") # Debug print
    if rag_chain is not None:
        try:
            response = rag_chain.invoke(message)
            # Append message and response in the 'messages' format
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": response['result']})
            print(f"chat_with_rag returning history type: {type(history)}, content: {history}") # Debug print
            return history
        except Exception as e:
            error_message = f"An error occurred while querying the RAG chain: {e}"
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": error_message})
            print(f"chat_with_rag returning history with error type: {type(history)}, content: {history}") # Debug print
            return history
    else:
        error_message = "Please upload a document to start the chat."
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": error_message})
        print(f"chat_with_rag returning history with no RAG chain type: {type(history)}, content: {history}") # Debug print
        return history

def process_file(file):
    """Processes the uploaded file and builds the RAG chain."""
    if file is None:
        print("process_file received None file.") # Debug print
        return "Please upload a file.", None
    print(f"Processing file: {file.name}") # Debug print
    success, message = build_rag_from_pdf(file.name)
    if success:
        print("File processing successful, returning message and None for history.") # Debug print
        # Return None for chatbot history to clear it on successful upload
        return message, None
    else:
        print(f"File processing failed, returning message and empty list for history. Message: {message}") # Debug print
        return message, []

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## RAG Chatbot with Document Upload")
    file_upload = gr.File(label="Upload your document (PDF)")
    output_message = gr.Textbox(label="Status")
    chatbot = gr.Chatbot(type='messages') # Explicitly set type to 'messages'
    msg = gr.Textbox(label="Your Question")
    clear = gr.ClearButton([msg, chatbot])

    file_upload.upload(process_file, inputs=[file_upload], outputs=[output_message, chatbot])
    msg.submit(chat_with_rag, inputs=[msg, chatbot], outputs=[chatbot])

if __name__ == "__main__":
    demo.launch(share=True)