Merge branch '49-august-9-human-in-the-loop' into 'master'

Full Regeneration of document, send email to Risklet teams for incomplete answers from AI

See merge request kbr4/riskletpy!55

backend/accounts/tasks.py

@@ -2,15 +2,21 @@ from celery import shared_task
 from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl
 from backend.core.utils import get_top_risk, get_controls_for_risk, generate_key_findings, generate_recommendations
 from django.shortcuts import get_object_or_404, render
-from .utils import send_payment_email
+from .utils import send_payment_email, send_document_to_reviewer
 from backend.core.tables import get_risk_table
+import logging
+
+logger = logging.getLogger(__name__)
 
 @shared_task
 def create_document_for_organization(confirmation_email):
-    
+    is_incomplete = False
-    organization = get_object_or_404(Organization, email=confirmation_email)
 
+    organization = get_object_or_404(Organization, email=confirmation_email)
     top_risk_ids = get_top_risk(organization)
+    if len(top_risk_ids) != 10:
+        is_incomplete = True
+        top_risk_ids = get_top_risk(organization)
     top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)
     organization.risks.set(top_risks)
 
@@ -36,7 +42,10 @@ def create_document_for_organization(confirmation_email):
         controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
 
         selected_controls = get_controls_for_risk(risk ,organization=organization)
-        
+        if len(selected_controls) < 10:
+            is_incomplete = True
+            selected_controls = get_controls_for_risk(risk, organization=organization)
+
         for control_id, weight, likelihood in selected_controls:
             control = Control.objects.filter(id=control_id).first() 
             if control:
@@ -65,4 +74,9 @@ def create_document_for_organization(confirmation_email):
         document.recomendations = recommendations
         document.save()
 
+    if is_incomplete:
+        logger.info("Marking document %s as INCOMPLETE (is_incomplete=%s)", document.id, is_incomplete)
+        document.status = Document.STATUS_INCOMPLETE
+        document.save(update_fields=['status', 'modified_at'])
+    
     send_payment_email(confirmation_email)

backend/accounts/utils.py

@@ -63,6 +63,27 @@ def send_documet_to_expert(email, document):
         fail_silently=False,
     )
 
+def send_document_to_reviewer(email, document):
+    subject = "Incomplete Document Review Needed"
+    document_link = f"{site_domain}{reverse('core:document', args=[document.id])}"
+    edit_link = f"{site_domain}{reverse('admin:core_document_change', args=[document.id])}"
+
+    message = f"""
+        We had some problems generating document, please review and complete it as needed.
+        When u are done, please save document, and send it to the customer.
+        Document ID: {document.id}
+        Document Link: {document_link}
+        Edit Document: {edit_link}
+    """
+
+    send_mail(
+        subject=subject,
+        message=message,
+        from_email=settings.DEFAULT_FROM_EMAIL,
+        recipient_list=[email],
+        fail_silently=False,
+    )
+
 def send_document_email(email, document_link, document):
     image_io = generate_first_page_image(document)
 

backend/core/admin.py

@@ -2,7 +2,7 @@ from django.contrib import admin
 from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate, DocumentRiskControl, DemoCode
 from django.urls import reverse, path
 from django.utils.html import format_html
-from .utils import generate_demo_code, get_controls_for_risk, generate_key_findings, generate_recommendations
+from .utils import generate_demo_code, get_top_risk, get_controls_for_risk, generate_key_findings, generate_recommendations
 from .tables import get_risk_table
 from django.shortcuts import render, redirect
 from .forms import GenerateCodesForm
@@ -10,6 +10,13 @@ from django.conf import settings
 from backend.accounts.utils import send_document_email
 from django import forms
 from django.contrib.admin.widgets import FilteredSelectMultiple
+import logging
+from django.contrib import messages
+from django.db import transaction
+
+
+logger = logging.getLogger(__name__)
+
 
 class DocumentRiskControlInline(admin.TabularInline):
     model = DocumentRiskControl
@@ -27,7 +34,7 @@ class DocumentRiskControlInline(admin.TabularInline):
             elif obj:
                 formset.form.base_fields['risk'].queryset = obj.organization.risks.all()
         except Exception:
