Update legal_document_analysis.py

legal_document_analysis.py

@@ -15,7 +15,6 @@ from fpdf import FPDF
 import getpass
 import pandas as pd
 import seaborn as sns
-import tempfile
 
 # Load environment variables from .env file
 load_dotenv()
@@ -338,13 +337,6 @@ def plot_risk_heatmap(detected_risks):
 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):
     pdf = FPDF()
@@ -379,15 +371,15 @@ def generate_pdf_analysis(document_text, summary, detected_clauses, hidden_oblig
     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.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)
 
-    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.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(save_base64_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
+    pdf.image(base64_to_image(risk_heatmap), x=10, y=pdf.get_y(), w=190)  # Fit image to width
     pdf.ln(10)
 
     # Footer
@@ -395,12 +387,7 @@ 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')
 
-    # 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
+    return pdf
 
 # Function to handle chatbot interaction
 def chatbot_query(user_input):
@@ -603,7 +590,10 @@ def display_legal_analysis_page():
 
             # 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 = 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(