Update app.py

app.py

@@ -14,15 +14,22 @@ 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')
+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)
@@ -40,44 +47,42 @@ def generate_embeddings(docs):
 # 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 a GDPR regulation document.
+    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 document.
-
+    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():
+def create_rag_chain(vector_store):
     prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are an AI assistant helping with GDPR-related queries. Use the following context from the official GDPR regulation document to answer the user's question:\n\n{context}"),
+        ("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):
+def gemini_response(query, full_content):
     prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are an AI assistant helping with GDPR-related queries. Use the following full content of the official GDPR regulation document to answer the user's question:\n\n{context}"),
+        ("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_pdf_content, "input": query}).content
+    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 AI data protection and GDPR compliance.
-    Your goal is to provide simple, practical explanation of and advice on how to meet GDPR requirements based on the given responses.
+    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 form both.
+    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.
     """)
@@ -87,6 +92,36 @@ def generate_final_response(response1, response2):
 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):
     preprocessed_query = preprocess_query(user_query)
     
@@ -94,7 +129,7 @@ def process_query(user_query):
     rag_response = rag_chain.invoke({"input": preprocessed_query})["answer"]
     
     # Get Gemini response with full PDF content
-    gemini_resp = gemini_response(preprocessed_query)
+    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)
@@ -102,25 +137,45 @@ def process_query(user_query):
     return rag_response, gemini_resp, html_content
 
 # Initialize
-GDPR_PDF_PATH = "GDPR.pdf"
-full_pdf_content = extract_pdf(GDPR_PDF_PATH)
-extracted_text = extract_pdf(GDPR_PDF_PATH)
-documents = split_text(extracted_text)
-vector_store = generate_embeddings(documents)
-rag_chain = create_rag_chain()
+full_pdf_content = ""
+vector_store = None
+rag_chain = None
 
 # Gradio interface
-iface = gr.Interface(
-    fn=process_query,
-    inputs=gr.Textbox(label="Ask your data protection related question"),
-    outputs=[
-        gr.Textbox(label="RAG Pipeline (Llama3.1) Response"),
-        gr.Textbox(label="Long Context (Gemini 1.5 Pro) Response"),
-        gr.HTML(label="Final (GPT-4o) Response")
-    ],
-    title="Data Protection Team",
-    description="Get responses combining advanced RAG, Long Context, and SOTA models to data protection related questions .",
-    allow_flagging="never" 
-)
+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")
+        ferpa_checkbox = gr.Checkbox(label="FERPA")
+        coppa_checkbox = gr.Checkbox(label="COPPA")
+    
+    additional_pdfs = gr.File(file_count="multiple", label="Upload additional regulations (PDF)")
+    
+    load_button = gr.Button("Load PDFs")
+    load_output = gr.Textbox(label="Load Status")
+    
+    query_input = gr.Textbox(label="Ask your data protection related question")
+    query_button = gr.Button("Submit Query")
+    
+    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(debug=True)