-            pass
+            logger.exception("Error building DocumentRiskControlInline formset")
         return formset
 
 
@@ -65,13 +72,13 @@ class DocumentAdmin(admin.ModelAdmin):
     list_filter = ('status', 'created_at')
     search_fields = ['organization__name', 'organization__email']
     readonly_fields = (
-        'created_at', 'modified_at',
+        'created_at', 'modified_at', 'regen_note_action',
-        'regen_controls_action', 'regen_keyfindings_action', 'regen_recommendations_action',
+        'regen_document_action', 'regen_top_risks_action', 'regen_controls_action', 'regen_keyfindings_action', 'regen_recommendations_action',
     )
 
     fieldsets = (
         ('Organization & Risks', {
-            'fields': ('organization', 'organization_risks', 'regen_controls_action')
+            'fields': ('organization', 'regen_note_action', 'regen_document_action', 'organization_risks', 'regen_top_risks_action')
         }),
         ('Key Findings', {
             'fields': ('key_findings', 'regen_keyfindings_action')
@@ -87,18 +94,47 @@ class DocumentAdmin(admin.ModelAdmin):
         }),
     )
 
+    def regen_note_action(self, obj):
+        return format_html(
+            '<div class="ai-callout">'
+            '⏰ <strong>Tip:</strong> Regenerating document can take some time since we depend on AI models to generate content.'
+            '</div>'
+            '<br/>'
+        )
+    regen_note_action.short_description = ''
+
+    def regen_top_risks_action(self,obj):
+        return format_html(
+            '<div class="ai-callout ai-callout-info">'
+            '💡 <strong>Recommended:</strong> after regenerating the Top 10 risks, also update Controls, Key Findings, and Recommendations.'
+            '</div>'
+            '<br/>'
+            '<button type="submit" name="_regen_top_risks" class="button">Regenerate Top 10 Risks using AI</button>'
+        )
+    regen_top_risks_action.short_description = ''
+
     def regen_controls_action(self, obj):
         return format_html('<button type="submit" name="_regen_controls" class="button">Regenerate Controls using AI</button>')
     regen_controls_action.short_description = ''
 
     def regen_keyfindings_action(self, obj):
-        return format_html('<button type="submit" name="_regen_key_findings" class="button">Regenerate Key Findings using AI</button>')
+        return format_html('<br/><button type="submit" name="_regen_key_findings" class="button">Regenerate Key Findings using AI</button>')
     regen_keyfindings_action.short_description = ''
 
     def regen_recommendations_action(self, obj):
-        return format_html('<button type="submit" name="_regen_recommendations" class="button">Regenerate Recommendations using AI</button>')
+        return format_html('<br/><button type="submit" name="_regen_recommendations" class="button">Regenerate Recommendations using AI</button>')
     regen_recommendations_action.short_description = ''
 
+    def regen_document_action(self, obj):
+        return format_html(
+            '<div class="ai-callout">'
+            '🧹 <strong>Warning:</strong> this will clear the current document (segments, mapped risks/controls, key findings, and recommendations) and regenerate everything as if the document was newly created.'
+            '</div>'
+            '<br/>'
+            '<button type="submit" name="_regen_document" class="button">Regenerate Entire Document</button>'
+        )
+    regen_document_action.short_description = ''
+
     def save_model(self, request, obj, form, change):
         super().save_model(request, obj, form, change)
         org_risks = form.cleaned_data.get('organization_risks')
@@ -111,27 +147,79 @@ class DocumentAdmin(admin.ModelAdmin):
                 risk_ids = [int(pk) for pk in request.POST.getlist('organization_risks') if pk]
                 obj.organization.risks.set(Risk.objects.filter(pk__in=risk_ids))
         except Exception:
