Update utils/risk_detector.py

utils/risk_detector.py

@@ -5,27 +5,58 @@ from transformers import pipeline
 # ⚖️ Load zero-shot classification model
 classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")
 
-# 🎯 Define risk-related labels
-labels = ["Indemnity", "Exclusivity", "Termination", "Jurisdiction", "Confidentiality"]
+# 🎯 Define risk-related labels (can expand as needed)
+labels = ["Indemnity", "Exclusivity", "Termination", "Jurisdiction", "Confidentiality", "Fees"]
+
+# Optional fallback suggestions
+fallbacks = {
+    "Indemnity": "Consider adding a mutual indemnification clause or capping liability.",
+    "Exclusivity": "Suggest clarifying duration and scope of exclusivity.",
+    "Termination": "Check for balanced termination rights and notice period.",
+    "Jurisdiction": "Ensure forum is neutral or matches your operational base.",
+    "Confidentiality": "Include a clear definition of confidential information and duration.",
+    "Fees": "Ensure clarity on payment structure, late fees, and reimbursement terms."
+}
+
+# ========== Core Function ==========
 
 def detect_risks(text, verbose=False):
     """
-    Classify clauses into predefined legal risk categories.
-    If verbose=True, include detailed scores for each label.
+    Detect and classify legal risks across multiple clauses.
 
     Returns:
-        - List of (label, score) tuples (default)
-        - Or dict of full model output if verbose
+        - List of tuples (clause_text, label, score, fallback) if verbose=True
+        - Otherwise: List of (label, score) tuples aggregated
     """
     if not text.strip():
         return []
 
-    result = classifier(text[:1000], candidate_labels=labels, multi_label=True)
+    # Break into clauses (simple split by period, can be improved)
+    clauses = [c.strip() for c in text.split(".") if len(c.strip()) > 20]
+    all_results = []
+
+    for clause in clauses:
+        result = classifier(clause[:1000], candidate_labels=labels, multi_label=True)
+        top_labels = list(zip(result["labels"], result["scores"]))
+
+        if verbose:
+            top_risks = [(lbl, score) for lbl, score in top_labels if score >= 0.5]
+            for lbl, score in top_risks:
+                all_results.append({
+                    "clause": clause,
+                    "label": lbl,
+                    "score": round(score, 3),
+                    "suggestion": fallbacks.get(lbl, "")
+                })
+        else:
+            all_results.extend(top_labels)
 
     if verbose:
-        return {
-            "sequence": result["sequence"],
-            "predictions": list(zip(result["labels"], result["scores"]))
-        }
-
-    return list(zip(result["labels"], result["scores"]))
+        return all_results
+    else:
+        # Return aggregated top risks (non-verbose mode)
+        from collections import Counter
+        agg = Counter()
+        for lbl, score in all_results:
+            agg[lbl] += score
+        return sorted(agg.items(), key=lambda x: x[1], reverse=True)