diff --git a/backend/core/utils.py b/backend/core/utils.py index 55af846..8f4bf2b 100644 --- a/backend/core/utils.py +++ b/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") - time.sleep(2) + print(f"Recived {len(selected_controls)} controls. Retrying for missing ones...\n") + 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): + : + : + ### 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 []