feat(LAB-3105): add param for new llm.export path

FannyGaudin · baptiste-olivier · commit d5573e21eaee · 2024-10-18T10:42:46.000+02:00
diff --git a/src/kili/llm/presentation/client/llm.py b/src/kili/llm/presentation/client/llm.py
@@ -44,6 +44,7 @@ def export(
         disable_tqdm: Optional[bool] = False,
         asset_ids: Optional[List[str]] = None,
         external_ids: Optional[List[str]] = None,
+        include_sent_back_labels: Optional[bool] = False,
     ) -> Optional[List[Dict[str, Union[List[str], str]]]]:
         # pylint: disable=line-too-long
         """Returns an export of llm assets with valid labels.
@@ -53,6 +54,7 @@ def export(
             asset_ids: Optional list of the assets internal IDs from which to export the labels.
             disable_tqdm: Disable the progress bar if True.
             external_ids: Optional list of the assets external IDs from which to export the labels.
+            include_sent_back_labels: Include sent back labels if True.
 
         !!! Example
             ```python
@@ -83,6 +85,7 @@ def export(
                 project_id=ProjectId(project_id),
                 asset_filter=asset_filter,
                 disable_tqdm=disable_tqdm,
+                include_sent_back_labels=include_sent_back_labels,
             )
         except NoCompatibleJobError as excp:
             warnings.warn(str(excp), stacklevel=2)
diff --git a/src/kili/llm/services/export/__init__.py b/src/kili/llm/services/export/__init__.py
@@ -2,31 +2,120 @@
 
 from typing import Dict, List, Optional, Union
 
+from kili.adapters.kili_api_gateway.helpers.queries import QueryOptions
 from kili.adapters.kili_api_gateway.kili_api_gateway import KiliAPIGateway
 from kili.domain.asset.asset import AssetFilters
 from kili.domain.project import ProjectId
 
 from .dynamic import LLMDynamicExporter
 from .static import LLMStaticExporter
 
+CHAT_ITEMS_NEEDED_FIELDS = [
+    "id",
+    "content",
+    "createdAt",
+    "modelId",
+    "parentId",
+    "role",
+]
+
+LABELS_NEEDED_FIELDS = [
+    "annotations.id",
+    "author.id",
+    "author.email",
+    "author.firstname",
+    "author.lastname",
+    *(f"chatItems.{field}" for field in CHAT_ITEMS_NEEDED_FIELDS),
+    "createdAt",
+    "id",
+    "isLatestLabelForUser",
+    "isSentBackToQueue",
+    "jsonResponse",  # This is needed to keep annotations
+    "labelType",
+    "modelName",
+]
+
+ASSET_DYNAMIC_NEEDED_FIELDS = [
+    "assetProjectModels.id",
+    "assetProjectModels.configuration",
+    "assetProjectModels.name",
+    "content",
+    "externalId",
+    "jsonMetadata",
+    *(f"labels.{field}" for field in LABELS_NEEDED_FIELDS),
+    "status",
+]
+
+ASSET_STATIC_NEEDED_FIELDS = [
+    "content",
+    "externalId",
+    "jsonMetadata",
+    "labels.jsonResponse",
+    "labels.author.id",
+    "labels.author.email",
+    "labels.author.firstname",
+    "labels.author.lastname",
+    "labels.createdAt",
+    "labels.isLatestLabelForUser",
+    "labels.isSentBackToQueue",
+    "labels.labelType",
+    "labels.modelName",
+    "status",
+]
+
 
 def export(  # pylint: disable=too-many-arguments, too-many-locals
     kili_api_gateway: KiliAPIGateway,
     project_id: ProjectId,
     asset_filter: AssetFilters,
     disable_tqdm: Optional[bool],
+    include_sent_back_labels: Optional[bool],
 ) -> Optional[List[Dict[str, Union[List[str], str]]]]:
     """Export the selected assets with their labels into the required format, and save it into a file archive."""
     project = kili_api_gateway.get_project(project_id, ["id", "inputType", "jsonInterface"])
     input_type = project["inputType"]
 
+    fields = get_fields_to_fetch(input_type)
+    asset_filter.status_in = ["LABELED", "REVIEWED", "TO_REVIEW"]
+    assets = list(
+        kili_api_gateway.list_assets(asset_filter, fields, QueryOptions(disable_tqdm=disable_tqdm))
+    )
+    cleaned_assets = preprocess_assets(assets, include_sent_back_labels or False)
     if input_type == "LLM_RLHF":
-        return LLMStaticExporter(kili_api_gateway, disable_tqdm).export(
-            project_id, asset_filter, project["jsonInterface"]
+        return LLMStaticExporter(kili_api_gateway).export(
+            cleaned_assets, project_id, project["jsonInterface"]
         )
     if input_type == "LLM_INSTR_FOLLOWING":
-        asset_filter.status_in = ["LABELED", "REVIEWED", "TO_REVIEW"]
-        return LLMDynamicExporter(kili_api_gateway, disable_tqdm).export(
-            asset_filter, project["jsonInterface"]
-        )
+        return LLMDynamicExporter(kili_api_gateway).export(cleaned_assets, project["jsonInterface"])
     raise ValueError(f'Project Input type "{input_type}" cannot be used for llm exports.')
+
+
+def get_fields_to_fetch(input_type: str) -> List[str]:
+    """Return the fields to fetch depending on the export type."""
+    if input_type == "LLM_RLHF":
+        return ASSET_STATIC_NEEDED_FIELDS
+    return ASSET_DYNAMIC_NEEDED_FIELDS
+
+
+def preprocess_assets(assets: List[Dict], include_sent_back_labels: bool) -> List[Dict]:
+    """Format labels in the requested format, and filter out autosave labels."""
+    assets_in_format = []
+    for asset in assets:
+        if "labels" in asset:
+            labels_of_asset = []
+            for label in asset["labels"]:
+                labels_of_asset.append(label)
+                if not include_sent_back_labels:
+                    labels_of_asset = list(
+                        filter(lambda label: label["isSentBackToQueue"] is False, labels_of_asset)
+                    )
+            if len(labels_of_asset) > 0:
+                asset["labels"] = labels_of_asset
+                assets_in_format.append(asset)
+        if "latestLabel" in asset:
+            label = asset["latestLabel"]
+            if label is not None:
+                asset["latestLabel"] = label
+            if include_sent_back_labels or asset["latestLabel"]["isSentBackToQueue"] is False:
+                assets_in_format.append(asset)
+    return assets_in_format
diff --git a/src/kili/llm/services/export/dynamic.py b/src/kili/llm/services/export/dynamic.py
@@ -1,9 +1,8 @@
 """Handle LLM_INSTR_FOLLOWING project exports."""
 
 import logging
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Union
 
-from kili.adapters.kili_api_gateway.helpers.queries import QueryOptions
 from kili.adapters.kili_api_gateway.kili_api_gateway import KiliAPIGateway
 from kili.domain.asset.asset import AssetFilters
 
@@ -46,16 +45,13 @@
 class LLMDynamicExporter:
     """Handle exports of LLM_RLHF projects."""
 
-    def __init__(self, kili_api_gateway: KiliAPIGateway, disable_tqdm: Optional[bool]):
+    def __init__(self, kili_api_gateway: KiliAPIGateway):
         self.kili_api_gateway = kili_api_gateway
-        self.disable_tqdm = disable_tqdm
 
     def export(
-        self, asset_filter: AssetFilters, json_interface: Dict
+        self, assets: List[Dict], json_interface: Dict
     ) -> List[Dict[str, Union[List[str], str]]]:
         """Asset content depends of each label."""
-        options = QueryOptions(disable_tqdm=self.disable_tqdm)
-        assets = self.kili_api_gateway.list_assets(asset_filter, ASSET_NEEDED_FIELDS, options)
         export_res = []
         for asset in assets:
             # obfuscate models here
diff --git a/src/kili/llm/services/export/static.py b/src/kili/llm/services/export/static.py
@@ -5,48 +5,24 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Union
 
-from kili.adapters.kili_api_gateway.helpers.queries import QueryOptions
 from kili.adapters.kili_api_gateway.kili_api_gateway import KiliAPIGateway
-from kili.domain.asset.asset import AssetFilters
 from kili.domain.project import ProjectId
 from kili.services.asset_import.helpers import SEPARATOR
 from kili.services.export.format.llm.types import ExportLLMItem
 from kili.use_cases.asset.media_downloader import MediaDownloader
 from kili.utils.tempfile import TemporaryDirectory
 
-ASSET_NEEDED_FIELDS = [
-    "content",
-    "externalId",
-    "jsonMetadata",
-    "labels.jsonResponse",
-    "labels.author.id",
-    "labels.author.email",
-    "labels.author.firstname",
-    "labels.author.lastname",
-    "labels.createdAt",
-    "labels.isLatestLabelForUser",
-    "labels.labelType",
-    "labels.modelName",
-    "status",
-]
-
 
 class LLMStaticExporter:
     """Handle exports of LLM_RLHF projects."""
 
-    def __init__(self, kili_api_gateway: KiliAPIGateway, disable_tqdm: Optional[bool]):
+    def __init__(self, kili_api_gateway: KiliAPIGateway):
         self.kili_api_gateway = kili_api_gateway
-        self.disable_tqdm = disable_tqdm
 
     def export(
-        self, project_id: ProjectId, asset_filter: AssetFilters, json_interface: Dict
+        self, assets: List[Dict], project_id: ProjectId, json_interface: Dict
     ) -> List[Dict[str, Union[List[str], str]]]:
         """Assets are static, with n labels."""
-        assets = list(
-            self.kili_api_gateway.list_assets(
-                asset_filter, ASSET_NEEDED_FIELDS, QueryOptions(disable_tqdm=self.disable_tqdm)
-            )
-        )
         with TemporaryDirectory() as tmpdirname:
             assets = MediaDownloader(
                 tmpdirname,
diff --git a/src/kili/services/export/__init__.py b/src/kili/services/export/__init__.py
@@ -57,7 +57,7 @@ def export_labels(  # pylint: disable=too-many-arguments, too-many-locals
         asset_filter_kwargs=asset_filter_kwargs,
         normalized_coordinates=normalized_coordinates,
         label_type_in=label_type_in,
-        include_sent_back_labels=include_sent_back_labels,
+        include_sent_back_labels=include_sent_back_labels if label_format != "llm_v1" else False,
     )
 
     logger = get_logger(log_level)
diff --git a/tests/unit/llm/services/export/test_dynamic.py b/tests/unit/llm/services/export/test_dynamic.py
@@ -185,9 +185,10 @@
                 ],
                 "createdAt": "2024-08-06T12:30:42.122Z",
                 "isLatestLabelForUser": True,
+                "isSentBackToQueue": False,
                 "id": "clzief6q2003e7tc91jm46uii",
                 "jsonResponse": {},
-                "labelType": "AUTOSAVE",
+                "labelType": "DEFAULT",
                 "modelName": None,
             }
         ],
@@ -261,7 +262,7 @@
                 {
                     "author": "test+admin@kili-technology.com",
                     "created_at": "2024-08-06T12:30:42.122Z",
-                    "label_type": "AUTOSAVE",
+                    "label_type": "DEFAULT",
                     "label": {"COMPARISON_JOB": "A_3", "CLASSIFICATION_JOB": ["BOTH_ARE_GOOD"]},
                 }
             ],
@@ -329,7 +330,7 @@
                 {
                     "author": "test+admin@kili-technology.com",
                     "created_at": "2024-08-06T12:30:42.122Z",
-                    "label_type": "AUTOSAVE",
+                    "label_type": "DEFAULT",
                     "label": {"COMPARISON_JOB": "B_1"},
                 }
             ],
@@ -411,7 +412,7 @@
                 {
                     "author": "test+admin@kili-technology.com",
                     "created_at": "2024-08-06T12:30:42.122Z",
-                    "label_type": "AUTOSAVE",
+                    "label_type": "DEFAULT",
                     "label": {"COMPARISON_JOB": "A_2"},
                 }
             ],
diff --git a/tests/unit/llm/services/export/test_static.py b/tests/unit/llm/services/export/test_static.py
@@ -54,6 +54,7 @@
                 },
                 "createdAt": "2024-08-05T13:03:00.051Z",
                 "isLatestLabelForUser": True,
+                "isSentBackToQueue": False,
                 "labelType": "DEFAULT",
                 "modelName": None,
             }
@@ -77,6 +78,7 @@
                 },
                 "createdAt": "2024-08-05T13:03:03.061Z",
                 "isLatestLabelForUser": True,
+                "isSentBackToQueue": False,
                 "labelType": "DEFAULT",
                 "modelName": None,
             }
@@ -101,6 +103,7 @@
                 },
                 "createdAt": "2024-08-05T13:03:16.028Z",
                 "isLatestLabelForUser": True,
+                "isSentBackToQueue": True,
                 "labelType": "DEFAULT",
                 "modelName": None,
             }
@@ -257,66 +260,6 @@
             }
         ],
     },
-    {
-        "raw_data": [
-            {
-                "role": "user",
-                "content": "BLABLABLA",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-            {
-                "role": "assistant",
-                "content": "response A1",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-            {
-                "role": "assistant",
-                "content": "response B1",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-            {
-                "role": "user",
-                "content": "BLIBLIBLI",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-            {
-                "role": "assistant",
-                "content": "response A2",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-            {
-                "role": "assistant",
-                "content": "response B2",
-                "id": None,
-                "chat_id": None,
-                "model": None,
-            },
-        ],
-        "status": "LABELED",
-        "external_id": "asset#2",
-        "metadata": {},
-        "labels": [
-            {
-                "author": "test+admin@kili-technology.com",
-                "created_at": "2024-08-05T13:03:16.028Z",
-                "label_type": "DEFAULT",
-                "label": {
-                    "CLASSIFICATION_JOB": ["TIE"],
-                    "TRANSCRIPTION_JOB": "There is only some formatting changes\n",
-                },
-            }
-        ],
-    },
 ]
 
 
diff --git a/tests/unit/services/export/test_export.py b/tests/unit/services/export/test_export.py

Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@ def export_labels( # pylint: disable=too-many-arguments, too-many-locals`
`57`	`57`	`asset_filter_kwargs=asset_filter_kwargs,`
`58`	`58`	`normalized_coordinates=normalized_coordinates,`
`59`	`59`	`label_type_in=label_type_in,`
`60`		`- include_sent_back_labels=include_sent_back_labels,`
	`60`	`+ include_sent_back_labels=include_sent_back_labels if label_format != "llm_v1" else False,`
`61`	`61`	`)`
`62`	`62`
`63`	`63`	`logger = get_logger(log_level)`
Original file line number	Diff line number	Diff line change
`@@ -185,9 +185,10 @@`
`185`	`185`	`],`
`186`	`186`	`"createdAt": "2024-08-06T12:30:42.122Z",`
`187`	`187`	`"isLatestLabelForUser": True,`
	`188`	`+ "isSentBackToQueue": False,`
`188`	`189`	`"id": "clzief6q2003e7tc91jm46uii",`
`189`	`190`	`"jsonResponse": {},`
`190`		`- "labelType": "AUTOSAVE",`
	`191`	`+ "labelType": "DEFAULT",`
`191`	`192`	`"modelName": None,`
`192`	`193`	`}`
`193`	`194`	`],`
`@@ -261,7 +262,7 @@`
`261`	`262`	`{`
`262`	`263`	`"author": "test+admin@kili-technology.com",`
`263`	`264`	`"created_at": "2024-08-06T12:30:42.122Z",`
`264`		`- "label_type": "AUTOSAVE",`
	`265`	`+ "label_type": "DEFAULT",`
`265`	`266`	`"label": {"COMPARISON_JOB": "A_3", "CLASSIFICATION_JOB": ["BOTH_ARE_GOOD"]},`
`266`	`267`	`}`
`267`	`268`	`],`
`@@ -329,7 +330,7 @@`
`329`	`330`	`{`
`330`	`331`	`"author": "test+admin@kili-technology.com",`
`331`	`332`	`"created_at": "2024-08-06T12:30:42.122Z",`
`332`		`- "label_type": "AUTOSAVE",`
	`333`	`+ "label_type": "DEFAULT",`
`333`	`334`	`"label": {"COMPARISON_JOB": "B_1"},`
`334`	`335`	`}`
`335`	`336`	`],`
`@@ -411,7 +412,7 @@`
`411`	`412`	`{`
`412`	`413`	`"author": "test+admin@kili-technology.com",`
`413`	`414`	`"created_at": "2024-08-06T12:30:42.122Z",`
`414`		`- "label_type": "AUTOSAVE",`
	`415`	`+ "label_type": "DEFAULT",`
`415`	`416`	`"label": {"COMPARISON_JOB": "A_2"},`
`416`	`417`	`}`
`417`	`418`	`],`