fix(LAB-4269): LLM project need annotations as the jsonResponseUrl is not compute for those project

FannyGaudin · FannyGaudin · commit ef614db7ce40 · 2026-03-09T14:55:15.000+01:00
diff --git a/src/kili/adapters/kili_api_gateway/asset/operations.py b/src/kili/adapters/kili_api_gateway/asset/operations.py
@@ -30,3 +30,9 @@ def get_assets_query(fragment: str) -> str:
   external_ids: filterExistingAssets(projectID: $projectID, externalIDs: $externalIDs)
 }
 """
+
+GQL_COUNT_ASSET_ANNOTATIONS = """
+query countAssetAnnotations($where: AssetWhere!) {
+    data: countAssetAnnotations(where: $where)
+}
+"""
diff --git a/src/kili/adapters/kili_api_gateway/asset/operations_mixin.py b/src/kili/adapters/kili_api_gateway/asset/operations_mixin.py
@@ -2,11 +2,16 @@
 
 from collections.abc import Generator
 
+from kili_formats.tool.annotations_to_json_response import (
+    AnnotationsToJsonResponseConverter,
+)
+
 from kili.adapters.kili_api_gateway.asset.formatters import (
     load_asset_json_fields,
 )
 from kili.adapters.kili_api_gateway.asset.mappers import asset_where_mapper
 from kili.adapters.kili_api_gateway.asset.operations import (
+    GQL_COUNT_ASSET_ANNOTATIONS,
     GQL_COUNT_ASSETS,
     GQL_CREATE_UPLOAD_BUCKET_SIGNED_URLS,
     GQL_FILTER_EXISTING_ASSETS,
@@ -18,11 +23,18 @@
     QueryOptions,
     fragment_builder,
 )
+from kili.adapters.kili_api_gateway.label.common import get_annotation_fragment
 from kili.adapters.kili_api_gateway.project.common import get_project
 from kili.core.graphql.operations.asset.mutations import GQL_SET_ASSET_CONSENSUS
 from kili.domain.asset import AssetFilters
 from kili.domain.types import ListOrTuple
 
+# Threshold for batching based on number of annotations
+# This is used to determine whether to use a single batch or multiple batches
+# when fetching assets. If the number of annotations counted exceeds this threshold,
+# the asset fetch will be done in multiple smaller batches to avoid performance issues.
+THRESHOLD_FOR_BATCHING = 200
+
 
 class AssetOperationMixin(BaseOperationMixin):
     """Mixin extending Kili API Gateway class with Assets related operations."""
@@ -66,30 +78,73 @@ def list_assets(
 
         yield from assets_gen
 
-    def list_assets_split(
+    def list_assets_split(  # pylint: disable=too-many-branches
         self,
         filters: AssetFilters,
         fields: ListOrTuple[str],
         options: QueryOptions,
         project_info,
     ) -> Generator[dict, None, None]:
         """List assets with given options."""
+        # For LLM projects, we need to fetch annotations and rebuild jsonResponse
+        # because LLM projects don't have jsonResponseUrl
+        is_llm_project = project_info["inputType"] in {
+            "LLM_RLHF",
+            "LLM_INSTR_FOLLOWING",
+            "LLM_STATIC",
+        }
+
         assets_batch_max_amount = 10 if project_info["inputType"] == "VIDEO" else 50
         batch_size_to_use = min(options.batch_size, assets_batch_max_amount)
 
-        options = QueryOptions(options.disable_tqdm, options.first, options.skip, batch_size_to_use)
+        # For LLM projects fetching annotations, adjust batch size based on annotation count
+        if is_llm_project and (
+            "labels.jsonResponse" in fields or "latestLabel.jsonResponse" in fields
+        ):
+            nb_annotations = self.count_assets_annotations(filters)
+            batch_size = (
+                1
+                if nb_annotations / batch_size_to_use > THRESHOLD_FOR_BATCHING
+                else batch_size_to_use
+            )
+        else:
+            batch_size = batch_size_to_use
+
+        options = QueryOptions(options.disable_tqdm, options.first, options.skip, batch_size)
+
+        requested_labels_json_response = "labels.jsonResponse" in fields
+        requested_latest_label_json_response = "latestLabel.jsonResponse" in fields
+        needs_json_response = requested_labels_json_response or requested_latest_label_json_response
 
         required_fields = {"content", "jsonContent", "resolution.width", "resolution.height"}
-        if "labels.jsonResponse" in fields:
-            required_fields.add("labels.jsonResponseUrl")
-        if "latestLabel.jsonResponse" in fields:
-            required_fields.add("latestLabel.jsonResponseUrl")
         fields = list(fields)
+
+        static_fragments = {}
+        if is_llm_project and needs_json_response:
+            # For LLM projects: fetch annotations and rebuild jsonResponse client-side
+            inner_annotation_fragment = get_annotation_fragment()
+            annotation_fragment = f"""
+                annotations {{
+                    {inner_annotation_fragment}
+                }}
+            """
+            static_fragments = {"labels": annotation_fragment, "latestLabel": annotation_fragment}
+
+            fields = list(fields)
+            for field in required_fields:
+                if field not in fields:
+                    fields.append(field)
+        else:
+            if requested_labels_json_response:
+                required_fields.add("labels.jsonResponseUrl")
+            if requested_latest_label_json_response:
+                required_fields.add("latestLabel.jsonResponseUrl")
+
         for field in required_fields:
             if field not in fields:
                 fields.append(field)
 
-        fragment = fragment_builder(fields)
+        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(
@@ -99,7 +154,26 @@ def list_assets_split(
             load_asset_json_fields(asset, fields, self.http_client) for asset in assets_gen
         )
 
-        yield from assets_gen
+        if is_llm_project and needs_json_response:
+            # Rebuild jsonResponse from annotations for LLM projects
+            converter = AnnotationsToJsonResponseConverter(
+                json_interface=project_info["jsonInterface"],
+                project_input_type=project_info["inputType"],
+            )
+            for asset in assets_gen:
+                if requested_latest_label_json_response and asset.get("latestLabel"):
+                    converter.patch_label_json_response(
+                        asset, asset["latestLabel"], asset["latestLabel"]["annotations"]
+                    )
+                    asset["latestLabel"].pop("annotations", None)
+
+                if requested_labels_json_response:
+                    for label in asset.get("labels", []):
+                        converter.patch_label_json_response(asset, label, label["annotations"])
+                        label.pop("annotations", None)
+                yield asset
+        else:
+            yield from assets_gen
 
     def count_assets(self, filters: AssetFilters) -> int:
         """Send a GraphQL request calling countIssues resolver."""
@@ -109,6 +183,14 @@ def count_assets(self, filters: AssetFilters) -> int:
         count: int = count_result["data"]
         return count
 
+    def count_assets_annotations(self, filters: AssetFilters) -> int:
+        """Count the number of annotations for assets matching the filters."""
+        where = asset_where_mapper(filters)
+        payload = {"where": where}
+        count_result = self.graphql_client.execute(GQL_COUNT_ASSET_ANNOTATIONS, payload)
+        count: int = count_result["data"]
+        return count
+
     def create_upload_bucket_signed_urls(self, file_paths: list[str]) -> list[str]:
         """Send a GraphQL request calling createUploadBucketSignedUrls resolver."""
         payload = {
diff --git a/src/kili/adapters/kili_api_gateway/label/common.py b/src/kili/adapters/kili_api_gateway/label/common.py
@@ -10,6 +10,42 @@
 from kili.exceptions import NotFound
 
 
+def get_annotation_fragment() -> str:
+    """Generate a basic annotation fragment for querying annotations.
+
+    This is used for LLM projects.
+    """
+    return """
+        __typename
+        id
+        job
+        path
+        labelId
+        ... on ClassificationAnnotation {
+            annotationValue {
+                categories
+            }
+            chatItemId
+        }
+        ... on ComparisonAnnotation {
+            annotationValue {
+                choice {
+                    code
+                    firstId
+                    secondId
+                }
+            }
+            chatItemId
+        }
+        ... on TranscriptionAnnotation {
+            annotationValue {
+                text
+            }
+            chatItemId
+        }
+    """
+
+
 def get_asset(
     graphql_client: GraphQLClient,
     http_client: HttpClient,
diff --git a/src/kili/adapters/kili_api_gateway/label/operations_mixin.py b/src/kili/adapters/kili_api_gateway/label/operations_mixin.py
@@ -4,6 +4,8 @@
 from collections.abc import Generator
 from typing import Optional
 
+from kili_formats.tool.annotations_to_json_response import AnnotationsToJsonResponseConverter
+
 from kili.adapters.kili_api_gateway.base import BaseOperationMixin
 from kili.adapters.kili_api_gateway.helpers.queries import (
     PaginatedGraphQLQuery,
@@ -19,6 +21,7 @@
 from kili.domain.types import ListOrTuple
 from kili.utils.tqdm import tqdm
 
+from .common import get_annotation_fragment
 from .formatters import load_label_json_fields
 from .mappers import append_label_data_mapper, append_to_labels_data_mapper, label_where_mapper
 from .operations import (
@@ -87,20 +90,53 @@ def list_labels_split(
                 options.disable_tqdm, options.first, options.skip, min(options.batch_size, 20)
             )
 
+        # For LLM projects, we need to fetch annotations and rebuild jsonResponse
+        # because LLM projects don't have jsonResponseUrl
+        is_llm_project = project_info["inputType"] in {
+            "LLM_RLHF",
+            "LLM_INSTR_FOLLOWING",
+            "LLM_STATIC",
+        }
+        needs_json_response = "jsonResponse" in fields
+
         fields = list(fields)
-        if "jsonResponse" in fields and "jsonResponseUrl" not in fields:
-            fields.append("jsonResponseUrl")
 
-        fragment = fragment_builder(fields)
-        query = get_labels_query(fragment)
+        if is_llm_project and needs_json_response:
+            # For LLM projects: fetch annotations and rebuild jsonResponse client-side
+            inner_annotation_fragment = get_annotation_fragment()
+            full_fragment = f"""
+                {fragment_builder([f for f in fields if f not in {"jsonResponse", "jsonResponseUrl"}])}
+                annotations {{
+                    {inner_annotation_fragment}
+                }}
+            """
+        else:
+            if "jsonResponse" in fields and "jsonResponseUrl" not in fields:
+                fields.append("jsonResponseUrl")
+            full_fragment = fragment_builder(fields)
+
+        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, self.http_client) for label in labels_gen
         )
-        yield from labels_gen
+
+        if is_llm_project and needs_json_response:
+            # Rebuild jsonResponse from annotations for LLM projects
+            converter = AnnotationsToJsonResponseConverter(
+                json_interface=project_info["jsonInterface"],
+                project_input_type=project_info["inputType"],
+            )
+            for label in labels_gen:
+                asset = None
+                converter.patch_label_json_response(asset, label, label["annotations"])
+                label.pop("annotations", None)
+                yield label
+        else:
+            yield from labels_gen
 
     def delete_labels(
         self, ids: ListOrTuple[LabelId], disable_tqdm: Optional[bool]