-            pass
+            logger.exception("Failed to apply posted organization_risks")
+
+    def _clear_segments(self, obj, startswith=None, exact=None):
+        try:
+            if startswith:
+                for s in startswith:
+                    obj.segments.filter(content__startswith=s).delete()
+            if exact:
+                obj.segments.filter(content__in=exact).delete()
+        except Exception:
+            logger.exception("Failed to clear segments for document %s", getattr(obj, 'pk', None))
+
+    def _risk_content(self, risks):
+        return "\n\n".join([
+            f"Risk: {risk.risk_id} - {risk.risk_name} \n"
+            f"Category: {risk.category}\n"
+            f"Primary Impact: {risk.primary_impact} \n"
+            f"Secondary Impact: {risk.secondary_impact}\n"
+            f"Tertiary Impact: {risk.tretiary_impact} \n"
+            f"Detection Difficulty: {risk.detection_difficulty} \n"
+            f"Recovery Complexity: {risk.recovery_complexity} \n"
+            f"Business Impact Severity: {risk.businnes_impact_severity}\n"
+            for risk in risks
+        ])
+
+    def _add_identified_risks(self, obj, risks):
+        if not risks:
+            return
+        self._clear_segments(obj, startswith=["Identified Risks"], exact=["Top 10 Risks Identified"])
+        obj.add_segment('h1', "Top 10 Risks Identified")
+        obj.add_segment('body', f"Identified Risks: \n\n{self._risk_content(risks)}")
+
+    def _clear_document_mappings(self, obj, clear_org_risks=True):
+        try:
+            obj.segments.all().delete()
+            obj.documentriskcontrol_set.all().delete()
+            obj.key_findings = ''
+            obj.recomendations = ''
+            obj.status = Document.STATUS_WAITING
+            obj.save(update_fields=['key_findings', 'recomendations', 'status', 'modified_at'])
+            if clear_org_risks and getattr(obj, 'organization', None):
+                obj.organization.risks.clear()
+        except Exception:
+            logger.exception("Failed to clear document mappings for document %s", getattr(obj, 'pk', None))
+
+    def _regen_pipeline(self, obj):
+        ok = True
+        if not self._regen_top_risks(obj):
+            ok = False
+        else:
+            try:
+                self._regen_controls(obj)
+            except Exception:
+                logger.exception("_regen_controls failed")
+                ok = False
+            if not self._regen_key_findings(obj):
+                ok = False
+            if not self._regen_recommendations(obj):
+                ok = False
+        return ok
+
+    def _regen_top_risks(self, obj):
+        top_risk_ids = get_top_risk(obj.organization)
+        top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)
+        obj.organization.risks.set(top_risks)
+        self._add_identified_risks(obj, top_risks)
+        return True
 
     def _regen_controls(self, obj):
-        obj.segments.filter(content__startswith="Identified Risks").delete()
+        self._clear_segments(obj, startswith=["Identified Risks", "Mitigation Controls"], exact=["Top 10 Risks Identified", "Regenerated Controls"])
-        obj.segments.filter(content__startswith="Mitigation Controls").delete()
-        obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()
         obj.documentriskcontrol_set.all().delete()
         top_risks = list(obj.organization.risks.all())
-        obj.add_segment('h1', "Top 10 Risks Identified")
+        self._add_identified_risks(obj, top_risks)
-        risk_content = "\n\n".join([
-            f"Risk: {r.risk_id} - {r.risk_name} \n"
-            f"Category: {r.category}\n"
-            f"Primary Impact: {r.primary_impact} \n"
-            f"Secondary Impact: {r.secondary_impact}\n"
-            f"Tertiary Impact: {r.tretiary_impact} \n"
-            f"Detection Difficulty: {r.detection_difficulty} \n"
-            f"Recovery Complexity: {r.recovery_complexity} \n"
-            f"Business Impact Severity: {r.businnes_impact_severity}\n"
-            for r in top_risks
-        ])
-        obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")
         controls_content = "Mitigation Controls:\n\n"
         for risk in top_risks:
             controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
@@ -170,27 +258,55 @@ class DocumentAdmin(admin.ModelAdmin):
         return False
 
     def changeform_view(self, request, object_id=None, form_url='', extra_context=None):
