Update legal_document_analysis.py

legal_document_analysis.py

@@ -15,6 +15,7 @@ from fpdf import FPDF
 import getpass
 import pandas as pd
 import seaborn as sns
+import tempfile
 
 # Load environment variables from .env file
 load_dotenv()
@@ -333,9 +334,12 @@ def plot_risk_heatmap(detected_risks):
 
     return img_str
 
-# Function to convert base64 to image
-def base64_to_image(data):
-    return io.BytesIO(base64.b64decode(data))
+# 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 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):
@@ -371,15 +375,15 @@ def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_oblig
     pdf.ln(10)
 
     # 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.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)
 
-    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.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)
 
-    pdf.image(base64_to_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
+    pdf.image(save_base64_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
     pdf.ln(10)
 
     # Footer
@@ -387,7 +391,12 @@ def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_oblig
     pdf.set_font("Arial", 'I', 8)
     pdf.cell(0, 10, f'Page {pdf.page_no()}', 0, 0, 'C')
 
-    return pdf
+    # Save PDF to a BytesIO object
+    pdf_buffer = io.BytesIO()
+    pdf.output(pdf_buffer, 'F')  # Save to BytesIO
+    pdf_buffer.seek(0)
+
+    return pdf_buffer
 
 # Function to handle chatbot interaction
 def chatbot_query(user_input):
@@ -590,10 +599,7 @@ def display_legal_analysis_page():
 
             # 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)
+            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)
 
             # Add download button for PDF
             st.download_button(