terrakio-core 0.4.98.1b4__py3-none-any.whl → 0.4.98.1b6__py3-none-any.whl

terrakio_core/endpoints/mass_stats.py

@@ -1,34 +1,34 @@
+import json
 import os
 import time
-import typer
 from typing import Dict, Any, Optional, List
-from dateutil import parser
 
+import aiohttp
+import typer
+from dateutil import parser
 from rich.console import Console
 from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeElapsedColumn
 
 from ..exceptions import (
+    CancelAllTasksError,
+    CancelCollectionTasksError,
+    CancelTaskError,
+    CollectionAlreadyExistsError,
+    CollectionNotFoundError,
     CreateCollectionError,
+    DeleteCollectionError,
+    DownloadFilesError,
     GetCollectionError,
-    ListCollectionsError,
-    CollectionNotFoundError,
-    CollectionAlreadyExistsError,
+    GetTaskError,
     InvalidCollectionTypeError,
-    DeleteCollectionError,
+    ListCollectionsError,
     ListTasksError,
-    UploadRequestsError,
-    UploadArtifactsError,
-    GetTaskError,
     TaskNotFoundError,
-    DownloadFilesError,
-    CancelTaskError,
-    CancelCollectionTasksError,
-    CancelAllTasksError,
+    UploadArtifactsError,
+    UploadRequestsError,
 )
 from ..helper.decorators import require_api_key
 
-import aiohttp  # Make sure this is imported at the top
-
 
 class MassStats:
     def __init__(self, client):
@@ -156,6 +156,7 @@ class MassStats:
         
         self.console.print(f"[bold green]All {number_of_jobs} jobs finished![/bold green]")
 
