fix(LAB-3244): remove order on already sorted chat_items

src/kili/llm/services/export/dynamic.py

@@ -8,7 +8,6 @@
 CHAT_ITEMS_NEEDED_FIELDS = [
     "id",
     "content",
-    "createdAt",
     "modelId",
     "parentId",
     "role",
@@ -123,42 +122,72 @@ def _init_round(self, context):
 
     def _build_rounds(self, chat_items, annotations, json_interface):
         """A round is composed of a prompt with n pre-prompts and n completions."""
-        ordered_chat_items = sorted(chat_items, key=lambda x: x["createdAt"])
+        dict_chat_items = {}
+        for chat_item in chat_items:
+            if dict_chat_items.get(chat_item["parentId"]) is None:
+                dict_chat_items[chat_item["parentId"]] = []
+            dict_chat_items[chat_item["parentId"]].append(chat_item)
         rounds = []
+        parent_target = None
+        has_children = True
         current_round = self._init_round([])
-        for chat_item in ordered_chat_items:
-            role = chat_item["role"].lower() if chat_item["role"] else None
-            if role == "user" or role == "system":
-                if current_round["prompt"] is not None:
-                    rounds.append(current_round)
-                    new_context = (
-                        current_round["context"]
-                        + current_round["pre_prompts"]
-                        + [
-                            current_round["prompt"],
-                            self._get_round_winner(
-                                current_round["completion"],
-                                current_round["annotations"],
-                                json_interface,
-                            ),
-                        ]
-                    )
-                    current_round = self._init_round(new_context)
-
-            if role == "user":
-                current_round["prompt"] = chat_item
-            elif role == "system":
-                current_round["pre_prompts"].append(chat_item)
-            elif role == "assistant":
-                current_round["completion"].append(chat_item)
-            else:
-                raise ValueError(f"Role {chat_item['role']} not supported")
-            current_round["annotations"] += [
-                annotation
-                for annotation in annotations
-                if annotation["chatItemId"] == chat_item["id"]
-            ]
-        rounds.append(current_round)
+
+        while has_children:
+            node = dict_chat_items[parent_target][0]
+            if node["role"].lower() == "system":
+                current_round["pre_prompts"].append(node)
+                parent_target = node["id"]
+                current_round["annotations"] += [
+                    annotation
+                    for annotation in annotations
+                    if annotation["chatItemId"] == node["id"]
+                ]
+                continue
+
+            if node["role"].lower() == "user":
+                current_round["prompt"] = node
+                parent_target = node["id"]
+                current_round["annotations"] += [
+                    annotation
+                    for annotation in annotations
+                    if annotation["chatItemId"] == node["id"]
+                ]
+                continue
+
+            if node["role"].lower() == "assistant":
+                has_children = False
+                if dict_chat_items.get(parent_target) is None:
+                    continue
+                for chat_item in dict_chat_items[parent_target]:
+                    current_round["completion"].append(chat_item)
+                    current_round["annotations"] += [
+                        annotation
+                        for annotation in annotations
+                        if annotation["chatItemId"] == chat_item["id"]
+                    ]
+                    if not has_children and dict_chat_items.get(chat_item["id"]) is not None:
+                        has_children = True
+                        parent_target = chat_item["id"]
+
+                rounds.append(current_round)
+                new_context = (
+                    current_round["context"]
+                    + current_round["pre_prompts"]
+                    + [
+                        current_round["prompt"],
+                        self._get_round_winner(
+                            current_round["completion"],
+                            current_round["annotations"],
+                            json_interface,
+                        ),
+                    ]
+                )
+                current_round = self._init_round(new_context)
+                continue
+
+            raise ValueError(f"Role {node['role']} not supported")
+        if current_round["prompt"] is not None:
+            rounds.append(current_round)
         return rounds
 
 

tests/unit/llm/services/export/test_dynamic.py

@@ -111,12 +111,12 @@
                 },
                 "chatItems": [
                     {
-                        "id": "cm2u6kgcc001aj7ja1stsbrvu",
-                        "content": "You are a helpful assistant",
-                        "createdAt": "2024-08-06T12:28:52.170Z",
-                        "modelId": None,
-                        "parentId": None,
-                        "role": "SYSTEM",
+                        "id": "clziefeoe003m7tc976xwbh58",
+                        "content": "Turtles are reptiles known for their protective shells, which act as both home and armor. They are slow-moving on land but can be agile in water, with many species being excellent swimmers. Turtles are omnivorous, feeding on plants, insects, and small animals. They are cold-blooded and rely on external heat sources to regulate their body temperature. Turtles are long-lived, with some species living over 100 years. They are found on every continent except Antarctica, inhabiting a variety of environments including oceans, freshwater lakes, and even deserts. Many turtle species are endangered due to habitat loss and other factors.",
+                        "createdAt": "2024-08-06T12:30:52.430Z",
+                        "modelId": "clzief6pr003c7tc99680e8yj",
+                        "parentId": "clziefeh6003k7tc99abderkk",
+                        "role": "ASSISTANT",
                     },
                     {
                         "id": "clziefeh6003k7tc99abderkk",
@@ -127,12 +127,12 @@
                         "role": "USER",
                     },
                     {
-                        "id": "clziefeoe003m7tc976xwbh58",
-                        "content": "Turtles are reptiles known for their protective shells, which act as both home and armor. They are slow-moving on land but can be agile in water, with many species being excellent swimmers. Turtles are omnivorous, feeding on plants, insects, and small animals. They are cold-blooded and rely on external heat sources to regulate their body temperature. Turtles are long-lived, with some species living over 100 years. They are found on every continent except Antarctica, inhabiting a variety of environments including oceans, freshwater lakes, and even deserts. Many turtle species are endangered due to habitat loss and other factors.",
-                        "createdAt": "2024-08-06T12:30:52.430Z",
-                        "modelId": "clzief6pr003c7tc99680e8yj",
-                        "parentId": "clziefeh6003k7tc99abderkk",
-                        "role": "ASSISTANT",
+                        "id": "cm2u6kgcc001aj7ja1stsbrvu",
+                        "content": "You are a helpful assistant",
+                        "createdAt": "2024-08-06T12:28:52.170Z",
+                        "modelId": None,
+                        "parentId": None,
+                        "role": "SYSTEM",
                     },
                     {
                         "id": "clziefepk003n7tc9fyx49vei",