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)