Update app.py

app.py

@@ -5,16 +5,25 @@ import numpy as np
 import os
 
 def generate_heatmap(results):
-    # Sample risk levels
-    risk_levels = {"High": 3, "Medium": 2, "Low": 1}
-    
-    clauses = [r['clause'] for r in results]
-    risk_values = [risk_levels.get(r['risk_level'], 1) for r in results]
+    # Check the structure of the results
+    print("Results:", results)
 
+    # If the results are strings (e.g., just the clauses)
+    if isinstance(results, list) and isinstance(results[0], str):
+        clauses = results  # Directly use clauses
+        # For simplicity, assume all clauses are "high risk" here for testing purposes
+        risk_values = [3 for _ in clauses]  # Replace with actual risk assessment logic
+    else:
+        # Assuming results are in the format [{'clause': ..., 'risk_level': ...}, ...]
+        clauses = [r['clause'] for r in results]  # Extract clause text
+        risk_levels = {"High": 3, "Medium": 2, "Low": 1}
+        risk_values = [risk_levels.get(r['risk_level'], 1) for r in results]  # Map risk level to value
+
+    # Plot heatmap
     fig = plt.figure(figsize=(10, 6))
     sns.heatmap([risk_values], annot=True, xticklabels=clauses, yticklabels=["Risk Levels"], cmap="YlOrRd")
 
-    # Save heatmap image to current directory
+    # Save heatmap image
     heatmap_path = os.path.join(os.getcwd(), 'contract_risk_heatmap.png')
     plt.savefig(heatmap_path)