Update crs_arena/battle_manager.py

crs_arena/battle_manager.py

@@ -4,7 +4,9 @@ Contains helper functions to select fighters for a battle and generate
 unique user ids.
 """
 
+import itertools
 import logging
+import random
 import uuid
 from collections import defaultdict
 from typing import Optional, Tuple
@@ -36,8 +38,8 @@ CONVERSATION_COUNTS.update(
         "kbrd_opendialkg": 7,
         "unicrs_redial": 4,
         "unicrs_opendialkg": 3,
-        "barcor_redial": 2,
-        "barcor_opendialkg": 2,
+        "barcor_redial": 4,
+        "barcor_opendialkg": 6,
         "crbcrs_redial": 4,
     }
 )
@@ -50,12 +52,37 @@ def get_crs_fighters() -> Tuple[CRSFighter, CRSFighter]:
     The selection is based on the number of conversations collected per model.
     The ones with the least conversations are selected.
 
+    Raises:
+        Exception: If there is an error selecting the fighters.
+        
     Returns:
         CRS models to battle.
     """
-    pair = sorted(CONVERSATION_COUNTS.items(), key=lambda x: x[1])[:2]
-    fighter1 = CRSFighter(1, pair[0][0], CRS_MODELS[pair[0][0]])
-    fighter2 = CRSFighter(2, pair[1][0], CRS_MODELS[pair[1][0]])
+    sorted_count = sorted(CONVERSATION_COUNTS.items(), key=lambda x: x[1])
+    # Group models by conversation count.
+    groups = [
+        list(group)
+        for _, group in itertools.groupby(sorted_count, key=lambda x: x[1])
+    ]
+
+    model_1, model_2 = None, None
+
+    try:
+        if len(groups[0]) >= 2:
+            model_1 = groups[0].pop(random.randint(0, len(groups[0]) - 1))[0]
+            model_2 = groups[0].pop(random.randint(0, len(groups[0]) - 1))[0]
+        else:
+            model_1 = groups[0].pop(random.randint(0, len(groups[0]) - 1))[0]
+            model_2 = groups[1].pop(random.randint(0, len(groups[1]) - 1))[0]
+    except Exception as e:
+        logging.error(f"Error selecting CRS fighters: {e}")
+        if model_1 is None:
+            model_1 = sorted_count[0][0]
+        if model_2 is None:
+            model_2 = sorted_count[1][0]
+
+    fighter1 = CRSFighter(1, model_1, CRS_MODELS[model_1])
+    fighter2 = CRSFighter(2, model_2, CRS_MODELS[model_2])
     return fighter1, fighter2