-        if request.method == 'POST' and any(k in request.POST for k in ("_regen_controls", "_regen_key_findings", "_regen_recommendations")):
+        if request.method == 'POST' and any(k in request.POST for k in ("_regen_controls", "_regen_key_findings", "_regen_recommendations", "_regen_top_risks", "_regen_document")):
             obj = self.get_object(request, object_id)
             if obj is None:
                 return super().changeform_view(request, object_id, form_url, extra_context)
             try:
                 self._apply_post_org_risks(request, obj)
-                if "_regen_controls" in request.POST:
+                if "_regen_top_risks" in request.POST:
+                    if not obj.organization_id:
+                        self.message_user(request, "Please select an organization first.", level=messages.WARNING)
+                    else:
+                        if self._regen_top_risks(obj):
+                            self.message_user(request, "Top risks regenerated and risk segment updated.")
+                        else:
+                            self.message_user(request, "Top risks could not be generated.", level=messages.WARNING)
+                elif "_regen_controls" in request.POST:
                     self._regen_controls(obj)
                     self.message_user(request, "Risks and controls regenerated successfully.")
                 elif "_regen_key_findings" in request.POST:
                     if self._regen_key_findings(obj):
                         self.message_user(request, "Key Findings regenerated.")
                     else:
-                        self.message_user(request, "Key Findings could not be generated.", level='warning')
+                        self.message_user(request, "Key Findings could not be generated.", level=messages.WARNING)
                 elif "_regen_recommendations" in request.POST:
                     if self._regen_recommendations(obj):
                         self.message_user(request, "Recommendations regenerated.")
                     else:
-                        self.message_user(request, "Recommendations could not be generated.", level='warning')
+                        self.message_user(request, "Recommendations could not be generated.", level=messages.WARNING)
+                elif "_regen_document" in request.POST:
+                    if not obj.organization_id:
+                        self.message_user(request, "Please select an organization first.", level=messages.WARNING)
+                    else:
+                        try:
+                            with transaction.atomic():
+                                self._clear_document_mappings(obj, clear_org_risks=True)
+                                regen_ok = self._regen_pipeline(obj)
+                                try:
+                                    self.log_change(request, obj, "Full document regeneration triggered")
+                                except Exception:
+                                    logger.exception("Failed to log admin change for full regen")
+                                if regen_ok:
+                                    self.message_user(request, "Document fully regenerated.")
+                                else:
+                                    self.message_user(request, "Document regeneration finished with warnings or missing outputs.", level=messages.WARNING)
+                        except Exception as e:
+                            logger.exception("Full regeneration failed")
+                            self.message_user(request, f"Full regeneration failed: {e}", level=messages.ERROR)
             except Exception as e:
-                self.message_user(request, f"Action failed: {e}", level='error')
+                logger.exception("changeform_view action failed")
+                self.message_user(request, f"Action failed: {e}", level=messages.ERROR)
             return redirect(reverse('admin:core_document_change', args=[obj.pk]))
         return super().changeform_view(request, object_id, form_url, extra_context)
 
@@ -209,32 +325,17 @@ class DocumentAdmin(admin.ModelAdmin):
                 obj.save(update_fields=['status', 'modified_at'])
                 self.message_user(request, "Document sent and marked as done.")
             except Exception as e:
-                self.message_user(request, f"Failed to send document: {e}", level='error')
+                logger.exception("Failed to send document email")
+                self.message_user(request, f"Failed to send document: {e}", level=messages.ERROR)
             return redirect(reverse('admin:core_document_change', args=[obj.pk]))
         return super().response_change(request, obj)
 
     def _refresh_segments_from_current_mappings(self, obj):
-
+        self._clear_segments(obj, startswith=["Identified Risks", "Mitigation Controls"], exact=["Top 10 Risks Identified", "Regenerated Controls"])
-        obj.segments.filter(content__startswith="Identified Risks").delete()
-        obj.segments.filter(content__startswith="Mitigation Controls").delete()
-        obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()
 
         top_risks = list(obj.organization.risks.all())
-
         if top_risks:
