feat: [DEV-12997]: Bundles (#391)

Sung96kim · web-flow · commit 6552464e0f9f · 2025-07-14T15:40:07.000-04:00
diff --git a/indico/client/request.py b/indico/client/request.py
@@ -70,6 +70,29 @@ def process_response(self, response: "AnyDict") -> "ResponseType":
         return cast("ResponseType", raw_response)
 
 
+def _parse_nested_response(
+    response: "AnyDict", nested_keys: "List[str | int]"
+) -> "Any":
+    composite: "Any" = response
+    for key in nested_keys:
+        if isinstance(composite, list):
+            if not isinstance(key, int):
+                raise IndicoInputError(
+                    f"Invalid nested key type: {type(key)}",
+                )
+            composite = composite[int(key)]
+            continue
+
+        if isinstance(composite, dict):
+            if key not in composite.keys():
+                raise IndicoInputError(
+                    f"Nested key not found in response: {key}",
+                )
+            composite = composite[key]
+
+    return composite
+
+
 class PagedRequest(GraphQLRequest[ResponseType]):
     """
     To enable pagination, query must include $after as an argument
@@ -100,19 +123,11 @@ def __init__(self, query: str, variables: "Optional[AnyDict]" = None):
         super().__init__(query, variables=variables)
 
     def parse_payload(
-        self, response: "AnyDict", nested_keys: "Optional[List[str]]" = None
+        self, response: "AnyDict", nested_keys: "Optional[List[str | int]]" = None
     ) -> "Any":
         raw_response: "AnyDict" = cast("AnyDict", super().parse_payload(response))
-
         if nested_keys:
-            composite = raw_response
-            for key in nested_keys:
-                if key not in composite.keys():
-                    raise IndicoInputError(
-                        f"Nested key not found in response: {key}",
-                    )
-                composite = composite[key]
-
+            composite: "Any" = _parse_nested_response(raw_response, nested_keys)
             _pg = composite.get("pageInfo")
         else:
             _pg = next(iter(raw_response.values())).get("pageInfo")
@@ -157,19 +172,12 @@ def __init__(self, query: str, variables: "Optional[AnyDict]" = None):
         super().__init__(query, variables=variables)
 
     def parse_payload(
-        self, response: "AnyDict", nested_keys: "Optional[List[str]]" = None
+        self, response: "AnyDict", nested_keys: "Optional[List[str | int]]" = None
     ) -> "Any":
         raw_response: "AnyDict" = cast("AnyDict", super().parse_payload(response))
 
         if nested_keys:
-            composite = raw_response
-            for key in nested_keys:
-                if key not in composite.keys():
-                    raise IndicoInputError(
-                        f"Nested key not found in response: {key}",
-                    )
-                composite = composite[key]
-
+            composite: "Any" = _parse_nested_response(raw_response, nested_keys)
             pagination_data = composite
         else:
             pagination_data = next(iter(raw_response.values()))
diff --git a/indico/filters/__init__.py b/indico/filters/__init__.py
@@ -147,26 +147,50 @@ class ModelGroupExampleFilter(Filter):
 
     Args:
         file_name (str): examples with input file names containing this string
+        file_type (str): examples with this file type (e.g. "PDF")
+        labeler (List[int]): examples labeled by these users
+        include_class (List[int]): examples that are or are not included in this class
+        exclude_class (List[int]): examples that are or are not excluded from this class
         partial (bool): examples that are or are not partially labeled
+        autolabeled (bool): examples that are or are not autolabeled
         status (str): submissions in this status. Options:
             [COMPLETE, INCOMPLETE]
         text_search (bool): examples that contain this substring in their text
     Returns:
         dict containing query filter parameters
     """
 
-    # TODO: extend to support full filter list
-    __options__ = ("file_name", "partial", "status", "text_search")
+    __options__ = (
+        "file_name",
+        "file_type",
+        "labeler",
+        "include_class",
+        "exclude_class",
+        "partial",
+        "autolabeled",
+        "status",
+        "text_search",
+    )
 
     def __init__(
         self,
         file_name: "Optional[str]" = None,
+        file_type: "Optional[str]" = None,
+        labeler: "Optional[List[int]]" = None,
+        include_class: "Optional[List[int]]" = None,
+        exclude_class: "Optional[List[int]]" = None,
+        autolabeled: "Optional[bool]" = None,
         partial: "Optional[bool]" = None,
         status: "Optional[str]" = None,
         text_search: "Optional[str]" = None,
     ):
-        kwargs: "Dict[str, Optional[Union[bool, str]]]" = {
+        kwargs: "Dict[str, Optional[Union[bool, str, List[int]]]]" = {
             "fileName": file_name,
+            "fileType": file_type,
+            "labeler": labeler,
+            "includeClass": include_class,
+            "excludeClass": exclude_class,
+            "autolabeled": autolabeled,
             "partial": partial,
             "textSearch": text_search,
         }
diff --git a/indico/queries/example.py b/indico/queries/example.py
@@ -18,24 +18,28 @@ class ListModelGroupExamples(PagedRequest["List[Example]"]):
     Options:
         model_group_ids (List[int]): Model group ids to filter by
         filters (ModelGroupExampleFilter or Dict): Example attributes to filter by
-        limit (int, default=1000): Maximum number of Examples to return
+        limit (int, default=1000): Maximum number of Examples to return per page
         orderBy (str, default="ID"): Example attribute to filter by
         desc: (bool, default=True): List in descending order
+        after (int): id to start the pagination from
+        before (int): id to end the pagination at
 
     Returns:
         List[Example]: All the found Example objects
         If paginated, yields results one at a time
     """
 
     query = """
-        query GetExamples($modelGroupId:Int, $orderBy: ExampleOrder, $desc: Boolean, $limit: Int, $filters: ModelGroupExampleFilter, $after: Int, $before: Int) {
+        query GetExamples($modelGroupId:Int, $orderBy: ExampleOrder, $desc: Boolean, $limit: Int, $filters: ExampleFilter, $after: Int, $before: Int) {
             modelGroups(modelGroupIds: [$modelGroupId]) {
                 modelGroups {
-                    pagedExamples(orderBy:$orderBy, desc:$desc, limit: $limit, filters: $filters, after: $after, before: $before) {
+                    pagedExamples(orderBy: $orderBy, desc: $desc, limit: $limit, filters: $filters, after: $after, before: $before) {
                         examples {
                             id
                             status
-                            datafileId
+                            datafileIds
+                            originalDatafileId
+                            originalDatafileName
                         }
                         pageInfo {
                             startCursor
@@ -74,5 +78,12 @@ def __init__(
         )
 
     def process_response(self, response: "Payload") -> "List[Example]":
-        example_page = super().parse_payload(response)["modelGroups"]["modelGroups"][0]
-        return [Example(**s) for s in example_page["pagedExamples"]["examples"]]
+        example_page = super().parse_payload(
+            response, nested_keys=["modelGroups", "modelGroups", 0, "pagedExamples"]
+        )
+        return [
+            Example(**example)
+            for example in example_page["modelGroups"]["modelGroups"][0][
+                "pagedExamples"
+            ]["examples"]
+        ]
diff --git a/indico/queries/questionnaire.py b/indico/queries/questionnaire.py
@@ -75,9 +75,15 @@ class GetQuestionnaireExamples(GraphQLRequest["List[Example]"]):
         questionnaires(questionnaireIds: [$questionnaire_id]) {
             questionnaires {
                 examples(numExamples: $num_examples, datafileId: $datafile_id) {
-                    rowIndex
-                    datafileId
-                    source
+                    datafileIds
+                    originalDatafileId
+                    originalDatafileName
+                    contexts {
+                        id
+                        datafileId
+                        source
+                    }
+                    status
                     id
                 }
             }
diff --git a/indico/types/questionnaire.py b/indico/types/questionnaire.py
@@ -26,11 +26,18 @@ class Questionnaire(BaseType):
     question: Question
 
 
-class Example(BaseType):
-    row_index: int
+class ExampleContext(BaseType):
+    id: int
     datafile_id: int
-    status: str
     source: str
+
+
+class Example(BaseType):
+    datafile_ids: List[int]
+    original_datafile_id: int
+    original_datafile_name: str
+    contexts: List[ExampleContext]
+    status: str
     id: int
 
 
diff --git a/tests/integration/queries/test_questionnaire.py b/tests/integration/queries/test_questionnaire.py
@@ -139,9 +139,10 @@ def test_get_examples(indico, unlabeled_questionnaire):
     assert len(examples) == 3
     for example in examples:
         assert isinstance(example, Example)
-        assert isinstance(example.source, str)
-        assert isinstance(example.row_index, int)
-        assert isinstance(example.datafile_id, int)
+        assert isinstance(example.contexts, list)
+        assert isinstance(example.datafile_ids, list)
+        assert isinstance(example.original_datafile_id, int)
+        assert isinstance(example.original_datafile_name, str)
 
 
 def test_add_labels(indico, unlabeled_questionnaire):
