Update legal_document_analysis.py

legal_document_analysis.py

@@ -15,7 +15,8 @@ from fpdf import FPDF
 import getpass
 import pandas as pd
 import seaborn as sns
-import tempfile
+import requests
+from bs4 import BeautifulSoup
 
 # Load environment variables from .env file
 load_dotenv()
@@ -169,6 +170,22 @@ def detect_hidden_obligations_or_dependencies(text, summary):
     
     return hidden_dependencies
 
+# Function to answer questions about the document
+def answer_question(question, document_text):
+    prompt = f"The following is a legal document:\n\n{document_text}\n\nBased on this document, answer the following question: {question}"
+    
+    try:
+        response = model.invoke(prompt)
+        if hasattr(response, 'content'):
+            answer = response.content
+        else:
+            answer = str(response)
+        
+        return answer.strip() if answer else "No answer available."
+    except Exception as e:
+        st.error(f"Error answering question: {str(e)}")
+        return None
+
 # Function to detect risks in the text
 def detect_risks(text, summary):
     risk_phrases = [
@@ -334,12 +351,9 @@ def plot_risk_heatmap(detected_risks):
 
     return img_str
 
-# Function to save base64 images to temporary files
-def save_base64_image(base64_data):
-    img_data = base64.b64decode(base64_data)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
-        temp_file.write(img_data)
-        return temp_file.name
+# Function to convert base64 to image
+def base64_to_image(data):
+    return io.BytesIO(base64.b64decode(data))
 
 # Function to generate PDF document with improved aesthetics
 def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix, risk_level_distribution, risks_by_type, stacked_bar_chart, risk_heatmap):
@@ -374,28 +388,38 @@ def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_oblig
         pdf.cell(0, 10, f"{risk['phrase']}: {risk['summary']} (Risk Level: {risk['risk_level']})", ln=True)
     pdf.ln(10)
 
-    # Add visualizations for risks
-    pdf.image(save_base64_image(risk_assessment_matrix), x=10, y=pdf.get_y(), w=90)
-    pdf.image(save_base64_image(risk_level_distribution), x=110, y=pdf.get_y()-50, w=90)  # Position next to the first image
-    pdf.ln(60)
+    # Save images temporarily and add them to the PDF
+    def save_base64_image(image_str, filename):
+        with open(filename, "wb") as img_file:
+            img_file.write(base64.b64decode(image_str))
+
+    # Save images
+    image_filenames = [
+        "risk_assessment_matrix.png",
+        "risk_level_distribution.png",
+        "risks_by_type.png",
+        "stacked_bar_chart.png",
+        "risk_heatmap.png"
+    ]
 
-    pdf.image(save_base64_image(risks_by_type), x=10, y=pdf.get_y(), w=90)
-    pdf.image(save_base64_image(stacked_bar_chart), x=110, y=pdf.get_y()-50, w=90)  # Position next to the previous image
-    pdf.ln(60)
+    images = [risk_assessment_matrix, risk_level_distribution, risks_by_type, stacked_bar_chart, risk_heatmap]
+    
+    for img_str, filename in zip(images, image_filenames):
+        save_base64_image(img_str, filename)
+        pdf.image(filename, x=10, y=pdf.get_y(), w=90)  # Use the saved image file
 
-    pdf.image(save_base64_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
     pdf.ln(10)
 
-    # Footer
-    pdf.set_y(-15)
-    pdf.set_font("Arial", 'I', 8)
-    pdf.cell(0, 10, f'Page {pdf.page_no()}', 0, 0, 'C')
+    # Save PDF to a temporary file
+    temp_pdf_path = "legal_document_analysis.pdf"
+    pdf.output(temp_pdf_path, 'F')
 
-    # Save PDF to a BytesIO object
-    pdf_buffer = io.BytesIO()
-    pdf.output(name=None, dest='S')  # Save to BytesIO with dest='S'
-    pdf_buffer.write(pdf.output(dest='S').encode('latin1'))  # Write output to buffer
-    pdf_buffer.seek(0)
+    # Load the PDF into a BytesIO object
+    with open(temp_pdf_path, "rb") as f:
+        pdf_buffer = io.BytesIO(f.read())
+
+    # Optionally, delete the temporary file
+    os.remove(temp_pdf_path)
 
     return pdf_buffer
 
@@ -460,8 +484,11 @@ def send_pdf_via_email(pdf_buffer, recipient_email):
     msg.attach(MIMEText("Please find the attached analysis of your legal document.", 'plain'))
 
     # Attach the PDF
-    pdf_attachment = io.BytesIO(pdf_buffer.getvalue())
+    pdf_attachment = io.BytesIO()
+    pdf_buffer.seek(0)
+    pdf_attachment.write(pdf_buffer.read())
     pdf_attachment.seek(0)
+
     part = MIMEApplication(pdf_attachment.read(), Name='legal_document_analysis.pdf')
     part['Content-Disposition'] = 'attachment; filename="legal_document_analysis.pdf"'
     msg.attach(part)
@@ -538,7 +565,6 @@ def display_legal_analysis_page():
                     with st.expander(clause['clause'], expanded=False):
                         st.write(f"*Summary:* {clause['summary']}")
                         st.write(f"*Context:* {clause['explanation']}")
-                
             else:
                 st.write("No key clauses detected.")
 
@@ -588,19 +614,21 @@ def display_legal_analysis_page():
             for suggestion in suggestions:
                 st.write(f"- {suggestion}")
 
-            st.subheader("Chatbot for Analysis")
-            user_input = st.text_input("Ask the chatbot about your document:")
-            if st.button("Send"):
-                if user_input:
-                    chatbot_response = chatbot_query(user_input)
-                    st.write("*Chatbot Response:*")
-                    st.write(chatbot_response)
-                else:
-                    st.warning("Please enter a question.")
+            # Chatbot Tab
+            st.subheader("🤖 Chatbot")
+            question = st.text_input("Ask a question about the document:")
+            if question:
+                with st.spinner("Getting answer..."):
+                    answer = answer_question(question, document_text)
+                    if answer:
+                        st.write(f"Answer: {answer}")
+                    else:
+                        st.write("Sorry, I couldn't find an answer to that question.")
 
             # Download PDF Analysis Button
             st.subheader("Download Analysis as PDF")
             pdf_buffer = generate_pdf_analysis(document_text, summary, detected_clauses, hidden_obligations, detected_risks, risk_assessment_matrix, risk_level_distribution, risks_by_type, stacked_bar_chart, risk_heatmap)
+            pdf_buffer.seek(0)
 
             # Add download button for PDF
             st.download_button(
@@ -635,9 +663,13 @@ def display_legal_analysis_page():
                         with st.expander(update['update'], expanded=False):
                             suggestion = get_update_suggestion(update['update'])
                             st.write(f"*Suggestion:* {suggestion}")
+                            # Additional functionality
+                            if st.button(f"Mark '{update['update']}' as addressed"):
+                                st.success(f"'{update['update']}' has been marked as addressed.")
                 else:
                     st.write("No updates detected.")
 
+
 # Run the application
 if __name__ == "__main__":
     display_legal_analysis_page()