Update app.py

app.py

@@ -3,35 +3,6 @@ from utils import extract_text_from_pdf, simple_clause_split
 from model import classify_clauses
 from report import generate_pdf
 from salesforce_stub import send_to_salesforce
-from plotly.graph_objects import go
-
-# Heatmap function
-def generate_heatmap(results):
-    risk_levels = {'High': 3, 'Medium': 2, 'Low': 1}
-    x = []
-    y = []
-    z = []
-
-    for idx, r in enumerate(results):
-        x.append(r['clause'])  # Add clause text as x-axis
-        y.append(r['risk_level'])  # Add risk level as y-axis
-        z.append(risk_levels.get(r['risk_level'], 1))  # Map risk level to numeric value
-
-    # Create a heatmap
-    fig = go.Figure(data=go.Heatmap(
-        z=[z],
-        x=x,
-        y=y,
-        colorscale='YlOrRd'
-    ))
-
-    fig.update_layout(
-        title="Contract Risk Heatmap",
-        xaxis_title="Clause",
-        yaxis_title="Risk Level",
-    )
-
-    return fig.to_html(full_html=False)
 
 # Initialize Gradio interface
 def analyze_contract(file):
@@ -43,21 +14,29 @@ def analyze_contract(file):
     overall_score = sum(r['risk_score'] for r in results) / len(results) if results else 0
     report_path = generate_pdf(results, overall_score)
 
-    # High-risk clauses separately
+    # Create clause-level risk highlights with heatmap-style intensity
+    highlight_output = ""
     high_risk_clauses = ""
+    
     for r in results:
+        color = (
+            "red" if r['risk_level'] == "High" else
+            "orange" if r['risk_level'] == "Medium" else
+            "green"
+        )
+
+        # High-risk clauses separately
         if r['risk_level'] == "High":
-            high_risk_clauses += f"<div style='color:red'><b>{r['risk_level']}</b>: {r['clause']}</div><br>"
+            high_risk_clauses += f"<div style='color:{color}'><b>{r['risk_level']}</b>: {r['clause']}</div><br>"
 
-    # Generate heatmap
-    heatmap_html = generate_heatmap(results)
+        highlight_output += f"<div style='color:{color}'><b>{r['risk_level']}</b>: {r['clause']}</div><br>"
 
     send_to_salesforce({
         "clauses": results,
         "overall_score": overall_score
     })
 
-    return f"Overall Risk Score: {overall_score:.2f}", high_risk_clauses, heatmap_html, report_path
+    return f"Overall Risk Score: {overall_score:.2f}", high_risk_clauses, highlight_output, report_path
 
 # Create Gradio Interface
 iface = gr.Interface(
@@ -66,7 +45,7 @@ iface = gr.Interface(
     outputs=[
         gr.Textbox(label="Overall Risk Score"),
         gr.HTML(label="High-Risk Clauses"),  # Display high-risk clauses here
-        gr.HTML(label="Heatmap Visualization"),  # Plotly heatmap here
+        gr.HTML(label="Clause Risk Highlight"),  # Heatmap-style highlighting
         gr.File(label="Download Risk Report (PDF)")
     ],
     title="📜 Contract Risk Heatmap Generator",