Zamenio na while petlju dok ne popuni svi 10 kontrola za rizik

backend/core/utils.py

@@ -63,14 +63,17 @@ def get_top_risk(organization):
 def get_controls_for_risk(risk, organization):
     client = OpenAI(api_key=settings.OPENAI_API_KEY)
     all_controls = Control.objects.all()
-    control_list = []
-
     organization_details = extract_organization_details(organization)
+    control_list = [f"Control ID: {control.id}, Control Name: {control.name}" for control in all_controls]
     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:
+    def fetch_controls(prompt):
-        control_list.append(f"Control ID: {control.id}, Control Name: {control.name}")
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "system", "content": prompt}]
+        )
+        return response.choices[0].message.content.strip()
 
     prompt = f"""
     You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated organization details "{organization_details}",
@@ -85,95 +88,97 @@ def get_controls_for_risk(risk, organization):
     {control_list}
 
     ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
-    ```
     <control_id> : <weight>
     <control_id> : <weight>
-
-    ```
-
     ### 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.**
     """
 
-    for attempt in range(10):
+    selected_controls = []
-        response = client.chat.completions.create(
+    control_ids_seen = set()
-            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}")
 
-        selected_controls = []
+    result = fetch_controls(prompt)
-        valid = True
+        
-        control_ids_seen = set()
+    for line in result.split("\n"):
+        line = line.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()
+            print(f"Control:{control_id_str} Weight:{weight_str}")
+            print(f"ControlType: {type(control_id_str)} WeightType: {type(weight_str)}")
 
+            control_id_str = ''.join(filter(str.isdigit, control_id_str))
+            weight_str = ''.join(filter(str.isdigit, weight_str))
+        
+            if control_id_str and weight_str:
+                try:
+                    control_id = int(control_id_str)
+                    weight = int(weight_str)
+                
+                    if control_id in valid_control_ids and 1 <= weight <= 10 and control_id not in control_ids_seen:
+                        selected_controls.append((control_id, weight))
+                        control_ids_seen.add(control_id)
+                except ValueError:
+                    continue
+        
+    if len(selected_controls) == 10:
+        return selected_controls
+    
+    while len(selected_controls) < 10:
+        missing_count = 10 - len(selected_controls)
+        remaining_controls = valid_control_ids - control_ids_seen
+        remaining_controls_list = [f"Control ID: {cid}, Control Name: {control_map[cid]}" for cid in remaining_controls]
+
+        retry_prompt = f"""
+        You are an expert in cybersecurity risk management. Given the risk "{risk.risk_name}" and its associated organization details "{organization_details}",
+        your task is to select **exactly {missing_count} 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:
+        {remaining_controls_list}
+
+        ### Expected Response Format (STRICTLY FOLLOW THIS FORMAT):
+        <control_id> : <weight>
+        <control_id> : <weight>
+        ### 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.**
+        """
+
+        result = fetch_controls(retry_prompt)
         for line in result.split("\n"):
             line = line.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()
-
+                print(f"Control:{control_id} Weight:{weight_str}")
+                print(f"ControlType: {type(control_id)} WeightType: {type(weight_str)}")
+        
                 control_id_str = ''.join(filter(str.isdigit, control_id_str))
                 weight_str = ''.join(filter(str.isdigit, weight_str))
-
+        
                 if control_id_str and weight_str:
-                    control_id = int(control_id_str)
+                    try:
-                    weight = int(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:
+                        if control_id in valid_control_ids and 1 <= weight <= 10 and control_id not in control_ids_seen:
-                            continue
+                            selected_controls.append((control_id, weight))
-                        selected_controls.append((control_id, weight))
+                            control_ids_seen.add(control_id)
-                        control_ids_seen.add(control_id)
+                    except ValueError:
-                    else:
                         continue
 
-        if valid and len(selected_controls) == 10:
+            if not remaining_controls: 
-            return selected_controls
+                break
-
+        return selected_controls if len(selected_controls) == 10 else []
-        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):
-            <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 []