-            obj.add_segment('h1', "Top 10 Risks Identified")
+            self._add_identified_risks(obj, top_risks)
-            risk_content = "\n\n".join([
-                f"Risk: {r.risk_id} - {r.risk_name} \n"
-                f"Category: {r.category}\n"
-                f"Primary Impact: {r.primary_impact} \n"
-                f"Secondary Impact: {r.secondary_impact}\n"
-                f"Tertiary Impact: {r.tretiary_impact} \n"
-                f"Detection Difficulty: {r.detection_difficulty} \n"
-                f"Recovery Complexity: {r.recovery_complexity} \n"
-                f"Business Impact Severity: {r.businnes_impact_severity}\n"
-                for r in top_risks
-            ])
-            obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")
 
         from collections import defaultdict
         controls_by_risk = defaultdict(list)
@@ -260,7 +361,7 @@ class DocumentAdmin(admin.ModelAdmin):
         try:
             self._refresh_segments_from_current_mappings(obj)
         except Exception:
-            pass
+            logger.exception("Failed to refresh segments from current mappings for document %s", getattr(obj, 'pk', None))
 
 class DocumentTemplateAdmin(admin.ModelAdmin):
     list_display = ['name', 'created_at', 'updated_at', 'preview_button']

backend/core/migrations/0025_alter_document_status.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.1.3 on 2025-08-26 14:46
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0024_document_status'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='document',
+            name='status',
+            field=models.CharField(choices=[('waiting', 'Waiting'), ('done', 'Done'), ('incomplete', 'Incomplete')], default='waiting', max_length=16),
+        ),
+    ]

backend/core/models.py

@@ -109,9 +109,11 @@ class Document(models.Model):
 
     STATUS_WAITING = 'waiting'
     STATUS_DONE = 'done'
+    STATUS_INCOMPLETE = 'incomplete'
     STATUS_CHOICES = (
         (STATUS_WAITING, 'Waiting'),
         (STATUS_DONE, 'Done'),
+        (STATUS_INCOMPLETE, 'Incomplete'),
     )
     status = models.CharField(max_length=16, choices=STATUS_CHOICES, default=STATUS_WAITING)
 

backend/core/templates/admin/core/document/change_form.html

@@ -22,7 +22,23 @@
     /* Hide labels (and colons) for regen action rows */
     .form-row.field-regen_controls_action label,
     .form-row.field-regen_keyfindings_action label,
+    .form-row.field-regen_document_action label,
+    .form-row.field-regen_note_action label,
     .form-row.field-regen_recommendations_action label { display: none; }
+    .form-row.field-regen_top_risks_action label { display: none; }
+    .form-row.field-regen_controls_action { display: none !important; }
+
+    .ai-callout {
+      margin: 8px 0 0;
+      padding: 10px 12px;
+      border-left: 4px solid #2f80ed;
+      border-radius: 4px;
+      background: #e8f4ff;
+      color: #0c3b66;
+      box-shadow: 0 1px 0 rgba(0,0,0,0.04), 0 0 0 1px rgba(13, 110, 253, 0.06) inset;
+      font-size: 13px;
+      line-height: 1.35;
+    }
   </style>
 {% endblock %}
 
@@ -31,6 +47,19 @@
   <p>Use the Regenerate buttons inside each section to update content. When ready, Save and Send.</p>
   {{ block.super }}
 {% endblock %}
+{% block field_sets %}
+  {% for fieldset in adminform %}
+    {% include "admin/includes/fieldset.html" %}
+    {% if forloop.first %}
+      {% for inline_admin_formset in inline_admin_formsets %}
+        {% include inline_admin_formset.opts.template %}
+      {% endfor %}
+    {% endif %}
+  {% endfor %}
+{% endblock %}
+
+{% block inline_field_sets %}{% endblock %}
+
 
 {% block submit_buttons_bottom %}
   {{ block.super }}
@@ -48,7 +77,36 @@
           }
         });
       }
