updated get top risks prompt
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user