Update app.py

app.py

@@ -4,31 +4,37 @@ import networkx as nx
 import matplotlib.pyplot as plt
 import pdfkit
 import base64
+import shutil
+import logging
 from networkx.drawing.nx_pydot import graphviz_layout
 
-# Function to generate the organization chart and PDF
+# Set up logging to see which layout is used
+logging.basicConfig(level=logging.INFO)
+
 def generate_chart(file, title):
     if file is None:
         return None, "Please upload a CSV file."
-    
-    # Read CSV file
-    df = pd.read_csv(file)
-    
-    # Validate CSV columns
+
+    # Handle the uploaded file: either a dict (Gradio v3.x) or a file path
+    file_path = file["name"] if isinstance(file, dict) and "name" in file else file
+
+    try:
+        df = pd.read_csv(file_path)
+    except Exception as e:
+        return None, f"Error reading CSV: {e}"
+
+    # Validate required columns
     expected_columns = {"Name", "Role", "Reporting To"}
     if not expected_columns.issubset(set(df.columns)):
         return None, "CSV must contain Name, Role, and Reporting To columns."
-    
-    # Build graph from CSV (structure based on CSV data)
+
+    # Build a directed graph from CSV data
     G = nx.DiGraph()
     for _, row in df.iterrows():
-        # Create node label
         node_label = f"{row['Name']}\n({row['Role']})"
         G.add_node(node_label)
-        # Add edge if "Reporting To" is specified
         if pd.notna(row['Reporting To']) and row['Reporting To'] != "":
             rep_name = row['Reporting To']
-            # Try to find the reporting person's role from the CSV
             rep_row = df[df['Name'] == rep_name]
             if not rep_row.empty:
                 rep_role = rep_row.iloc[0]['Role']
@@ -36,47 +42,58 @@ def generate_chart(file, title):
             else:
                 rep_label = rep_name
             G.add_edge(rep_label, node_label)
-    
-    # Create layout using graphviz (via pydot)
-    pos = graphviz_layout(G, prog="dot")
-    
-    # Create the plot
+
+    if len(G.nodes) == 0:
+        return None, "No nodes to display in the chart."
+
+    # Determine layout method: use 'dot' if available, else fall back to spring_layout
+    if shutil.which("dot") is not None:
+        try:
+            logging.info("Using dot for layout.")
+            pos = graphviz_layout(G, prog="dot")
+        except Exception as e:
+            logging.error(f"Error using dot for layout: {e}")
+            logging.info("Falling back to spring_layout.")
+            pos = nx.spring_layout(G)
+    else:
+        logging.info("dot not found. Using spring_layout as fallback.")
+        pos = nx.spring_layout(G)
+
+    # Create the plot with a shadow effect for a 3D look
     plt.figure(figsize=(10, 6))
     ax = plt.gca()
-    
-    # Draw nodes with a shadow for a 3D effect
     for node, (x, y) in pos.items():
-        # Draw shadow (offset rectangle)
-        shadow = plt.Rectangle((x - 40 + 5, y - 20 - 5), 80, 40, 
+        # Draw shadow rectangle for a 3D effect
+        shadow = plt.Rectangle((x - 40 + 5, y - 20 - 5), 80, 40,
                                color='gray', alpha=0.3, zorder=1, ec="none")
         ax.add_patch(shadow)
-        # Draw node rectangle
-        rect = plt.Rectangle((x - 40, y - 20), 80, 40, color='lightblue', 
+        # Draw node rectangle with rounded edges
+        rect = plt.Rectangle((x - 40, y - 20), 80, 40, color='lightblue',
                              ec='black', lw=1.5, zorder=2, joinstyle="round")
         ax.add_patch(rect)
-        # Draw node text
-        plt.text(x, y, node, horizontalalignment='center', verticalalignment='center', 
+        plt.text(x, y, node, horizontalalignment='center', verticalalignment='center',
                  fontsize=10, zorder=3)
-    
-    # Draw edges manually
+
+    # Draw edges between nodes
     for (u, v) in G.edges():
-        x1, y1 = pos[u]
-        x2, y2 = pos[v]
-        plt.plot([x1, x2], [y1, y2], color="gray", zorder=1)
-    
+        if u in pos and v in pos:
+            x1, y1 = pos[u]
+            x2, y2 = pos[v]
+            plt.plot([x1, x2], [y1, y2], color="gray", zorder=1)
+
     plt.axis('off')
-    
-    # Save the chart as an image in /tmp/
+
+    # Save the chart image
     chart_path = "/tmp/chart.png"
     plt.savefig(chart_path, format="png", dpi=300, bbox_inches="tight")
     plt.close()
-    
-    # Convert image to base64 for embedding in HTML
+
+    # Convert the image to base64 for embedding in the PDF
     with open(chart_path, "rb") as img_file:
         encoded_string = base64.b64encode(img_file.read()).decode()
     img_html = f"<img src='data:image/png;base64,{encoded_string}' style='width:100%;'/>"
-    
-    # Build HTML content for PDF (includes user-defined title and the chart image)
+
+    # Build HTML for PDF generation including the user-defined title and the chart image
     html_content = f"""
     <html>
       <head>
@@ -89,14 +106,22 @@ def generate_chart(file, title):
       </body>
     </html>
     """
-    
-    # Generate PDF from HTML content and save to /tmp/
+
     pdf_path = "/tmp/chart.pdf"
-    pdfkit.from_string(html_content, pdf_path)
-    
+
+    # Locate wkhtmltopdf and configure pdfkit
+    wkhtmltopdf_path = shutil.which("wkhtmltopdf")
+    if wkhtmltopdf_path is None:
+        return chart_path, "Error generating PDF: wkhtmltopdf executable not found. Ensure it is installed."
+    config = pdfkit.configuration(wkhtmltopdf=wkhtmltopdf_path)
+
+    try:
+        pdfkit.from_string(html_content, pdf_path, configuration=config)
+    except Exception as e:
+        return chart_path, f"Error generating PDF: {e}"
+
     return chart_path, pdf_path
 
-# Function to provide CSV template
 def download_template():
     template_path = "/tmp/template.csv"
     template_data = "Name,Role,Reporting To\nAlice,CEO,\nBob,Manager,Alice\nCharlie,Engineer,Bob"
@@ -107,18 +132,22 @@ def download_template():
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## Organization Chart Generator")
-    gr.Markdown("Upload a CSV file (with columns: **Name, Role, Reporting To**) to generate an interactive Organization Chart. Download the template to see the required format.")
+    gr.Markdown(
+        "Upload a CSV file (with columns: **Name, Role, Reporting To**) to generate an interactive Organization Chart. "
+        "Use the button below to download a CSV template."
+    )
 
     with gr.Row():
         template_button = gr.Button("Download CSV Template")
         file_input = gr.File(label="Upload CSV File")
     title_input = gr.Textbox(label="Enter PDF Title", placeholder="Company Org Chart")
     submit_button = gr.Button("Generate Chart")
-    
+
     image_output = gr.Image(label="Generated Chart")
-    download_pdf = gr.File(label="Download PDF")
-    
-    template_button.click(download_template, outputs=download_pdf)
-    submit_button.click(generate_chart, inputs=[file_input, title_input], outputs=[image_output, download_pdf])
+    pdf_output = gr.File(label="Download PDF")
+    template_output = gr.File(label="CSV Template")
+
+    template_button.click(download_template, outputs=template_output)
+    submit_button.click(generate_chart, inputs=[file_input, title_input], outputs=[image_output, pdf_output])
 
 demo.launch()