Changed to gpt-4o-mini, Added organisation info

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() 

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 .models import Risk, Control
-from django.shortcuts import get_object_or_404
+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
+    You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated factors "{risk_factors}",
-    from the following list and assign a weight (1-10) based on how much they reduce risks.
+    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.
-    Available Controls (only respond with control IDs and weights):
+    ### 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:
+    ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
-    ID: <control_id> Weight: <weight>
+    ```
-    Example:
+    <control_id> : <weight>
-    1: 9
+    <control_id> : <weight>
-    2: 6
+
-    3: 4
+    ```
+
+    ### 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(
+    for attempt in range(5):
-        model="gpt-4",
+        response = client.chat.completions.create(
-        messages=[{"role": "system", "content": prompt}]
+            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"):
+        for line in result.split("\n"):
-        line = line.strip()
+            line = line.strip()
 
-        parts = line.split("Weight:")
+            parts = line.split(":")
-        if len(parts) == 2:
+            if len(parts) == 2:
-            control_id_str = parts[0].replace("ID:", "").replace("id:", "").replace("Id:", "").strip() 
+                control_id_str = parts[0].replace("ID:", "").replace("id:", "").replace("Id:", "").strip()
-            weight_str = parts[1].strip().replace("Weight:", "").replace("weight:","").strip()
+                weight_str = parts[1].strip().replace("Weight:", "").replace("weight:", "").strip()
 
-            control_id_str = ''.join(filter(str.isdigit, control_id_str))
+                control_id_str = ''.join(filter(str.isdigit, control_id_str))
-            weight_str = ''.join(filter(str.isdigit, weight_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 = Control.objects.filter(id=control_id).first()
+                if control_id_str and weight_str:
-            if control:
+                    control_id = int(control_id_str)
-                selected_controls.append((control_id, weight))
+                    weight = int(weight_str)
-    return selected_controls[:10]  
+
+                    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 []