From f9f9a2bdc64e2259a24437d6dfe8827003c586cf Mon Sep 17 00:00:00 2001
From: Michal Noszczak <mnoszczak@labelbox.com>
Date: Thu, 2 Jan 2025 15:48:46 +0000
Subject: [PATCH] Support include predictions in catalog exports

---
 libs/labelbox/src/labelbox/schema/catalog.py       | 1 +
 libs/labelbox/src/labelbox/schema/data_row.py      | 1 +
 libs/labelbox/src/labelbox/schema/dataset.py       | 1 +
 libs/labelbox/src/labelbox/schema/export_params.py | 1 +
 libs/labelbox/src/labelbox/schema/slice.py         | 1 +
 5 files changed, 5 insertions(+)

diff --git a/libs/labelbox/src/labelbox/schema/catalog.py b/libs/labelbox/src/labelbox/schema/catalog.py
index 382aab683..7202e1ca4 100644
--- a/libs/labelbox/src/labelbox/schema/catalog.py
+++ b/libs/labelbox/src/labelbox/schema/catalog.py
@@ -104,6 +104,7 @@ def _export(
                 "interpolated_frames": False,
                 "all_projects": False,
                 "all_model_runs": False,
+                "predictions": False,
             }
         )
         validate_catalog_export_params(_params)
diff --git a/libs/labelbox/src/labelbox/schema/data_row.py b/libs/labelbox/src/labelbox/schema/data_row.py
index 48bc6da85..7450f4e0b 100644
--- a/libs/labelbox/src/labelbox/schema/data_row.py
+++ b/libs/labelbox/src/labelbox/schema/data_row.py
@@ -315,6 +315,7 @@ def _export(
                 "interpolated_frames": False,
                 "all_projects": False,
                 "all_model_runs": False,
+                "predictions": False,
             }
         )
 
diff --git a/libs/labelbox/src/labelbox/schema/dataset.py b/libs/labelbox/src/labelbox/schema/dataset.py
index 3d83b0549..ce4d8416a 100644
--- a/libs/labelbox/src/labelbox/schema/dataset.py
+++ b/libs/labelbox/src/labelbox/schema/dataset.py
@@ -392,6 +392,7 @@ def _export(
                 "interpolated_frames": False,
                 "all_projects": False,
                 "all_model_runs": False,
+                "predictions": False,
             }
         )
         validate_catalog_export_params(_params)
diff --git a/libs/labelbox/src/labelbox/schema/export_params.py b/libs/labelbox/src/labelbox/schema/export_params.py
index f921cdd31..9976db796 100644
--- a/libs/labelbox/src/labelbox/schema/export_params.py
+++ b/libs/labelbox/src/labelbox/schema/export_params.py
@@ -29,6 +29,7 @@ class CatalogExportParams(DataRowParams):
     interpolated_frames: Optional[bool]
     all_projects: Optional[bool]
     all_model_runs: Optional[bool]
+    predictions: Optional[bool]
 
 
 class ModelRunExportParams(DataRowParams):
diff --git a/libs/labelbox/src/labelbox/schema/slice.py b/libs/labelbox/src/labelbox/schema/slice.py
index f6bb53d75..a47a9da6d 100644
--- a/libs/labelbox/src/labelbox/schema/slice.py
+++ b/libs/labelbox/src/labelbox/schema/slice.py
@@ -160,6 +160,7 @@ def _export(
                 "interpolated_frames": False,
                 "all_projects": False,
                 "all_model_runs": False,
+                "predictions": False,
             }
         )
         validate_catalog_export_params(_params)