+    # below are functions related to collection
     @require_api_key
     async def create_collection(
         self, 
@@ -202,45 +203,6 @@ class MassStats:
 
         return response
 
-    @require_api_key
-    async def delete_collection(
-        self, 
-        collection: str, 
-        full: Optional[bool] = False, 
-        outputs: Optional[list] = [], 
-        data: Optional[bool] = False
-    ) -> Dict[str, Any]:
-        """
-        Delete a collection by name.
-
-        Args:
-            collection: The name of the collection to delete (required)
-            full: Delete the full collection (optional, defaults to False)
-            outputs: Specific output folders to delete (optional, defaults to empty list)
-            data: Whether to delete raw data (xdata folder) (optional, defaults to False)
-            
-        Returns:
-            API response as a dictionary confirming deletion
-            
-        Raises:
-            CollectionNotFoundError: If the collection is not found
-            DeleteCollectionError: If the API request fails due to unknown reasons
-        """
-        payload = {
-            "full": full,
-            "outputs": outputs,
-            "data": data
-        }
-        
-        response, status = await self._client._terrakio_request("DELETE", f"collections/{collection}", json=payload)
-
-        if status != 200:
-            if status == 404:
-                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise DeleteCollectionError(f"Delete collection failed with status {status}", status_code=status)
-
-        return response
-
     @require_api_key
     async def get_collection(self, collection: str) -> Dict[str, Any]:
         """
@@ -304,95 +266,72 @@ class MassStats:
         return response
 
     @require_api_key
-    async def list_tasks(
-        self,
-        limit: Optional[int] = 10,
-        page: Optional[int] = 0
-    ) -> List[Dict[str, Any]]:
+    async def delete_collection(
+        self, 
+        collection: str, 
+        full: Optional[bool] = False, 
+        outputs: Optional[list] = [], 
+        data: Optional[bool] = False
+    ) -> Dict[str, Any]:
         """
-        List tasks for the current user.
+        Delete a collection by name.
 
         Args:
-            limit: Number of tasks to return (optional, defaults to 10)
-            page: Page number (optional, defaults to 0)
-        
+            collection: The name of the collection to delete (required)
+            full: Delete the full collection (optional, defaults to False)
+            outputs: Specific output folders to delete (optional, defaults to empty list)
+            data: Whether to delete raw data (xdata folder) (optional, defaults to False)
+            
         Returns:
-            API response as a list of dictionaries containing task information
+            API response as a dictionary confirming deletion
             
         Raises:
-            ListTasksError: If the API request fails due to unknown reasons
+            CollectionNotFoundError: If the collection is not found
+            DeleteCollectionError: If the API request fails due to unknown reasons
         """
-        params = {
-            "limit": limit,
-            "page": page
+        payload = {
+            "full": full,
+            "outputs": outputs,
+            "data": data
         }
-        response, status = await self._client._terrakio_request("GET", "tasks", params=params)
-
-        if status != 200:
-            raise ListTasksError(f"List tasks failed with status {status}", status_code=status)
-
-        return response
-
-    @require_api_key
-    async def upload_requests(
-        self,
-        collection: str
-    ) -> Dict[str, Any]:
-        """
-        Retrieve signed url to upload requests for a collection.
-
-        Args:
-            collection: Name of collection
         
-        Returns:
-            API response as a dictionary containing the upload URL
-
-        Raises:
-            CollectionNotFoundError: If the collection is not found
-            UploadRequestsError: If the API request fails due to unknown reasons
-        """
-        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/upload/requests")
+        response, status = await self._client._terrakio_request("DELETE", f"collections/{collection}", json=payload)
 
         if status != 200:
             if status == 404:
                 raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise UploadRequestsError(f"Upload requests failed with status {status}", status_code=status)
+            raise DeleteCollectionError(f"Delete collection failed with status {status}", status_code=status)
 
         return response
 
+    # below are functions related to tasks
     @require_api_key
-    async def upload_artifacts(
+    async def list_tasks(
         self,
-        collection: str,
-        file_type: str,
-        compressed: Optional[bool] = True
-    ) -> Dict[str, Any]:
+        limit: Optional[int] = 10,
+        page: Optional[int] = 0
+    ) -> List[Dict[str, Any]]:
         """
-        Retrieve signed url to upload artifact file to a collection.
+        List tasks for the current user.
 
         Args:
-            collection: Name of collection
-            file_type: The extension of the file
-            compressed: Whether to compress the file using gzip or not (defaults to True)
+            limit: Number of tasks to return (optional, defaults to 10)
+            page: Page number (optional, defaults to 0)
         
         Returns:
-            API response as a dictionary containing the upload URL
-
+            API response as a list of dictionaries containing task information
+            
         Raises:
-            CollectionNotFoundError: If the collection is not found
-            UploadArtifactsError: If the API request fails due to unknown reasons
+            ListTasksError: If the API request fails due to unknown reasons
         """
         params = {
-            "file_type": file_type,
-            "compressed": str(compressed).lower(),
+            "limit": limit,
+            "page": page
         }
-
-        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/upload", params=params)
+        response, status = await self._client._terrakio_request("GET", "tasks", params=params)
 
         if status != 200:
-            if status == 404:
-                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise UploadArtifactsError(f"Upload artifacts failed with status {status}", status_code=status)
+            raise ListTasksError(f"List tasks failed with status {status}", status_code=status)
 
         return response
         
@@ -424,205 +363,114 @@ class MassStats:
         return response
 
     @require_api_key
-    async def generate_data(
+    async def cancel_task(
         self,
-        collection: str,
-        output: str,
-        skip_existing: Optional[bool] = True,
-        force_loc: Optional[bool] = None,
-        server: Optional[str] = None
+        task_id: str
     ) -> Dict[str, Any]:
         """
-        Generate data for a collection.
+        Cancel a task by task ID.
 
         Args:
-            collection: Name of collection
-            output: Output type (str)
-            force_loc: Write data directly to the cloud under this folder
-            skip_existing: Skip existing data
-            server: Server to use
-        
+            task_id: ID of task to cancel
+
         Returns:
             API response as a dictionary containing task information
 
         Raises:
-            CollectionNotFoundError: If the collection is not found
-            GetTaskError: If the API request fails due to unknown reasons
+            TaskNotFoundError: If the task is not found
+            CancelTaskError: If the API request fails due to unknown reasons
         """
-        payload = {"output": output, "skip_existing": skip_existing}
-        
-        if force_loc is not None:
-            payload["force_loc"] = force_loc
-        if server is not None:
-            payload["server"] = server
+        response, status = await self._client._terrakio_request("POST", f"tasks/cancel/{task_id}")
         
-        response, status = await self._client._terrakio_request("POST", f"collections/{collection}/generate_data", json=payload)
-
         if status != 200:
             if status == 404:
-                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise GetTaskError(f"Generate data failed with status {status}", status_code=status)
-        
+                raise TaskNotFoundError(f"Task {task_id} not found", status_code=status)
+            raise CancelTaskError(f"Cancel task failed with status {status}", status_code=status)
+
         return response
 
     @require_api_key
-    async def training_samples(
+    async def cancel_collection_tasks(
         self,
-        collection: str,
-        expressions: list[str],
-        filters: list[str],
-        aoi: dict,
-        samples: int,
-        crs: str,
-        tile_size: int,
-        res: float,
-        output: str,
-        year_range: Optional[list[int]] = None,
-        server: Optional[str] = None
+        collection: str
     ) -> Dict[str, Any]:
         """
-        Generate training samples for a collection.
+        Cancel all tasks for a collection.
 
         Args:
             collection: Name of collection
-            expressions: List of expressions for each sample
-            filters: Expressions to filter sample areas
-            aoi: AOI to sample from (geojson dict)
-            samples: Number of samples to generate
-            crs: CRS of AOI
-            tile_size: Pixel width and height of samples
-            res: Resolution of samples
-            output: Sample output type
-            year_range: Optional year range filter
-            server: Server to use
 
         Returns:
-            API response as a dictionary containing task information
+            API response as a dictionary containing task information for the collection
 
         Raises:
             CollectionNotFoundError: If the collection is not found
-            GetTaskError: If the API request fails due to unknown reasons
+            CancelCollectionTasksError: If the API request fails due to unknown reasons
         """
-        payload = {
-            "expressions": expressions,
-            "filters": filters,
-            "aoi": aoi,
-            "samples": samples,
-            "crs": crs,
-            "tile_size": tile_size,
-            "res": res,
-            "output": output
-        }
-        
-        if year_range is not None:
-            payload["year_range"] = year_range
-        if server is not None:
-            payload["server"] = server
+        response, status = await self._client._terrakio_request("POST", f"collections/{collection}/cancel")
         
-        response, status = await self._client._terrakio_request("POST", f"collections/{collection}/training_samples", json=payload)
-
         if status != 200:
             if status == 404:
                 raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise GetTaskError(f"Training sample failed with status {status}", status_code=status)
-        
+            raise CancelCollectionTasksError(f"Cancel collection tasks failed with status {status}", status_code=status)
+    
         return response
 
-    # @require_api_key
-    # async def post_processing(
-    #     self,
-    #     collection: str,
-    #     folder: str,
-    #     consumer: str
-    # ) -> Dict[str, Any]:
-    #     """
-    #     Run post processing for a collection.
-
-    #     Args:
-    #         collection: Name of collection
-    #         folder: Folder to store output
-    #         consumer: Post processing script
-
-    #     Returns:
-    #         API response as a dictionary containing task information
-
-    #     Raises:
-    #         CollectionNotFoundError: If the collection is not found
-    #         GetTaskError: If the API request fails due to unknown reasons
-    #     """
-    #     # payload = {
-    #     #     "folder": folder,
-    #     #     "consumer": consumer
-    #     # }
-    #     # we have the consumer as a string, we need to read in the file and then pass in the content
-    #     with open(consumer, 'rb') as f:
-    #         files = {
-    #             'consumer': ('consumer.py', f.read(), 'text/plain')
-    #         }
-    #         data = {
-    #             'folder': folder
-    #         }
-        
-    #     # response, status = await self._client._terrakio_request("POST", f"collections/{collection}/post_process", json=payload)
-    #     response, status = await self._client._terrakio_request(
-    #         "POST",
-    #         f"collections/{collection}/post_process",
-    #         files=files,
-    #         data=data
-    #     )
-    #     if status != 200:
-    #         if status == 404:
-    #             raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-    #         raise GetTaskError(f"Post processing failed with status {status}", status_code=status)
-        
-    #     return response
+    @require_api_key
+    async def cancel_all_tasks(
+        self
+    ) -> Dict[str, Any]:
+        """
+        Cancel all tasks for the current user.
+
+        Returns:
+            API response as a dictionary containing task information for all tasks
+
+        Raises:
+            CancelAllTasksError: If the API request fails due to unknown reasons
+        """
+        response, status = await self._client._terrakio_request("POST", "tasks/cancel")
+
+        if status != 200:
+            raise CancelAllTasksError(f"Cancel all tasks failed with status {status}", status_code=status)
 
+        return response
 
+    # below are functions related to the web ui and needs to be deleted in the future
     @require_api_key
-    async def post_processing(
+    async def upload_artifacts(
         self,
         collection: str,
-        folder: str,
-        consumer: str
+        file_type: str,
+        compressed: Optional[bool] = True
     ) -> Dict[str, Any]:
         """
-        Run post processing for a collection.
+        Retrieve signed url to upload artifact file to a collection.
 
         Args:
             collection: Name of collection
-            folder: Folder to store output
-            consumer: Path to post processing script
-
+            file_type: The extension of the file
+            compressed: Whether to compress the file using gzip or not (defaults to True)
+        
         Returns:
-            API response as a dictionary containing task information
+            API response as a dictionary containing the upload URL
 
         Raises:
             CollectionNotFoundError: If the collection is not found
-            GetTaskError: If the API request fails due to unknown reasons
+            UploadArtifactsError: If the API request fails due to unknown reasons
         """
-        # Read file and build multipart form data
-        with open(consumer, 'rb') as f:
-            form = aiohttp.FormData()
-            form.add_field('folder', folder)  # Add text field
-            form.add_field(
-                'consumer',  # Field name
-                f.read(),  # File content
-                filename='consumer.py',  # Filename
-                content_type='text/x-python'  # MIME type
-            )
-        
-        # Send using data= with FormData object (NOT files=)
-        response, status = await self._client._terrakio_request(
-            "POST",
-            f"collections/{collection}/post_process",
-            data=form  # ✅ Pass FormData as data
-        )
-        
+        params = {
+            "file_type": file_type,
+            "compressed": str(compressed).lower(),
+        }
+
+        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/upload", params=params)
+
         if status != 200:
             if status == 404:
                 raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise GetTaskError(f"Post processing failed with status {status}", status_code=status)
-        
+            raise UploadArtifactsError(f"Upload artifacts failed with status {status}", status_code=status)
+
         return response
 
     @require_api_key
@@ -716,120 +564,425 @@ class MassStats:
         
         return response
 
+    async def _upload_requests(
+        self,
+        collection: str
+    ) -> Dict[str, Any]:
+        """
+        Retrieve signed url to upload requests for a collection.
+
+        Args:
+            collection: Name of collection
+        
+        Returns:
+            API response as a dictionary containing the upload URL
+
+        Raises:
+            CollectionNotFoundError: If the collection is not found
+            UploadRequestsError: If the API request fails due to unknown reasons
+        """
+        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/upload/requests")
+
+        if status != 200:
+            if status == 404:
+                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
+            raise UploadRequestsError(f"Upload requests failed with status {status}", status_code=status)
+
+        return response
+
+    @require_api_key
+    async def _upload_file(self, file_path: str, url: str, use_gzip: bool = True):
+        """
+        Helper method to upload a JSON file to a signed URL.
+        
+        Args:
+            file_path: Path to the JSON file
+            url: Signed URL to upload to
+            use_gzip: Whether to compress the file with gzip
+        """
+        try:
+            with open(file_path, 'r') as file:
+                json_data = json.load(file)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"JSON file not found: {file_path}")
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in file {file_path}: {e}")
+        
+        return await self._upload_json_data(json_data, url, use_gzip)
+
+    @require_api_key
+    async def _upload_json_data(self, json_data, url: str, use_gzip: bool = True):
+        """
+        Helper method to upload JSON data directly to a signed URL.
+        
+        Args:
+            json_data: JSON data (dict or list) to upload
+            url: Signed URL to upload to
+            use_gzip: Whether to compress the data with gzip
+        """
+        if hasattr(json, 'dumps') and 'ignore_nan' in json.dumps.__code__.co_varnames:
+            dumps_kwargs = {'ignore_nan': True}
+        else:
+            dumps_kwargs = {}
+        
+        if use_gzip:
+            import gzip
+            body = gzip.compress(json.dumps(json_data, **dumps_kwargs).encode('utf-8'))
+            headers = {
+                'Content-Type': 'application/json',
+                'Content-Encoding': 'gzip'
+            }
+        else:
+            body = json.dumps(json_data, **dumps_kwargs).encode('utf-8')
+            headers = {
+                'Content-Type': 'application/json'
+            }
+        
+        response = await self._client._regular_request("PUT", url, data=body, headers=headers)
+        return response
+
     @require_api_key
-    async def download_files(
+    async def generate_data(
         self,
         collection: str,
-        file_type: str,
-        page: Optional[int] = 0,
-        page_size: Optional[int] = 100,
-        folder: Optional[str] = None
+        file_path: str,
+        output: str,
+        skip_existing: Optional[bool] = True,
+        force_loc: Optional[bool] = None,
+        server: Optional[str] = None
     ) -> Dict[str, Any]:
         """
-        Get list of signed urls to download files in collection.
+        Generate data for a collection.
 
         Args:
             collection: Name of collection
-            file_type: Whether to return raw or processed (after post processing) files
-            page: Page number (optional, defaults to 0)
-            page_size: Number of files to return per page (optional, defaults to 100)
-            folder: If processed file type, which folder to download files from (optional)
-
+            file_path: Path to the file to upload
+            output: Output type (str)
+            force_loc: Write data directly to the cloud under this folder
+            skip_existing: Skip existing data
+            server: Server to use
+        
         Returns:
-            API response as a dictionary containing list of download URLs
+            API response as a dictionary containing task information
 
         Raises:
             CollectionNotFoundError: If the collection is not found
-            DownloadFilesError: If the API request fails due to unknown reasons
+            GetTaskError: If the API request fails due to unknown reasons
         """
-        params = {"file_type": file_type}
+        await self.create_collection(
+            collection = collection
+        )
+
+        upload_urls = await self._upload_requests(
+            collection = collection
+        )
         
-        if page is not None:
-            params["page"] = page
-        if page_size is not None:
-            params["page_size"] = page_size
-        if folder is not None:
-            params["folder"] = folder
+        url = upload_urls['url']
 
-        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/download", params=params)
+        await self._upload_file(file_path, url)
+        
+        payload = {"output": output, "skip_existing": skip_existing}
+        
+        if force_loc is not None:
+            payload["force_loc"] = force_loc
+        if server is not None:
+            payload["server"] = server
+        
+        response, status = await self._client._terrakio_request("POST", f"collections/{collection}/generate_data", json=payload)
 
         if status != 200:
             if status == 404:
                 raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise DownloadFilesError(f"Download files failed with status {status}", status_code=status)
+            raise GetTaskError(f"Generate data failed with status {status}", status_code=status)
         
         return response
 
     @require_api_key
-    async def cancel_task(
+    async def post_processing(
         self,
-        task_id: str
-    ):
+        collection: str,
+        folder: str,
+        consumer: str
+    ) -> Dict[str, Any]:
         """
-        Cancel a task by task ID.
+        Run post processing for a collection.
 
         Args:
-            task_id: ID of task to cancel
+            collection: Name of collection
+            folder: Folder to store output
+            consumer: Path to post processing script
 
         Returns:
             API response as a dictionary containing task information
 
         Raises:
-            TaskNotFoundError: If the task is not found
-            CancelTaskError: If the API request fails due to unknown reasons
+            CollectionNotFoundError: If the collection is not found
+            GetTaskError: If the API request fails due to unknown reasons
         """
-        response, status = await self._client._terrakio_request("POST", f"tasks/cancel/{task_id}")
+        await self.create_collection(
+            collection = collection
+        )
+
+        with open(consumer, 'rb') as f:
+            form = aiohttp.FormData()
+            form.add_field('folder', folder)
+            form.add_field(
+                'consumer',
+                f.read(),
+                filename='consumer.py',
+                content_type='text/x-python'
+            )
+        
+        response, status = await self._client._terrakio_request(
+            "POST",
+            f"collections/{collection}/post_process",
+            data=form
+        )
+        
         if status != 200:
             if status == 404:
-                raise TaskNotFoundError(f"Task {task_id} not found", status_code=status)
-            raise CancelTaskError(f"Cancel task failed with status {status}", status_code=status)
-
+                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
+            raise GetTaskError(f"Post processing failed with status {status}", status_code=status)
+        
         return response
 
     @require_api_key
-    async def cancel_collection_tasks(
+    async def training_samples(
         self,
-        collection: str
-    ):
-        """
-        Cancel all tasks for a collection.
+        name: str,
+        aoi: str,
+        expression_x: str,
+        filter_x: str = "skip",
+        filter_x_rate: float = 1,
+        expression_y: str = "skip",
+        filter_y: str = "skip",
+        filter_y_rate: float = 1,
+        samples: int = 1000,
+        tile_size: float = 256,
+        crs: str = "epsg:3577",
+        res: float = 10,
+        res_y: float = None,
+        skip_test: bool = False,
+        start_year: int = None,
+        end_year: int = None,
+        bucket: str = None,
+        server: str = None,
+        extra_filters: list[str] = None,
+        extra_filters_rate: list[float] = None,
+        extra_filters_res: list[float] = None
+    ) -> dict:
+        """
+        Generate an AI dataset using specified parameters.
 
         Args:
-            collection: Name of collection
+            name: Name of the collection to create
+            aoi: Path to GeoJSON file containing area of interest
+            expression_x: Expression for X data (features)
+            filter_x: Filter expression for X data (default: "skip")
+            filter_x_rate: Filter rate for X data (default: 1)
+            expression_y: Expression for Y data (labels) (default: "skip")
+            filter_y: Filter expression for Y data (default: "skip")
+            filter_y_rate: Filter rate for Y data (default: 1)
+            samples: Number of samples to generate (default: 1000)
+            tile_size: Size of tiles in pixels (default: 256)
+            crs: Coordinate reference system (default: "epsg:3577")
+            res: Resolution for X data (default: 10)
+            res_y: Resolution for Y data, defaults to res if None
+            skip_test: Skip expression validation test (default: False)
+            start_year: Start year for temporal filtering
+            end_year: End year for temporal filtering
+            bucket: Storage bucket name
+            server: Server to use for processing
+            extra_filters: Additional filter expressions
+            extra_filters_rate: Rates for additional filters
+            extra_filters_res: Resolutions for additional filters
 
         Returns:
-            API response as a dictionary containing task information for the collection
+            Response containing task_id and collection name
 
         Raises:
             CollectionNotFoundError: If the collection is not found
-            CancelCollectionTasksError: If the API request fails due to unknown reasons
+            GetTaskError: If the API request fails
+            TypeError: If extra filters have mismatched rate and resolution lists
         """
+        expressions = [{"expr": expression_x, "res": res, "prefix": "x"}]
+        
+        res_y = res_y or res
+        
+        if expression_y != "skip":
+            expressions.append({"expr": expression_y, "res": res_y, "prefix": "y"})
+        
+        filters = []
+        if filter_x != "skip":
+            filters.append({"expr": filter_x, "res": res, "rate": filter_x_rate})
+        
+        if filter_y != "skip":
+            filters.append({"expr": filter_y, "res": res_y, "rate": filter_y_rate})
+        
+        if extra_filters:
+            try:
+                extra_filters_combined = zip(extra_filters, extra_filters_res, extra_filters_rate, strict=True)
+            except TypeError:
+                raise TypeError("Extra filters must have matching rate and resolution.")
+            
+            for expr, filter_res, rate in extra_filters_combined:
+                filters.append({"expr": expr, "res": filter_res, "rate": rate})
+        
+        if start_year is not None:
+            for expr_dict in expressions:
+                expr_dict["expr"] = expr_dict["expr"].replace("{year}", str(start_year))
+            
+            for filter_dict in filters:
+                filter_dict["expr"] = filter_dict["expr"].replace("{year}", str(start_year))
+        
+        if not skip_test:
+            for expr_dict in expressions:
+                test_request = self._client.model._generate_test_request(expr_dict["expr"], crs, -1)
+                await self._client._terrakio_request("POST", "geoquery", json=test_request)
+            
+            for filter_dict in filters:
+                test_request = self._client.model._generate_test_request(filter_dict["expr"], crs, -1)
+                await self._client._terrakio_request("POST", "geoquery", json=test_request)
+        
+        with open(aoi, 'r') as f:
+            aoi_data = json.load(f)
+        
+        await self.create_collection(
+            collection = name,
+            bucket = bucket,
+            collection_type = "basic"
+        )
+
+        payload = {
+            "expressions": expressions,
+            "filters": filters,
+            "aoi": aoi_data,
+            "samples": samples,
+            "crs": crs,
+            "tile_size": tile_size,
+            "res": res,
+            "output": "nc",
+            "year_range": [start_year, end_year],
+            "server": server
+        }
+        
+        task_id_dict, status = await self._client._terrakio_request("POST", f"collections/{name}/training_samples", json=payload)
 
-        response, status = await self._client._terrakio_request("POST", f"collections/{collection}/cancel")
         if status != 200:
             if status == 404:
-                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
-            raise CancelCollectionTasksError(f"Cancel collection tasks failed with status {status}", status_code=status)
-    
-        return response
+                raise CollectionNotFoundError(f"Collection {name} not found", status_code=status)
+            raise GetTaskError(f"Training sample failed with status {status}", status_code=status)
+        
+        task_id = task_id_dict["task_id"]
+        
+        await self._client.mass_stats.track_progress(task_id)
+        
+        return {"task_id": task_id, "collection": name}
 
     @require_api_key
-    async def cancel_all_tasks(
-        self
-    ):
+    async def download_files(
+        self,
+        collection: str,
+        file_type: str,
+        page: Optional[int] = 0,
+        page_size: Optional[int] = 100,
+        folder: Optional[str] = None,
+        url: Optional[bool] = True
+    ) -> Dict[str, Any]:
         """
-        Cancel all tasks for the current user.
+        Get list of signed urls to download files in collection, or download the files directly.
+
+        Args:
+            collection: Name of collection
+            file_type: Type of files to download - must be either 'raw' or 'processed'
+            page: Page number (optional, defaults to 0)
+            page_size: Number of files to return per page (optional, defaults to 100)
+            folder: If processed file type, which folder to download files from (optional)
+            url: If True, return signed URLs; if False, download files directly (optional, defaults to True)
 
         Returns:
-            API response as a dictionary containing task information for all tasks
+            API response as a dictionary containing list of download URLs (if url=True),
+            or a dictionary with downloaded file information (if url=False)
 
         Raises:
-            CancelAllTasksError: If the API request fails due to unknown reasons
+            CollectionNotFoundError: If the collection is not found
+            DownloadFilesError: If the API request fails due to unknown reasons
+            ValueError: If file_type is not 'raw' or 'processed'
         """
+        if file_type not in ['raw', 'processed']:
+            raise ValueError(f"file_type must be either 'raw' or 'processed', got '{file_type}'")
+        
+        params = {"file_type": file_type}
+        
+        if page is not None:
+            params["page"] = page
+        if page_size is not None:
+            params["page_size"] = page_size
+        if folder is not None:
+            params["folder"] = folder
 
-        response, status = await self._client._terrakio_request("POST", "tasks/cancel")
+        response, status = await self._client._terrakio_request("GET", f"collections/{collection}/download", params=params)
 
         if status != 200:
-            raise CancelAllTasksError(f"Cancel all tasks failed with status {status}", status_code=status)
-
-        return response
+            if status == 404:
+                raise CollectionNotFoundError(f"Collection {collection} not found", status_code=status)
+            raise DownloadFilesError(f"Download files failed with status {status}", status_code=status)
+        
+        if url:
+            return response
+        
+        downloaded_files = []
+        files_to_download = response.get('files', []) if isinstance(response, dict) else []
+        
+        async with aiohttp.ClientSession() as session:
+            for file_info in files_to_download:
+                try:
+                    file_url = file_info.get('url')
+                    filename = file_info.get('file', '')
+                    group = file_info.get('group', '')
+                    
+                    if not file_url:
+                        downloaded_files.append({
+                            'filename': filename,
+                            'group': group,
+                            'error': 'No URL provided'
+                        })
+                        continue
+                    
+                    async with session.get(file_url) as file_response:
+                        if file_response.status == 200:
+                            content = await file_response.read()
+                            
+                            output_dir = folder if folder else "downloads"
+                            if group:
+                                output_dir = os.path.join(output_dir, group)
+                            os.makedirs(output_dir, exist_ok=True)
+                            filepath = os.path.join(output_dir, filename)
+                            
+                            with open(filepath, 'wb') as f:
+                                f.write(content)
+                            
+                            downloaded_files.append({
+                                'filename': filename,
+                                'group': group,
+                                'filepath': filepath,
+                                'size': len(content)
+                            })
+                        else:
+                            downloaded_files.append({
+                                'filename': filename,
+                                'group': group,
+                                'error': f"Failed to download: HTTP {file_response.status}"
+                            })
+                except Exception as e:
+                    downloaded_files.append({
+                        'filename': file_info.get('file', 'unknown'),
+                        'group': file_info.get('group', ''),
+                        'error': str(e)
+                    })
+        
+        return {
+            'collection': collection,
+            'downloaded_files': downloaded_files,
+            'total': len(downloaded_files)
+        }

terrakio_core/endpoints/model_management.py

@@ -59,131 +59,6 @@ class ModelManagement:
         }
         return req
 
-    @require_api_key
-    async def generate_ai_dataset(
-        self,
-        name: str,
-        aoi: str,
-        expression_x: str,
-        filter_x: str = "skip",
-        filter_x_rate: float = 1,
-        expression_y: str = "skip",
-        filter_y: str = "skip",
-        filter_y_rate: float = 1,
-        samples: int = 1000,
-        tile_size: float = 256,
-        crs: str = "epsg:3577",
-        res: float = 10,
-        res_y: float = None,
-        skip_test: bool = False,
-        start_year: int = None,
-        end_year: int = None,
-        bucket: str = None,
-        server: str = None,
-        extra_filters: list[str] = None,
-        extra_filters_rate: list[float] = None,
-        extra_filters_res: list[float] = None
-    ) -> dict:
-        """
-        Generate an AI dataset using specified parameters.
-
-        Args:
-            name (str): Name of the collection to create
-            aoi (str): Path to GeoJSON file containing area of interest
-            expression_x (str): Expression for X data (features)
-            filter_x (str): Filter expression for X data (default: "skip")
-            filter_x_rate (float): Filter rate for X data (default: 1)
-            expression_y (str): Expression for Y data (labels) (default: "skip")
-            filter_y (str): Filter expression for Y data (default: "skip")
-            filter_y_rate (float): Filter rate for Y data (default: 1)
-            samples (int): Number of samples to generate (default: 1000)
-            tile_size (float): Size of tiles in pixels (default: 256)
-            crs (str): Coordinate reference system (default: "epsg:3577")
-            res (float): Resolution for X data (default: 10)
-            res_y (float): Resolution for Y data, defaults to res if None
-            skip_test (bool): Skip expression validation test (default: False)
-            start_year (int): Start year for temporal filtering
-            end_year (int): End year for temporal filtering
-            bucket (str): Storage bucket name
-            server (str): Server to use for processing
-            extra_filters (list[str]): Additional filter expressions
-            extra_filters_rate (list[float]): Rates for additional filters
-            extra_filters_res (list[float]): Resolutions for additional filters
-
-        Returns:
-            dict: Response containing task_id and collection name
-
-        Raises:
-            APIError: If the API request fails
-            TypeError: If extra filters have mismatched rate and resolution lists
-        """
-        expressions = [{"expr": expression_x, "res": res, "prefix": "x"}]
-        
-        res_y = res_y or res
-        
-        if expression_y != "skip":
-            expressions.append({"expr": expression_y, "res": res_y, "prefix": "y"})
-        
-        filters = []
-        if filter_x != "skip":
-            filters.append({"expr": filter_x, "res": res, "rate": filter_x_rate})
-        
-        if filter_y != "skip":
-            filters.append({"expr": filter_y, "res": res_y, "rate": filter_y_rate})
-        
-        if extra_filters:
-            try:
-                extra_filters_combined = zip(extra_filters, extra_filters_res, extra_filters_rate, strict=True)
-            except TypeError:
-                raise TypeError("Extra filters must have matching rate and resolution.")
-            
-            for expr, filter_res, rate in extra_filters_combined:
-                filters.append({"expr": expr, "res": filter_res, "rate": rate})
-        
-        if start_year is not None:
-            for expr_dict in expressions:
-                expr_dict["expr"] = expr_dict["expr"].replace("{year}", str(start_year))
-            
-            for filter_dict in filters:
-                filter_dict["expr"] = filter_dict["expr"].replace("{year}", str(start_year))
-        
-        # this is making request to the server that is being used when doing the initialization
-        if not skip_test:
-            for expr_dict in expressions:
-                test_request = self._generate_test_request(expr_dict["expr"], crs, -1)
-                await self._client._terrakio_request("POST", "geoquery", json=test_request)
-            
-            for filter_dict in filters:
-                test_request = self._generate_test_request(filter_dict["expr"], crs, -1)
-                await self._client._terrakio_request("POST", "geoquery", json=test_request)
-        
-        with open(aoi, 'r') as f:
-            aoi_data = json.load(f)
-        
-        await self._client.mass_stats.create_collection(
-            collection=name,
-            bucket=bucket,
-            collection_type="basic"
-        )
-        
-        task_id_dict = await self._client.mass_stats.training_samples(
-            collection=name,
-            expressions=expressions,
-            filters=filters,
-            aoi=aoi_data,
-            samples=samples,
-            year_range=[start_year, end_year],
-            crs=crs,
-            tile_size=tile_size,
-            res=res,
-            output="nc",
-            server=server
-        )
-        
-        task_id = task_id_dict["task_id"]
-        
-        await self._client.mass_stats.track_progress(task_id)
-
     @require_api_key
     async def _get_url_for_upload_model_and_script(self, expression: str, model_name: str, script_name: str) -> str:
         """

{terrakio_core-0.4.98.1b4.dist-info → terrakio_core-0.4.98.1b6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.4.98.1b4
+Version: 0.4.98.1b6
 Summary: Core package for the terrakio-python-api
 Requires-Python: >=3.11
 Requires-Dist: aiofiles>=24.1.0

{terrakio_core-0.4.98.1b4.dist-info → terrakio_core-0.4.98.1b6.dist-info}/RECORD

@@ -11,13 +11,13 @@ terrakio_core/convenience_functions/zonal_stats.py,sha256=7PI--RI0hiF1pzZ7_7hqty
 terrakio_core/endpoints/auth.py,sha256=5WvAO39aLsbJAVtxISwiOZseKr3B9I5tHjamPRehDBQ,8327
 terrakio_core/endpoints/dataset_management.py,sha256=jpwftiKOI59NhXXypqm6wtILIkfyjy9NfYifRIvhZS0,16791
 terrakio_core/endpoints/group_management.py,sha256=V0KOGTXwmePBFeym55G_m3ONR-jVj2IU4OVLiL5UKz4,15869
-terrakio_core/endpoints/mass_stats.py,sha256=YOtV1_yiTJwtJF-hXY5SrOZ7_tpHeXFk5HLoXcL1Yuc,30064
-terrakio_core/endpoints/model_management.py,sha256=AqQkFPIdPuApYC58VLRG0GxD-s6gic7Ffa9InG8lb3I,56505
+terrakio_core/endpoints/mass_stats.py,sha256=Bvke7gmZpNCsbEDyskqV7ZlkztkoBWxxFUPP0DCXz7Y,36970
+terrakio_core/endpoints/model_management.py,sha256=tvJ4BOBsyXKKH430byGH25CkLIzXWgxaPaL0CvL8_0Y,51341
 terrakio_core/endpoints/space_management.py,sha256=YWb55nkJnFJGlALJ520DvurxDqVqwYtsvqQPWzxzhDs,2266
 terrakio_core/endpoints/user_management.py,sha256=L_g4ysrh2xyz_JbObUU_tCxgHxisrDUPntWgQOs15GE,7709
 terrakio_core/helper/bounded_taskgroup.py,sha256=wiTH10jhKZgrsgrFUNG6gig8bFkUEPHkGRT2XY7Rgmo,677
 terrakio_core/helper/decorators.py,sha256=L6om7wmWNgCei3Wy5U0aZ-70OzsCwclkjIf7SfQuhCg,2289
 terrakio_core/helper/tiles.py,sha256=lcLCO6KiP05lCI9vngo3zCZJ6Z9C0pUxHSQS4H58EHc,2699
-terrakio_core-0.4.98.1b4.dist-info/METADATA,sha256=OWd41QUCndnXlcVN5tdYyQD_QleJ6ts6Fa-XVqxz0mg,1184
-terrakio_core-0.4.98.1b4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-terrakio_core-0.4.98.1b4.dist-info/RECORD,,
+terrakio_core-0.4.98.1b6.dist-info/METADATA,sha256=zlli64p2DlufJj4q_xoIYbyuDPra2thNi5r2CtovLJA,1184
+terrakio_core-0.4.98.1b6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+terrakio_core-0.4.98.1b6.dist-info/RECORD,,