updated get top risks prompt

backend/accounts/tasks.py

@@ -13,14 +13,17 @@ def create_document_for_organization(confirmation_email):
     is_incomplete = False
 
     organization = get_object_or_404(Organization, email=confirmation_email)
-    top_risk_ids = get_top_risk(organization)
+    top_risks_with_explanation = get_top_risk(organization)
+    top_risk_ids = [r['risk_id'] for r in top_risks_with_explanation]
     if len(top_risk_ids) != 10:
         is_incomplete = True
-        top_risk_ids = get_top_risk(organization)
+        top_risks_with_explanation = get_top_risk(organization)
+        top_risk_ids = [r['risk_id'] for r in top_risks_with_explanation]
     top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)
     organization.risks.set(top_risks)
-
+    explanation_map = {r['risk_id']: r['explanation'] for r in top_risks_with_explanation}
     document = Document.objects.create(organization=organization)
+    document.risk_explanations = explanation_map
     document.add_segment('h1', "Top 10 Risks Identified")
 
     risk_content = "\n\n".join([
@@ -32,6 +35,7 @@ def create_document_for_organization(confirmation_email):
         f"Detection Difficulty: {risk.detection_difficulty} \n"
         f"Recovery Complexity: {risk.recovery_complexity} \n"
         f"Business Impact Severity: {risk.businnes_impact_severity}\n"
+        f"Explanation: {explanation_map.get(risk.risk_id, '')}\n"
         for risk in top_risks
     ])
     document.add_segment('body', f"Identified Risks: \n\n{risk_content}")

backend/accounts/tests/test_tasks.py

@@ -16,7 +16,6 @@ class CeleryTaskTests(TestCase):
             compliance_frameworks=["Ab", "Ba"],
             industry_sector="Technology",
             it_dependency=8,
-            data_sensitivity="High",
             network_infrastructure="Cloud-based",
             remote_workforce_percentage="50%",
             third_party_vendor_access="10-20",
@@ -36,7 +35,11 @@ class CeleryTaskTests(TestCase):
     @patch("backend.accounts.tasks.get_controls_for_risk")
     @patch("backend.accounts.tasks.send_payment_email")
     def test_create_document_for_organization(self, mock_send_payment_email, mock_get_controls_for_risk, mock_get_top_risk):
-        mock_get_top_risk.return_value = [self.risk.risk_id] 
+        mock_get_top_risk.return_value = [{
+            "risk_id": int(self.risk.risk_id),
+            "risk_name": self.risk.risk_name,
+            "explanation": "Test explanation"
+        }]
         mock_get_controls_for_risk.return_value = [(self.control.id, 5, 7)]  
         create_document_for_organization(self.organization.email)
 

backend/accounts/tests/test_utils.py

@@ -21,7 +21,6 @@ class EmailTests(TestCase):
             compliance_frameworks=["Ab", "Ba"],
             industry_sector="Technology",
             it_dependency=8,
-            data_sensitivity="High",
             network_infrastructure="Cloud-based",
             remote_workforce_percentage="50%",
             third_party_vendor_access="10-20",