From 7151cd0c4df2566bf7ce485c070a96ebad9bf6f1 Mon Sep 17 00:00:00 2001 From: Amir Date: Mon, 17 Feb 2025 21:40:08 +0100 Subject: [PATCH] Changed to gpt-4o-mini, Added organisation info --- backend/accounts/tasks.py | 2 +- backend/core/utils.py | 114 ++++++++++++++++++++++++++------------ 2 files changed, 80 insertions(+), 36 deletions(-) diff --git a/backend/accounts/tasks.py b/backend/accounts/tasks.py index a75c1cb..56cf39a 100644 --- a/backend/accounts/tasks.py +++ b/backend/accounts/tasks.py @@ -34,7 +34,7 @@ def create_document_for_organization(confirmation_email): for risk in top_risks: controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n" - selected_controls = get_controls_for_risk(risk) + selected_controls = get_controls_for_risk(risk ,organization=organization) for control_id, weight in selected_controls: control = Control.objects.filter(id=control_id).first() diff --git a/backend/core/utils.py b/backend/core/utils.py index b261e10..a2a7feb 100644 --- a/backend/core/utils.py +++ b/backend/core/utils.py @@ -1,7 +1,7 @@ from openai import OpenAI from django.conf import settings -from .models import Risk, Control, Document, DocumentRiskControl -from django.shortcuts import get_object_or_404 +from .models import Risk, Control +import time def extract_risk_factors(organization): excluded_fields={"name","email"} @@ -26,6 +26,11 @@ def get_top_risk(organization): Category: {risk.category} Name: {risk.risk_name} Primary Impact: {risk.primary_impact} + Secondary Impact: {risk.secondary_impact} + Tertiary Impact: {risk.tretiary_impact} + Detection Difficulty: {risk.detection_difficulty} + Recovery Complexity: {risk.recovery_complexity} + Business Impact Severity: {risk.businnes_impact_severity} """) risk_factors = extract_risk_factors(organization) @@ -44,59 +49,98 @@ def get_top_risk(organization): """ response = client.chat.completions.create( - model="gpt-4", + model="gpt-4o-mini", messages=[{"role": "system", "content": prompt}] ) risk_ids = response.choices[0].message.content.strip().split(",") + print(f"Risks: {risk_ids}") return [int(risk_id) for risk_id in risk_ids if risk_id.isdigit()] -def get_controls_for_risk(risk): +def get_controls_for_risk(risk, organization): client = OpenAI(api_key=settings.OPENAI_API_KEY) all_controls = Control.objects.all() control_list = [] + risk_factors = extract_risk_factors(organization) + valid_control_ids = {control.id for control in all_controls} + for control in all_controls: control_list.append(f"Control ID: {control.id}, Control Name: {control.name}") prompt = f""" - You are a cyber security expert. For the risk '{risk.risk_name}', select 10 relevant controls - from the following list and assign a weight (1-10) based on how much they reduce risks. - Available Controls (only respond with control IDs and weights): + You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated factors "{risk_factors}", + your task is to select **exactly 10 unique controls** from the provided list that best mitigate this risk. Each control should be assigned a weight between **1 and 10** based on its effectiveness in reducing the risk. + ### Rules: + 1. **Each control ID must be unique** (no duplicates). + 2. **Only return control IDs and weights** in the exact format below. + 3. **Weights must be between 1 and 10** (1 = low impact, 10 = high impact). + 4. **Do NOT add explanations, descriptions, or extra text.** + 5. **Ensure that control IDs are randomly distributed and diverse across different categories.** + ### Available Controls: {control_list} - Respond only with control IDs (numbers) and their corresponding weights (1-10). - Format: - ID: Weight: - Example: - 1: 9 - 2: 6 - 3: 4 + + ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT): + ``` + : + : + + ``` + + ### Example Correct Response (NO DUPLICATES): + ``` + 12 : 8 + 45 : 7 + + ``` + + ⚠️ **If you provide duplicate control IDs, your response will be rejected. Ensure all control IDs are unique.** + ⚠️ **Follow the response format exactly. Any deviation will be considered invalid.** """ - response = client.chat.completions.create( - model="gpt-4", - messages=[{"role": "system", "content": prompt}] - ) + for attempt in range(5): + response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "system", "content": prompt}] + ) + result = response.choices[0].message.content.strip() + print(f"AI Response (Attempt {attempt+1}):\n{result}") - result = response.choices[0].message.content.strip() - selected_controls = [] + selected_controls = [] + valid = True + control_ids_seen = set() - for line in result.split("\n"): - line = line.strip() + for line in result.split("\n"): + line = line.strip() - parts = line.split("Weight:") - if len(parts) == 2: - control_id_str = parts[0].replace("ID:", "").replace("id:", "").replace("Id:", "").strip() - weight_str = parts[1].strip().replace("Weight:", "").replace("weight:","").strip() + parts = line.split(":") + if len(parts) == 2: + control_id_str = parts[0].replace("ID:", "").replace("id:", "").replace("Id:", "").strip() + weight_str = parts[1].strip().replace("Weight:", "").replace("weight:", "").strip() - control_id_str = ''.join(filter(str.isdigit, control_id_str)) - weight_str = ''.join(filter(str.isdigit, weight_str)) - control_id = int(control_id_str) - weight = int(weight_str) - print(f"ID: {control_id}, Weight: {weight}") + control_id_str = ''.join(filter(str.isdigit, control_id_str)) + weight_str = ''.join(filter(str.isdigit, weight_str)) - control = Control.objects.filter(id=control_id).first() - if control: - selected_controls.append((control_id, weight)) - return selected_controls[:10] + if control_id_str and weight_str: + control_id = int(control_id_str) + weight = int(weight_str) + + if control_id in valid_control_ids and 1 <= weight <= 10: + if control_id in control_ids_seen: + valid = False + break + selected_controls.append((control_id, weight)) + control_ids_seen.add(control_id) + else: + valid = False + break + + if valid and len(selected_controls) == 10: + return selected_controls + + print("Invalid response or duplicate control IDs found, retrying...\n") + time.sleep(2) + + print("Failed to get a valid response after multiple attempts.") + return []