updated get top risks prompt

This commit is contained in:
2025-09-29 14:07:15 +02:00
parent 98f5801bad
commit d2ff9690a1
14 changed files with 202 additions and 32 deletions

View File

@@ -5,7 +5,7 @@ from django.utils.html import format_html
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
from .forms import GenerateCodesForm, RiskExplanationWidget
from django.conf import settings
from backend.accounts.utils import send_document_email
from django import forms
@@ -13,7 +13,7 @@ from django.contrib.admin.widgets import FilteredSelectMultiple
import logging
from django.contrib import messages
from django.db import transaction
import re
logger = logging.getLogger(__name__)
@@ -45,10 +45,16 @@ class DocumentAdminForm(forms.ModelForm):
widget=FilteredSelectMultiple(verbose_name="Risks", is_stacked=False),
help_text="Edit the AI-selected risks for this organization."
)
risk_explanations = forms.Field(
required=False,
widget=RiskExplanationWidget,
help_text="Edit explanations for each risk.",
label=''
)
class Meta:
model = Document
fields = ['organization', 'status', 'key_findings', 'recomendations']
fields = ['organization', 'risk_explanations', 'status', 'key_findings', 'recomendations']
class Media:
css = { 'all': ('admin/css/widgets.css',) }
@@ -62,6 +68,22 @@ class DocumentAdminForm(forms.ModelForm):
if self.instance and getattr(self.instance, 'organization_id', None):
self.fields['organization_risks'].initial = self.instance.organization.risks.all()
def clean_risk_explanations(self):
data = self.data
explanations = {}
for key in data:
if key.startswith('risk_explanations_risk_') and not key.startswith('risk_explanations_risk_new_'):
risk_id = data[key]
explanation_key = f"risk_explanations_explanation_{risk_id}"
explanation = data.get(explanation_key, "")
if risk_id:
explanations[str(risk_id)] = explanation
for i in range(10):
risk_id = data.get(f"risk_explanations_risk_new_{i}", "")
explanation = data.get(f"risk_explanations_explanation_new_{i}", "")
if risk_id:
explanations[str(risk_id)] = explanation
return explanations
class DocumentAdmin(admin.ModelAdmin):
change_form_template = "admin/core/document/change_form.html"
@@ -78,7 +100,7 @@ class DocumentAdmin(admin.ModelAdmin):
fieldsets = (
('Organization & Risks', {
'fields': ('organization', 'regen_note_action', 'regen_document_action', 'organization_risks', 'regen_top_risks_action')
'fields': ('organization', 'regen_note_action', 'regen_document_action', 'organization_risks', 'risk_explanations', 'regen_top_risks_action')
}),
('Key Findings', {
'fields': ('key_findings', 'regen_keyfindings_action')
@@ -138,9 +160,21 @@ class DocumentAdmin(admin.ModelAdmin):
def save_model(self, request, obj, form, change):
super().save_model(request, obj, form, change)
org_risks = form.cleaned_data.get('organization_risks')
explanations = form.cleaned_data.get('risk_explanations', {})
old_explanations = obj.risk_explanations or {}
if org_risks is not None and obj.organization_id:
obj.organization.risks.set(org_risks)
new_explanations = {}
for risk in org_risks:
key = str(risk.risk_id)
new_explanations[key] = explanations.get(key, old_explanations.get(key, ""))
obj.risk_explanations = new_explanations
obj.save(update_fields=['risk_explanations'])
else:
if explanations:
obj.risk_explanations = explanations
obj.save(update_fields=['risk_explanations'])
def _apply_post_org_risks(self, request, obj):
try:
if 'organization_risks' in request.POST and obj.organization_id:
@@ -159,7 +193,7 @@ class DocumentAdmin(admin.ModelAdmin):
except Exception:
logger.exception("Failed to clear segments for document %s", getattr(obj, 'pk', None))
def _risk_content(self, risks):
def _risk_content(self, risks, explanation_map):
return "\n\n".join([
f"Risk: {risk.risk_id} - {risk.risk_name} \n"
f"Category: {risk.category}\n"
@@ -169,6 +203,7 @@ class DocumentAdmin(admin.ModelAdmin):
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 risks
])
@@ -177,7 +212,8 @@ class DocumentAdmin(admin.ModelAdmin):
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)}")
explanation_map = obj.risk_explanations or {}
obj.add_segment('body', f"Identified Risks: \n\n{self._risk_content(risks, explanation_map)}")
def _clear_document_mappings(self, obj, clear_org_risks=True):
try:
@@ -209,9 +245,13 @@ class DocumentAdmin(admin.ModelAdmin):
return ok
def _regen_top_risks(self, obj):
top_risk_ids = get_top_risk(obj.organization)
top_risks_with_explanation = get_top_risk(obj.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)
explanation_map = {r['risk_id']: r['explanation'] for r in top_risks_with_explanation}
obj.organization.risks.set(top_risks)
obj.risk_explanations = explanation_map
obj.save(update_fields=['risk_explanations', 'modified_at'])
self._add_identified_risks(obj, top_risks)
return True