Reducing chance to minimum for AI response to fail

backend/core/utils.py

@@ -67,6 +67,7 @@ def get_controls_for_risk(risk, organization):
 
     organization_details = extract_organization_details(organization)
     valid_control_ids = {control.id for control in all_controls}
+    control_map = {control.id: control.name for control in all_controls}
 
     for control in all_controls:
         control_list.append(f"Control ID: {control.id}, Control Name: {control.name}")
@@ -101,7 +102,7 @@ def get_controls_for_risk(risk, organization):
     ⚠️ **Follow the response format exactly. Any deviation will be considered invalid.**
     """
 
-    for attempt in range(5):
+    for attempt in range(10):
         response = client.chat.completions.create(
             model="gpt-4o-mini",
             messages=[{"role": "system", "content": prompt}]
@@ -141,8 +142,41 @@ def get_controls_for_risk(risk, organization):
         if valid and len(selected_controls) == 10:
             return selected_controls
 
-        print("Invalid response or duplicate control IDs found, retrying...\n")
+        print(f"Recived {len(selected_controls)} controls. Retrying for missing ones...\n")
-        time.sleep(2)
+        remaining_controls = valid_control_ids - control_ids_seen
+        missing_count = 10 - len(selected_controls)
+
+        if missing_count > 0:
+            remaining_controls_list = [f"Control ID:{cid}, Control Name: {control_map[cid]}" for cid in remaining_controls]
+            prompt = f"""
+            You are an expert in cybersecurity risk management. Previously, you selected {len(selected_controls)} controls for the risk "{risk.risk_name}" 
+            and its associated organization details "{organization_details}". 
+
+            Now, your task is to select **exactly {missing_count} additional unique controls** from the remaining list that best mitigate this risk.
+            Each selected 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 with the previously selected controls).
+            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.**
+
+            ### Remaining Available Controls:
+            {remaining_controls_list}
+
+            ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
+            <control_id> : <weight> 
+            <control_id> : <weight> 
+            ### Example Correct Response (NO DUPLICATES):
+            23 : 7,
+            45 : 3,
+            **If you provide duplicate control IDs or controls outside the remaining list, your response will be rejected.**  
+            **Follow the response format exactly. Any deviation will be considered invalid.**
+            """
+
+            time.sleep(2)
+            continue
 
     print("Failed to get a valid response after multiple attempts.")
     return []