fix(LAB-4269): isolated LLM annotations specific code

Author: FannyGaudin

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

@@ -31,6 +31,8 @@ def get_assets_query(fragment: str) -> str:
 }
 """
 
+# TODO(LAB-4269): TEMPORARY WORKAROUND - Remove when backend handles jsonResponseUrl for LLM
+# This GraphQL query is ONLY used for LLM projects to count annotations for batch size optimization
 GQL_COUNT_ASSET_ANNOTATIONS = """
 query countAssetAnnotations($where: AssetWhere!) {
     data: countAssetAnnotations(where: $where)

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

@@ -29,16 +29,127 @@
 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.
+# TODO(LAB-4269): TEMPORARY WORKAROUND - Remove when backend handles jsonResponseUrl for LLM
+# Threshold for batching based on number of annotations.
+# This is used ONLY for LLM projects that need to fetch annotations instead of jsonResponseUrl.
+# 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."""
 
+    # ============================================================================
+    # TODO(LAB-4269): TEMPORARY WORKAROUND - Remove when backend handles jsonResponseUrl for LLM
+    # The following helper methods handle the case where LLM projects don't have jsonResponseUrl.
+    # Once the backend properly computes jsonResponseUrl for LLM projects, all these methods
+    # can be deleted as a unit.
+    # ============================================================================
+
+    @staticmethod
+    def _is_llm_project_without_json_response_url(project_input_type: str) -> bool:
+        """Check if project is LLM type that doesn't have jsonResponseUrl.
+
+        TODO(LAB-4269): Remove this method when backend handles jsonResponseUrl for LLM projects.
+
+        Args:
+            project_input_type: The input type of the project
+
+        Returns:
+            True if project is LLM type without jsonResponseUrl support
+        """
+        return project_input_type in {
+            "LLM_RLHF",
+            "LLM_INSTR_FOLLOWING",
+            "LLM_STATIC",
+        }
+
+    def _calculate_batch_size_for_llm_annotations(
+        self,
+        filters,
+        batch_size_to_use: int,
+    ) -> int:
+        """Calculate optimal batch size for LLM projects based on annotation count.
+
+        TODO(LAB-4269): Remove this method when backend handles jsonResponseUrl for LLM projects.
+
+        For LLM projects, we fetch annotations instead of jsonResponseUrl. If there are too many
+        annotations per asset, we reduce the batch size to avoid performance issues.
+
+        Args:
+            filters: Asset filters to count annotations for
+            batch_size_to_use: The default batch size to use
+
+        Returns:
+            Adjusted batch size (1 if annotations/batch > threshold, otherwise batch_size_to_use)
+        """
+        nb_annotations = self.count_assets_annotations(filters)
+        return (
+            1 if nb_annotations / batch_size_to_use > THRESHOLD_FOR_BATCHING else batch_size_to_use
+        )
+
+    @staticmethod
+    def _build_llm_annotation_fragment() -> dict[str, str]:
+        """Build GraphQL fragment for fetching annotations in LLM projects.
+
+        TODO(LAB-4269): Remove this method when backend handles jsonResponseUrl for LLM projects.
+
+        Returns:
+            Dictionary with static fragments for labels and latestLabel containing annotations
+        """
+        inner_annotation_fragment = get_annotation_fragment()
+        annotation_fragment = f"""
+            annotations {{
+                {inner_annotation_fragment}
+            }}
+        """
+        return {"labels": annotation_fragment, "latestLabel": annotation_fragment}
+
+    def _convert_llm_annotations_to_json_response(
+        self,
+        assets_gen: Generator[dict, None, None],
+        project_info: dict,
+        requested_labels_json_response: bool,
+        requested_latest_label_json_response: bool,
+    ) -> Generator[dict, None, None]:
+        """Convert annotations to jsonResponse for LLM projects.
+
+        TODO(LAB-4269): Remove this method when backend handles jsonResponseUrl for LLM projects.
+
+        For LLM projects, we need to rebuild jsonResponse from annotations client-side
+        because the backend doesn't compute jsonResponseUrl for these project types.
+
+        Args:
+            assets_gen: Generator of assets with annotations
+            project_info: Project information including jsonInterface and inputType
+            requested_labels_json_response: Whether labels.jsonResponse was requested
+            requested_latest_label_json_response: Whether latestLabel.jsonResponse was requested
+
+        Yields:
+            Assets with jsonResponse rebuilt from annotations
+        """
+        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
+
+    # ============================================================================
+    # END TODO(LAB-4269) - LLM workaround methods
+    # ============================================================================
+
     def list_assets(
         self,
         filters: AssetFilters,
@@ -86,54 +197,31 @@ def list_assets_split(  # pylint: disable=too-many-branches
         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",
-        }
+        # TODO(LAB-4269): LLM workaround - detect if we need special handling
+        is_llm_project = self._is_llm_project_without_json_response_url(project_info["inputType"])
 
         assets_batch_max_amount = 10 if project_info["inputType"] == "VIDEO" else 50
         batch_size_to_use = min(options.batch_size, assets_batch_max_amount)
 
-        # 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
-            )
+        # TODO(LAB-4269): LLM workaround - adjust batch size for annotation fetching
+        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
+
+        if is_llm_project and needs_json_response:
+            batch_size = self._calculate_batch_size_for_llm_annotations(filters, 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"}
         fields = list(fields)
 
+        # TODO(LAB-4269): LLM workaround - build annotation fragment instead of using jsonResponseUrl
         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)
+            static_fragments = self._build_llm_annotation_fragment()
         else:
             if requested_labels_json_response:
                 required_fields.add("labels.jsonResponseUrl")
@@ -154,24 +242,14 @@ def list_assets_split(  # pylint: disable=too-many-branches
             load_asset_json_fields(asset, fields, self.http_client) for asset in assets_gen
         )
 
+        # TODO(LAB-4269): LLM workaround - convert annotations to jsonResponse client-side
         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"],
+            yield from self._convert_llm_annotations_to_json_response(
+                assets_gen,
+                project_info,
+                requested_labels_json_response,
+                requested_latest_label_json_response,
             )
-            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
 
@@ -184,7 +262,11 @@ def count_assets(self, filters: AssetFilters) -> int:
         return count
 
     def count_assets_annotations(self, filters: AssetFilters) -> int:
-        """Count the number of annotations for assets matching the filters."""
+        """Count the number of annotations for assets matching the filters.
+
+        TODO(LAB-4269): Remove this method when backend handles jsonResponseUrl for LLM projects.
+        This method is ONLY used to calculate optimal batch sizes for LLM projects.
+        """
         where = asset_where_mapper(filters)
         payload = {"where": where}
         count_result = self.graphql_client.execute(GQL_COUNT_ASSET_ANNOTATIONS, payload)

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

@@ -9,11 +9,24 @@
 from kili.domain.types import ListOrTuple
 from kili.exceptions import NotFound
 
+# ============================================================================
+# TODO(LAB-4269): TEMPORARY WORKAROUND - Remove when backend handles jsonResponseUrl for LLM
+# This entire function is only used for LLM projects that don't have jsonResponseUrl.
+# Once the backend properly computes jsonResponseUrl for LLM projects, this function
+# can be deleted entirely.
+# ============================================================================
+
 
 def get_annotation_fragment() -> str:
     """Generate a basic annotation fragment for querying annotations.
 
-    This is used for LLM projects.
+    TODO(LAB-4269): Remove this function when backend handles jsonResponseUrl for LLM projects.
+
+    This is used ONLY for LLM projects that don't have jsonResponseUrl computed by the backend.
+    The fragment is used to fetch annotations so we can rebuild jsonResponse client-side.
+
+    Returns:
+        GraphQL fragment for querying LLM annotations (classification, comparison, transcription)
     """
     return """
         __typename
@@ -46,6 +59,11 @@ def get_annotation_fragment() -> str:
     """
 
 
+# ============================================================================
+# END TODO(LAB-4269) - LLM workaround function
+# ============================================================================
+
+
 def get_asset(
     graphql_client: GraphQLClient,
     http_client: HttpClient,