feat(LAB-4123): faster exports with the SDK

Author: FannyGaudin

src/kili/adapters/kili_api_gateway/asset/formatters.py

@@ -1,11 +1,17 @@
 """Formatters for assets retrieved from Kili API."""
 
 import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import requests
 
 from kili.adapters.http_client import HttpClient
 from kili.core.helpers import get_response_json, is_url, log_raise_for_status
 from kili.domain.types import ListOrTuple
 
+# Batch size for parallel JSON response downloads (same as export service)
+JSON_RESPONSE_BATCH_SIZE = 10
+
 
 def load_json_from_link(link: str, http_client: HttpClient) -> dict:
     """Load json from link."""
@@ -17,6 +23,72 @@ def load_json_from_link(link: str, http_client: HttpClient) -> dict:
     return get_response_json(response)
 
 
+def download_json_responses_parallel(
+    url_to_label_mapping: list[tuple[str, dict]], http_client: HttpClient
+) -> None:
+    """Download JSON responses in parallel and assign to labels.
+
+    Args:
+        url_to_label_mapping: List of tuples (url, label_dict) to download
+        http_client: HTTP client to use for downloads
+    """
+    if not url_to_label_mapping:
+        return
+
+    # Process in batches to limit concurrent connections
+    for i in range(0, len(url_to_label_mapping), JSON_RESPONSE_BATCH_SIZE):
+        batch = url_to_label_mapping[i : i + JSON_RESPONSE_BATCH_SIZE]
+
+        # Download all URLs in the batch in parallel
+        with ThreadPoolExecutor(max_workers=JSON_RESPONSE_BATCH_SIZE) as executor:
+            # Submit all download tasks
+            future_to_label = {
+                executor.submit(load_json_from_link, url, http_client): label
+                for url, label in batch
+            }
+
+            # Collect results as they complete
+            for future in as_completed(future_to_label):
+                label = future_to_label[future]
+                try:
+                    json_response = future.result()
+                    label["jsonResponse"] = json_response
+                    if "jsonResponseUrl" in label:
+                        del label["jsonResponseUrl"]
+                except (requests.RequestException, json.JSONDecodeError, TimeoutError):
+                    # Set empty dict to ensure consistent response format
+                    label["jsonResponse"] = {}
+                    if "jsonResponseUrl" in label:
+                        del label["jsonResponseUrl"]
+
+
+def _parse_label_json_response(label: dict) -> None:
+    """Parse jsonResponse string to dict for a single label.
+
+    Args:
+        label: Label dict to update in place
+    """
+    json_response_value = label.get("jsonResponse", "{}")
+    try:
+        label["jsonResponse"] = json.loads(json_response_value)
+    except json.JSONDecodeError:
+        label["jsonResponse"] = {}
+
+
+def _process_label_json_response(label: dict, url_to_label_mapping: list[tuple[str, dict]]) -> None:
+    """Process a single label's jsonResponse, either scheduling URL download or parsing.
+
+    Args:
+        label: Label dict to process
+        url_to_label_mapping: List to append URL mapping if download needed
+    """
+    json_response_url = label.get("jsonResponseUrl")
+    if json_response_url and is_url(json_response_url):
+        url_to_label_mapping.append((json_response_url, label))
+    else:
+        _parse_label_json_response(label)
+
+
 def load_asset_json_fields(asset: dict, fields: ListOrTuple[str], http_client: HttpClient) -> dict:
     """Load json fields of an asset."""
     if "jsonMetadata" in fields:
@@ -28,18 +100,17 @@ def load_asset_json_fields(asset: dict, fields: ListOrTuple[str], http_client: H
     if "ocrMetadata" in fields and asset.get("ocrMetadata") is not None:
         asset["ocrMetadata"] = load_json_from_link(asset.get("ocrMetadata", ""), http_client)
 
+    # Collect all URLs to download in parallel (similar to export service)
+    url_to_label_mapping = []
+
     if "labels.jsonResponse" in fields:
-        asset_labels = asset.get("labels", [])
-        for label in asset_labels:
-            try:
-                label["jsonResponse"] = json.loads(label["jsonResponse"])
-            except json.JSONDecodeError:
-                label["jsonResponse"] = {}
+        for label in asset.get("labels", []):
+            _process_label_json_response(label, url_to_label_mapping)
 
     if "latestLabel.jsonResponse" in fields and asset.get("latestLabel") is not None:
-        try:
-            asset["latestLabel"]["jsonResponse"] = json.loads(asset["latestLabel"]["jsonResponse"])
-        except json.JSONDecodeError:
-            asset["latestLabel"]["jsonResponse"] = {}
+        _process_label_json_response(asset["latestLabel"], url_to_label_mapping)
+
+    if url_to_label_mapping:
+        download_json_responses_parallel(url_to_label_mapping, http_client)
 
     return asset

src/kili/adapters/kili_api_gateway/asset/operations_mixin.py

@@ -46,6 +46,9 @@ def list_assets(
         options: QueryOptions,
     ) -> Generator[dict, None, None]:
         """List assets with given options."""
+        has_labels_url = "labels.jsonResponseUrl" in fields
+        has_latest_label_url = "latestLabel.jsonResponseUrl" in fields
+
         if "labels.jsonResponse" in fields or "latestLabel.jsonResponse" in fields:
             # Check if we can get the jsonResponse of if we need to rebuild it.
             project_info = get_project(
@@ -58,7 +61,10 @@ def list_assets(
                 "LLM_STATIC",
                 "GEOSPATIAL",
             }:
-                yield from self.list_assets_split(filters, fields, options, project_info)
+                fetch_annotations = not (has_labels_url or has_latest_label_url)
+                yield from self.list_assets_split(
+                    filters, fields, options, project_info, fetch_annotations
+                )
                 return
 
         fragment = fragment_builder(fields)
@@ -79,7 +85,12 @@ def list_assets(
         yield from assets_gen
 
     def list_assets_split(
-        self, filters: AssetFilters, fields: ListOrTuple[str], options: QueryOptions, project_info
+        self,
+        filters: AssetFilters,
+        fields: ListOrTuple[str],
+        options: QueryOptions,
+        project_info,
+        fetch_annotations: bool,
     ) -> Generator[dict, None, None]:
         """List assets with given options."""
         nb_annotations = self.count_assets_annotations(filters)
@@ -91,22 +102,23 @@ def list_assets_split(
 
         options = QueryOptions(options.disable_tqdm, options.first, options.skip, batch_size)
 
-        inner_annotation_fragment = get_annotation_fragment()
-        annotation_fragment = f"""
-            annotations {{
-                {inner_annotation_fragment}
-            }}
-        """
-        # Ensure 'content', 'resolution', and 'jsonContent' are present in fields
-        required_fields = {"content", "jsonContent", "resolution.width", "resolution.height"}
-        fields = list(fields)
-        for field in required_fields:
-            if field not in fields:
-                fields.append(field)
-
-        fragment = fragment_builder(
-            fields, {"labels": annotation_fragment, "latestLabel": annotation_fragment}
-        )
+        static_fragments = {}
+        if fetch_annotations:
+            inner_annotation_fragment = get_annotation_fragment()
+            annotation_fragment = f"""
+                annotations {{
+                    {inner_annotation_fragment}
+                }}
+            """
+            static_fragments = {"labels": annotation_fragment, "latestLabel": annotation_fragment}
+
+            required_fields = {"content", "jsonContent", "resolution.width", "resolution.height"}
+            fields = list(fields)
+            for field in required_fields:
+                if field not in fields:
+                    fields.append(field)
+
+        fragment = fragment_builder(fields, static_fragments if static_fragments else None)
         query = get_assets_query(fragment)
         where = asset_where_mapper(filters)
         assets_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
@@ -115,25 +127,29 @@ def list_assets_split(
         assets_gen = (
             load_asset_json_fields(asset, fields, self.http_client) for asset in assets_gen
         )
-        converter = AnnotationsToJsonResponseConverter(
-            json_interface=project_info["jsonInterface"],
-            project_input_type=project_info["inputType"],
-        )
-        is_requesting_annotations = any("annotations." in element for element in fields)
-        for asset in assets_gen:
-            if "latestLabel.jsonResponse" in fields and asset.get("latestLabel"):
-                converter.patch_label_json_response(
-                    asset, asset["latestLabel"], asset["latestLabel"]["annotations"]
-                )
-                if not is_requesting_annotations:
-                    asset["latestLabel"].pop("annotations")
 
-            if "labels.jsonResponse" in fields:
-                for label in asset.get("labels", []):
-                    converter.patch_label_json_response(asset, label, label["annotations"])
+        if fetch_annotations:
+            converter = AnnotationsToJsonResponseConverter(
+                json_interface=project_info["jsonInterface"],
+                project_input_type=project_info["inputType"],
+            )
+            is_requesting_annotations = any("annotations." in element for element in fields)
+            for asset in assets_gen:
+                if "latestLabel.jsonResponse" in fields and asset.get("latestLabel"):
+                    converter.patch_label_json_response(
+                        asset, asset["latestLabel"], asset["latestLabel"]["annotations"]
+                    )
                     if not is_requesting_annotations:
-                        label.pop("annotations")
-            yield asset
+                        asset["latestLabel"].pop("annotations")
+
+                if "labels.jsonResponse" in fields:
+                    for label in asset.get("labels", []):
+                        converter.patch_label_json_response(asset, label, label["annotations"])
+                        if not is_requesting_annotations:
+                            label.pop("annotations")
+                yield asset
+        else:
+            yield from assets_gen
 
     def count_assets(self, filters: AssetFilters) -> int:
         """Send a GraphQL request calling countIssues resolver."""

src/kili/adapters/kili_api_gateway/label/formatters.py

@@ -2,15 +2,33 @@
 
 import json
 
+from kili.adapters.http_client import HttpClient
+from kili.core.helpers import is_url
 from kili.domain.types import ListOrTuple
 
 
-def load_label_json_fields(label: dict, fields: ListOrTuple[str]) -> dict:
+def load_json_from_link(link: str, http_client: HttpClient) -> dict:
+    """Load json from link."""
+    if link == "" or not is_url(link):
+        return {}
+
+    response = http_client.get(link, timeout=30)
+    response.raise_for_status()
+    return response.json()
+
+
+def load_label_json_fields(label: dict, fields: ListOrTuple[str], http_client: HttpClient) -> dict:
     """Load json fields of a label."""
     if "jsonResponse" in fields:
-        try:
-            label["jsonResponse"] = json.loads(label.get("jsonResponse", "{}"))
-        except json.JSONDecodeError:
-            label["jsonResponse"] = {}
+        json_response_url = label.get("jsonResponseUrl")
+        if json_response_url and is_url(json_response_url):
+            label["jsonResponse"] = load_json_from_link(json_response_url, http_client)
+            del label["jsonResponseUrl"]
+        else:
+            json_response_value = label.get("jsonResponse", "{}")
+            try:
+                label["jsonResponse"] = json.loads(json_response_value)
+            except json.JSONDecodeError:
+                label["jsonResponse"] = {}
 
     return label

src/kili/adapters/kili_api_gateway/label/operations_mixin.py

@@ -52,6 +52,8 @@ def list_labels(
         options: QueryOptions,
     ) -> Generator[dict, None, None]:
         """List labels."""
+        has_json_response_url = "jsonResponseUrl" in fields
+
         if "jsonResponse" in fields:
             if "labelOf" not in fields:
                 fields = [*list(fields), "assetId"]
@@ -65,7 +67,10 @@ def list_labels(
                 "LLM_INSTR_FOLLOWING",
                 "LLM_STATIC",
             }:
-                yield from self.list_labels_split(filters, fields, options, project_info)
+                fetch_annotations = not has_json_response_url
+                yield from self.list_labels_split(
+                    filters, fields, options, project_info, fetch_annotations
+                )
                 return
 
         fragment = fragment_builder(fields)
@@ -74,11 +79,18 @@ def list_labels(
         labels_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
             query, where, options, "Retrieving labels", GQL_COUNT_LABELS
         )
-        labels_gen = (load_label_json_fields(label, fields) for label in labels_gen)
+        labels_gen = (
+            load_label_json_fields(label, fields, self.http_client) for label in labels_gen
+        )
         yield from labels_gen
 
     def list_labels_split(
-        self, filters: LabelFilters, fields: ListOrTuple[str], options: QueryOptions, project_info
+        self,
+        filters: LabelFilters,
+        fields: ListOrTuple[str],
+        options: QueryOptions,
+        project_info,
+        fetch_annotations: bool,
     ) -> Generator[dict, None, None]:
         """List labels."""
         if project_info["inputType"] == "VIDEO":
@@ -87,21 +99,28 @@ def list_labels_split(
             )
 
         fragment = fragment_builder(fields)
-        inner_annotation_fragment = get_annotation_fragment()
-        full_fragment = f"""
-            {fragment}
-            annotations {{
-                {inner_annotation_fragment}
-            }}
-        """
+
+        if fetch_annotations:
+            inner_annotation_fragment = get_annotation_fragment()
+            full_fragment = f"""
+                {fragment}
+                annotations {{
+                    {inner_annotation_fragment}
+                }}
+            """
+        else:
+            full_fragment = fragment
+
         query = get_labels_query(full_fragment)
         where = label_where_mapper(filters)
         labels_gen = PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
             query, where, options, "Retrieving labels", GQL_COUNT_LABELS
         )
-        labels_gen = (load_label_json_fields(label, fields) for label in labels_gen)
+        labels_gen = (
+            load_label_json_fields(label, fields, self.http_client) for label in labels_gen
+        )
 
-        if "jsonResponse" in fields:
+        if "jsonResponse" in fields and fetch_annotations:
             converter = AnnotationsToJsonResponseConverter(
                 json_interface=project_info["jsonInterface"],
                 project_input_type=project_info["inputType"],

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

@@ -36,6 +36,7 @@
     "isLatestLabelForUser",
     "isSentBackToQueue",
     "jsonResponse",  # This is needed to keep annotations
+    "jsonResponseUrl",
     "labelType",
     "modelName",
 ]
@@ -57,6 +58,7 @@
     "externalId",
     "jsonMetadata",
     "labels.jsonResponse",
+    "labels.jsonResponseUrl",
     "labels.author.id",
     "labels.author.email",
     "labels.author.firstname",

src/kili/services/copy_project/__init__.py

@@ -202,6 +202,7 @@ def _copy_labels_legacy(self, from_project_id: str, new_project_id: str) -> None
             fields=[
                 "author.email",
                 "jsonResponse",
+                "jsonResponseUrl",
                 "secondsToLabel",
                 "isLatestLabelForUser",
                 "labelOf.externalId",