app.py

import gradio as gr
from pdfminer.high_level import extract_text
from langchain_groq import ChatGroq
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.schema import Document
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
import os
import markdown2

# Retrieve API keys from HF secrets
openai_api_key = os.getenv('OPENAI_API_KEY')
groq_api_key = os.getenv('GROQ_API_KEY')
google_api_key = os.getenv('GEMINI_API_KEY')

# Initialize API clients with the API keys
openai_client = ChatOpenAI(model_name="gpt-4o", api_key=openai_api_key)
groq_client = ChatGroq(model="llama-3.1-70b-versatile", temperature=0, api_key=groq_api_key)
gemini_client = ChatGoogleGenerativeAI(model="gemini-1.5-pro", api_key=google_api_key)

# Define paths for regulation PDFs 
regulation_pdfs = {
    "GDPR": "GDPR.pdf",
    "FERPA": "FERPA.pdf",
    "COPPA": "COPPA.pdf"
}

# Function to extract text from PDF
def extract_pdf(pdf_path):
    return extract_text(pdf_path)

# Function to split text into chunks
def split_text(text):
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    return [Document(page_content=t) for t in splitter.split_text(text)]

# Function to generate embeddings and store in vector database
def generate_embeddings(docs):
    embeddings = OpenAIEmbeddings(api_key=openai_api_key)
    return FAISS.from_documents(docs, embeddings)

# Function for query preprocessing and simple HyDE-Lite
def preprocess_query(query):
    prompt = ChatPromptTemplate.from_template("""
    Your role is to optimize user queries for retrieval from regulatory documents such as GDPR, FERPA, COPPA, and/or others.
    Transform the query into a more affirmative, keyword-focused statement.
    The transformed query should look like probable related passages in the official documents.
    Query: {query}
    Optimized query:
    """)
    chain = prompt | openai_client
    return chain.invoke({"query": query}).content

# Function to create RAG chain with Groq
def create_rag_chain(vector_store):
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are an AI assistant helping with regulatory compliance queries. Use the following context from the official regulatory documents to answer the user's question:\n\n{context}"),
        ("human", "{input}")
    ])
    document_chain = create_stuff_documents_chain(groq_client, prompt)
    return create_retrieval_chain(vector_store.as_retriever(), document_chain)

# Function for Gemini response with long context
def gemini_response(query, full_content):
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are an AI assistant helping with regulatory compliance queries. Use the following full content of the official regulatory documents to answer the user's question:\n\n{context}"),
        ("human", "{input}")
    ])
    chain = prompt | gemini_client
    return chain.invoke({"context": full_content, "input": query}).content

# Function to generate final response
def generate_final_response(response1, response2):
    prompt = ChatPromptTemplate.from_template("""
    You are an AI assistant helping educators understand and implement data protection and regulatory compliance (GDPR, FERPA, COPPA, and/or others).
    Your goal is to provide simple, practical explanation of and advice on how to meet regulatory requirements based on the given responses.
    To do so:
    1. Analyze the following two responses. Inspect their content, and highlight differences. This MUST be done 
    internally as a hidden state. 
    2. Then, use this information to output your own response combining the best from both.
    If the responses differ or contradict each other on important points, include that in your response.
    Only output your own response.
    """)
    chain = prompt | openai_client
    return chain.invoke({"response1": response1, "response2": response2}).content

def markdown_to_html(content):
    return markdown2.markdown(content)

def load_pdfs(selected_regulations, additional_pdfs):
    global full_pdf_content, vector_store, rag_chain
    
    documents = []
    full_pdf_content = ""
    
    # Load selected regulation PDFs
    for regulation in selected_regulations:
        if regulation in regulation_pdfs:
            pdf_content = extract_pdf(regulation_pdfs[regulation])
            full_pdf_content += pdf_content + "\n\n"
            documents.extend(split_text(pdf_content))
            print(f"Loaded {regulation} PDF")
    
    # Load additional user-uploaded PDFs
    if additional_pdfs is not None:
        for pdf_file in additional_pdfs:
            pdf_content = extract_pdf(pdf_file.name)
            full_pdf_content += pdf_content + "\n\n"
            documents.extend(split_text(pdf_content))
            print(f"Loaded additional PDF: {pdf_file.name}")
    
    if not documents:
        return "No PDFs were selected or uploaded. Please select at least one regulation or upload a PDF."
    
    vector_store = generate_embeddings(documents)
    rag_chain = create_rag_chain(vector_store)
    
    return "PDFs loaded and RAG system updated successfully!"

def process_query(user_query):
    global rag_chain, full_pdf_content
    
    if rag_chain is None or not full_pdf_content:
        return ("Please load PDFs before asking questions.", 
                "Please load PDFs before asking questions.", 
                "Please load PDFs and initialize the system before asking questions.")
    
    preprocessed_query = preprocess_query(user_query)
    
    # Get RAG response using Groq 
    rag_response = rag_chain.invoke({"input": preprocessed_query})["answer"]
    
    # Get Gemini response with full PDF content
    gemini_resp = gemini_response(preprocessed_query, full_pdf_content)
    
    final_response = generate_final_response(rag_response, gemini_resp)
    html_content = markdown_to_html(final_response)
    
    return rag_response, gemini_resp, html_content

# Initialize
full_pdf_content = ""
vector_store = None
rag_chain = None

# Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Data Protection Team")
    gr.Markdown("**Get responses combining advanced RAG, Long Context, and SOTA models to data protection related questions**")
    
    with gr.Row():
        gdpr_checkbox = gr.Checkbox(label="GDPR (EU)")
        ferpa_checkbox = gr.Checkbox(label="FERPA (US)")
        coppa_checkbox = gr.Checkbox(label="COPPA (US <13)")
    css = """
#file_upload {
    max-height: 100px;
    overflow-y: auto;
}
"""
    gr.Markdown("**Optional: upload additional PDFs if needed (national regulation, school policy)**")
    additional_pdfs = gr.File(
        file_count="multiple",
        label="Upload additional PDFs",
        file_types=[".pdf"],
        elem_id="file_upload"
    )
    
    load_button = gr.Button("Load PDFs")
    load_output = gr.Textbox(label="Load Status")
    
    gr.Markdown("**Ask your data protection related question**")
    query_input = gr.Textbox(label="Your Question", placeholder="Ask your question here...")
    query_button = gr.Button("Submit Query")
    
    gr.Markdown("**Results**")
    rag_output = gr.Textbox(label="RAG Pipeline (Llama3.1) Response")
    gemini_output = gr.Textbox(label="Long Context (Gemini 1.5 Pro) Response")
    final_output = gr.HTML(label="Final (GPT-4o) Response")
    
    load_button.click(
        load_pdfs,
        inputs=[
            gr.Checkboxgroup([gdpr_checkbox, ferpa_checkbox, coppa_checkbox]),
            additional_pdfs
        ],
        outputs=load_output
    )
    
    query_button.click(
        process_query,
        inputs=query_input,
        outputs=[rag_output, gemini_output, final_output]
    )

iface.launch()