Update legal_document_analysis.py

legal_document_analysis.py

@@ -15,8 +15,6 @@ from fpdf import FPDF
 import getpass
 import pandas as pd
 import seaborn as sns
-import requests
-from bs4 import BeautifulSoup
 
 # Load environment variables from .env file
 load_dotenv()
@@ -372,50 +370,33 @@ 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)
 
-    # 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"
-    ]
+    # Add visualizations for risks
+    pdf.image(base64_to_image(risk_assessment_matrix), x=10, y=pdf.get_y(), w=90)
+    pdf.image(base64_to_image(risk_level_distribution), x=110, y=pdf.get_y()-50, w=90)  # Position next to the first 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(base64_to_image(risks_by_type), x=10, y=pdf.get_y(), w=90)
+    pdf.image(base64_to_image(stacked_bar_chart), x=110, y=pdf.get_y()-50, w=90)  # Position next to the previous image
+    pdf.ln(60)
 
+    pdf.image(base64_to_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
     pdf.ln(10)
 
-    # Save PDF to a temporary file
-    temp_pdf_path = "legal_document_analysis.pdf"
-    pdf.output(temp_pdf_path, 'F')
-
-    # Load the PDF into a BytesIO object
-    with open(temp_pdf_path, "rb") as f:
-        pdf_buffer = io.BytesIO(f.read())
+    # Footer
+    pdf.set_y(-15)
+    pdf.set_font("Arial", 'I', 8)
+    pdf.cell(0, 10, f'Page {pdf.page_no()}', 0, 0, 'C')
 
-    # Optionally, delete the temporary file
-    os.remove(temp_pdf_path)
-
-    return pdf_buffer
+    return pdf
 
 # Function to handle chatbot interaction
-def chatbot_query(user_input, document_text):
-    prompt = f"You are a legal expert. Answer the question based on the following document text:\n\n{document_text}\n\nQuestion: {user_input}"
+def chatbot_query(user_input):
     try:
-        response = model.invoke(prompt)
-        if hasattr(response, 'content'):
-            return response.content.strip()
+        response = model({"text": user_input})
+        if isinstance(response, dict) and 'text' in response:
+            return response['text']
         else:
-            return str(response).strip()
+            return "Error: Unexpected response format."
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -460,6 +441,8 @@ def send_feedback(feedback_content):
 def send_pdf_via_email(pdf_buffer, recipient_email):
     sender_email = os.getenv("SENDER_EMAIL")
     password = os.getenv("EMAIL_PASS")
+
+    msg = MIMEMultipart()
     msg['From'] = sender_email
     msg['To'] = recipient_email
     msg['Subject'] = "Legal Document Analysis PDF"
@@ -467,11 +450,8 @@ 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.seek(0)
-    pdf_attachment.write(pdf_buffer.read())
+    pdf_attachment = io.BytesIO(pdf_buffer.getvalue())
     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)
@@ -548,6 +528,7 @@ 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.")
 
@@ -601,12 +582,43 @@ def display_legal_analysis_page():
             user_input = st.text_input("Ask the chatbot about your document:")
             if st.button("Send"):
                 if user_input:
-                    chatbot_response = chatbot_query(user_input, document_text)
+                    chatbot_response = chatbot_query(user_input)
                     st.write("*Chatbot Response:*")
                     st.write(chatbot_response)
                 else:
                     st.warning("Please enter a question.")
 
+            # Download PDF Analysis Button
+            st.subheader("Download Analysis as PDF")
+            pdf_buffer = io.BytesIO()
+            pdf = 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.output(pdf_buffer, 'F')
+            pdf_buffer.seek(0)
+
+            # Add download button for PDF
+            st.download_button(
+                label="Download PDF Analysis",
+                data=pdf_buffer,
+                file_name="legal_document_analysis.pdf",
+                mime="application/pdf"
+            )
+
+            # Input for recipient email
+            recipient_email = st.text_input("Enter your email address to receive the PDF:")
+
+            # Button to send PDF via email
+            if st.button("Send PDF Analysis"):
+                if recipient_email:
+                    if send_pdf_via_email(pdf_buffer, recipient_email):
+                        st.success("PDF has been sent successfully!")
+                    else:
+                        st.error("Failed to send PDF. Please try again.")
+                else:
+                    st.warning("Please enter a valid email address.")
+
+            # Feedback Form Section
+            display_feedback_form()
+
         with tabs[6]:  # Update Tracker Tab
             st.subheader("Document Updates")
             updates = track_updates(document_text)
@@ -616,9 +628,6 @@ 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.")