+
+      function addRegenControlsInsideInline() {
+        var inlineGroup = document.getElementById('documentriskcontrol_set-group');
+        if (!inlineGroup) {
+          document.querySelectorAll('.inline-group').forEach(function(g){
+            if (inlineGroup) return;
+            var h2 = g.querySelector('h2');
+            if (h2 && /document risk controls/i.test(h2.textContent)) {
+              inlineGroup = g;
+            }
+          });
+        }
+        if (!inlineGroup) return;
+
+        if (inlineGroup.querySelector('button[name="_regen_controls"]')) return;
+
+        var h2 = inlineGroup.querySelector('h2') || inlineGroup.firstElementChild;
+        (h2 || inlineGroup).insertAdjacentHTML('afterend',
+          '<div class="inline-action">' +
+            '<div class="ai-callout ai-callout-info">' +
+              '💡 <strong>Recommended:</strong> after regenerating controls, also update Key Findings and Recommendations.' +
+            '</div>' +
+            '<br/>' +
+            '<button type="submit" name="_regen_controls" class="button">Regenerate Controls using AI</button>' +
+            '<br/>'+
+            '</div>'
+        );
+      }
       addSaveSendButton();
+      addRegenControlsInsideInline();
       document.addEventListener('formset:added', addSaveSendButton);
       document.addEventListener('formset:removed', addSaveSendButton);
     })();

backend/core/templates/payment_review.html

@@ -0,0 +1,19 @@
+{% extends "base.html" %}
+{% block content %}
+<section class="py-16 bg-secondary sm:py-24 p-body-full">
+  <div class="max-w-lg w-full mx-auto text-center shadow-lg border border-success rounded-xl p-8 bg-white">
+    <h2 class="text-3xl font-extrabold mb-4 text-success">Thank you.</h2>
+        <p class="mb-8 text-gray-700 text-lg">
+            Your report is awaiting expert review. We’ll deliver the final document to {{ email }} once the review is complete.
+        </p>
+        <a href="{% url 'core:index' %}" class="bg-transparent border-2 border-accent text-accent hover:bg-accent hover:text-primary font-semibold py-3 px-8 rounded-lg text-lg transition-all duration-300 ease-in-out">
+            Go back to Homepage
+        </a>
+  </div>
+</section>
+
+{% endblock content %}
+
+{% block bottom %}
+    <script src="/static/js/formHandling.js"></script>
+{% endblock bottom %}

backend/core/views.py

@@ -7,7 +7,7 @@ import time
 from django.shortcuts import render, redirect , get_object_or_404
 from .forms import OrganizationForm
 from .models import Organization,Document, DocumentTemplate, DemoCode, DocumentRiskControl, Risk, Control
-from backend.accounts.utils import send_confirmation_email, send_document_email, send_documet_to_expert
+from backend.accounts.utils import send_confirmation_email, send_document_email, send_documet_to_expert, send_document_to_reviewer
 from django.contrib.admin.views.decorators import staff_member_required
 from .utils import generate_pdf, generate_risk_graph, generate_residual_risk_graph
 from .tables import risk_matrix_table ,get_risk_table, get_safeguard_summary_table
@@ -135,10 +135,22 @@ def payment_page(request):
                 payment_code.save()
                 document = Document.objects.get(organization = org)
                 url = f"{site_domain}/pdf/{document.id}/"
+                expert_emails_qs = ExpertAnalysisEmails.objects.values_list('email', flat=True).distinct()
+
+                # If document is incomplete, notify reviewers
+                if document.status == Document.STATUS_INCOMPLETE:
+                    if expert_emails_qs:
+                        for email_addr in expert_emails_qs:
+                            try:
+                                send_document_to_reviewer(email_addr, document)
+                            except Exception:
+                                logger.exception("Failed to send incomplete document email to %s", email_addr)
+                    return render(request, 'payment_review.html', {'email': email, 'document': document})
+
+                # If organization requested expert analysis, mark waiting and notify experts
                 if org.expert_analysis:
                     document.status = Document.STATUS_WAITING
                     document.save(update_fields=['status', 'modified_at'])
-                    expert_emails_qs = ExpertAnalysisEmails.objects.values_list('email', flat=True).distinct()
                     expert_emails = [e for e in expert_emails_qs if e]
                     if expert_emails:
                         for email_addr in expert_emails: