terrakio-core 0.4.97__py3-none-any.whl → 0.4.98.1b1__py3-none-any.whl

terrakio_core/endpoints/model_management.py

@@ -1,27 +1,22 @@
-# Standard library imports
 import ast
 import json
 import textwrap
-import time
 from io import BytesIO
 from typing import Optional, Tuple
+
 import onnxruntime as ort
 
-# Internal imports
 from ..helper.decorators import require_api_key
 
-# Optional dependency flags
 TORCH_AVAILABLE = False
 SKL2ONNX_AVAILABLE = False
 
-# PyTorch imports
 try:
     import torch
     TORCH_AVAILABLE = True
 except ImportError:
     torch = None
 
-# Scikit-learn and ONNX conversion imports
 try:
     from sklearn.base import BaseEstimator
     from skl2onnx import convert_sklearn
@@ -36,124 +31,159 @@ class ModelManagement:
     def __init__(self, client):
         self._client = client
 
+    def _generate_test_request(self, expr: str, crs: str, resolution: float) -> dict:
+        """Generate test request using set polygon (Australia)"""
+        req = {
+            "feature": {
+                "type": "Feature",
+                "properties": {},
+                "geometry": {
+                    "coordinates": [
+                        [
+                            [150.57846438251084, -29.535000759011766],
+                            [150.57846438251084, -29.539538891448665],
+                            [150.5845432181327, -29.539538891448665],
+                            [150.5845432181327, -29.535000759011766],
+                            [150.57846438251084, -29.535000759011766]
+                        ]
+                    ],
+                    "type": "Polygon"
+                }
+            },
+            "in_crs": "epsg:4326",
+            "out_crs": crs,
+            "output": "csv",
+            "resolution": resolution,
+            "expr": expr,
+            "debug": True
+        }
+        return req
+
     @require_api_key
     async def generate_ai_dataset(
         self,
         name: str,
-        aoi_geojson: str,
+        aoi: str,
         expression_x: str,
-        filter_x_rate: float,
-        filter_y_rate: float,
-        samples: int,
-        tile_size: int,
-        expression_y: str = "skip",
         filter_x: str = "skip",
+        filter_x_rate: float = 1,
+        expression_y: str = "skip",
         filter_y: str = "skip",
-        crs: str = "epsg:4326",
-        res: float = 0.001,
-        region: str = None,
-        bucket: str = None,
+        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 dataset to generate
-            aoi_geojson (str): Path to GeoJSON file containing area of interest
-            expression_x (str): Expression for X variable (e.g. "MSWX.air_temperature@(year=2021, month=1)")
-            filter_x (str): Filter for X variable (e.g. "MSWX.air_temperature@(year=2021, month=1)")
-            filter_x_rate (float): Filter rate for X variable (e.g. 0.5)
-            expression_y (str): Expression for Y variable with {year} placeholder
-            filter_y (str): Filter for Y variable (e.g. "MSWX.air_temperature@(year=2021, month=1)")
-            filter_y_rate (float): Filter rate for Y variable (e.g. 0.5)
-            samples (int): Number of samples to generate
-            tile_size (int): Size of tiles in degrees
-            crs (str, optional): Coordinate reference system. Defaults to "epsg:4326"
-            res (float, optional): Resolution in degrees. Defaults to 0.001
-            region (str, optional): Region code. Defaults to None
-            bucket (str, optional): Bucket name. Defaults to None
-            start_year (int, optional): Start year for data generation. Required if end_year provided
-            end_year (int, optional): End year for data generation. Required if start_year provided
+            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 from the AI dataset generation API
+            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
         """
-        # Build config for expressions and filters
-        config = {
-            "expressions": [{"expr": expression_x, "res": res, "prefix": "x"}],
-            "filters": []
-        }
-
+        expressions = [{"expr": expression_x, "res": res, "prefix": "x"}]
+        
+        res_y = res_y or res
+        
         if expression_y != "skip":
-            config["expressions"].append({"expr": expression_y, "res": res, "prefix": "y"})
-
+            expressions.append({"expr": expression_y, "res": res_y, "prefix": "y"})
+        
+        filters = []
         if filter_x != "skip":
-            config["filters"].append({"expr": filter_x, "res": res, "rate": filter_x_rate})
+            filters.append({"expr": filter_x, "res": res, "rate": filter_x_rate})
+        
         if filter_y != "skip":
-            config["filters"].append({"expr": filter_y, "res": res, "rate": filter_y_rate})
-
-        # Replace year placeholders if start_year is provided
+            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:
-            expression_x = expression_x.replace("{year}", str(start_year))
-            if expression_y != "skip":
-                expression_y = expression_y.replace("{year}", str(start_year))
-            if filter_x != "skip":
-                filter_x = filter_x.replace("{year}", str(start_year))
-            if filter_y != "skip":
-                filter_y = filter_y.replace("{year}", str(start_year))
-
-        # Load AOI GeoJSON
-        with open(aoi_geojson, 'r') as f:
+            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)
-
-        task_response = await self._client.mass_stats.random_sample(
-            name=name,
-            config=config,
+        
+        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,
-            region=region,
-            output="netcdf",
-            server=self._client.url,
-            bucket=bucket,
-            overwrite=True
+            output="nc",
+            server=server
         )
-        task_id = task_response["task_id"]
-
-        while True:
-            result = await self._client.mass_stats.track_job(ids=[task_id])
-            status = result[task_id]['status']
-            completed = result[task_id].get('completed', 0)
-            total = result[task_id].get('total', 1)
-            
-            progress = completed / total if total > 0 else 0
-            bar_length = 50
-            filled_length = int(bar_length * progress)
-            bar = '█' * filled_length + '░' * (bar_length - filled_length)
-            percentage = progress * 100
-            
-            self._client.logger.info(f"Job status: {status} [{bar}] {percentage:.1f}% ({completed}/{total})")   
-
-            if status == "Completed":
-                self._client.logger.info("Job completed successfully!")
-                break
-            elif status == "Error":
-                self._client.logger.info("Job encountered an error")
-                raise Exception(f"Job {task_id} encountered an error")
-            
-            time.sleep(5)
+        
+        task_id = task_id_dict["task_id"]
+        
+        await self._client.mass_stats.track_progress(task_id)
 
-        task_id = await self._client.mass_stats.start_job(task_id)
-        return task_id
-    
     @require_api_key
     async def _get_url_for_upload_model_and_script(self, expression: str, model_name: str, script_name: str) -> str:
         """
@@ -252,7 +282,7 @@ class ModelManagement:
         return bucket_name
 
     @require_api_key
-    async def upload_and_deploy_model(self, model, virtual_dataset_name: str, virtual_product_name: str, input_expression: str, dates_iso8601: list, input_shape: Tuple[int, ...] = None, processing_script_path: Optional[str] = None, model_type: Optional[str] = None):
+    async def upload_and_deploy_model(self, model, virtual_dataset_name: str, virtual_product_name: str, input_expression: str, dates_iso8601: list, input_shape: Tuple[int, ...] = None, processing_script_path: Optional[str] = None, model_type: Optional[str] = None, padding: int = 0):
         """
         Upload a model to the bucket and deploy it.
         Args:
@@ -264,6 +294,7 @@ class ModelManagement:
             input_shape: Shape of input data for ONNX conversion (required for PyTorch models)
             processing_script_path: Path to the processing script, if not provided, no processing will be done
             model_type: The type of the model we want to upload
+            padding: Padding value for the dataset (default: 0)
 
         Raises:
             APIError: If the API request fails
@@ -278,12 +309,11 @@ class ModelManagement:
         uid = user_info["uid"]
         await self._client.datasets.create_dataset(
             name=virtual_dataset_name,
-            collection="terrakio-datasets",
             products=[virtual_product_name],
-            path=f"gs://{bucket_name}/{uid}/virtual_datasets/{virtual_dataset_name}/inference_scripts",
+            path=f"gs://{bucket_name}/{uid}/models/{virtual_dataset_name}/inference_scripts",
             input=input_expression,
             dates_iso8601=dates_iso8601,
-            padding=0
+            padding=padding
         )
     
     @require_api_key
@@ -1088,15 +1118,32 @@ class ModelManagement:
             raise ValueError(f"Failed to convert scikit-learn model to ONNX: {str(e)}")
 
     @require_api_key
-    def train_model(
-        self, 
-        model_name: str, 
-        training_dataset: str, 
-        task_type: str, 
-        model_category: str, 
-        architecture: str, 
-        region: str, 
-        hyperparameters: dict = None
+    async def train_model(
+        self,
+        model_name: str,
+        task_type: str,
+        model_category: str,
+        architecture: str,
+        hyperparameters: dict = None,
+        aoi: str = None,
+        expression_x: str = None,
+        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,
+        server: str = None,
+        extra_filters: list[str] = None,
+        extra_filters_rate: list[float] = None,
+        extra_filters_res: list[float] = None
     ) -> dict:
         """
         Train a model using the external model training API.
@@ -1107,7 +1154,6 @@ class ModelManagement:
             task_type (str): The type of ML task (e.g., regression, classification).
             model_category (str): The category of model (e.g., random_forest).
             architecture (str): The model architecture.
-            region (str): The region identifier.
             hyperparameters (dict, optional): Additional hyperparameters for training.
             
         Returns:
@@ -1116,13 +1162,71 @@ class ModelManagement:
         Raises:
             APIError: If the API request fails
         """
+        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._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=model_name,
+            bucket="terrakio-mass-requests",
+            collection_type="basic"
+        )
+
         payload = {
             "model_name": model_name,
-            "training_dataset": training_dataset,
             "task_type": task_type,
             "model_category": model_category,
             "architecture": architecture,
-            "region": region,
-            "hyperparameters": hyperparameters
+            "hyperparameters": hyperparameters,
+            "expressions": expressions,
+            "filters": filters,
+            "aoi": aoi_data,
+            "samples": samples,
+            "year_range": [start_year, end_year],
+            "crs": crs,
+            "tile_size": tile_size,
+            "res": res,
+            "server": server
         }
-        return self._client._terrakio_request("POST", "/train_model", json=payload)
+
+        task_id_dict, _ = await self._client._terrakio_request("POST", "models/train", json=payload)
+
+        await self._client.mass_stats.track_progress(task_id_dict["task_id"])

terrakio_core/endpoints/user_management.py

@@ -1,12 +1,13 @@
 from typing import Dict, Any, List, Optional
 from ..helper.decorators import require_token, require_api_key, require_auth
+from ..exceptions import UserNotFoundError, GetUserByIdError, GetUserByEmailError, ListUsersError, EditUserError, ResetQuotaError, DeleteUserError, GetUsersByRoleError, RoleDoNotExistError, ChangeRoleError
 
 class UserManagement:
     def __init__(self, client):
         self._client = client
 
     @require_api_key
-    def get_user_by_id(self, id: str) -> Dict[str, Any]:
+    async def get_user_by_id(self, id: str) -> Dict[str, Any]:
         """
         Get user by ID.
 
@@ -17,12 +18,19 @@ class UserManagement:
             User information
             
         Raises:
-            APIError: If the API request fails
+            GetUserByIdError: If the API request fails
+            UserNotFoundError: If the user is not found
         """
-        return self._client._terrakio_request("GET", f"admin/users/{id}")
-    
+        response, status = await self._client._terrakio_request("GET", f"admin/users/{id}")
+        if status != 200:
+            if status == 404:
+                raise UserNotFoundError(f"User {id} not found.", status_code = status)
+            raise GetUserByIdError(f"Get user by id failed with status {status}", status_code = status)
+        else:
+            return response
+
     @require_api_key
-    def get_user_by_email(self, email: str) -> Dict[str, Any]:
+    async def get_user_by_email(self, email: str) -> Dict[str, Any]:
         """
         Get user by email.
 
@@ -33,12 +41,19 @@ class UserManagement:
             User information
             
         Raises:
-            APIError: If the API request fails
+            GetUserByEmailError: If the API request fails
+            UserNotFoundError: If the user is not found
         """
-        return self._client._terrakio_request("GET", f"admin/users/email/{email}")
+        response, status = await self._client._terrakio_request("GET", f"admin/users/email/{email}")
+        if status != 200:
+            if status == 404:
+                raise UserNotFoundError(f"User {email} not found.", status_code = status)
+            raise GetUserByEmailError(f"Get user by email failed with status {status}", status_code = status)
+        else:
+            return response
     
     @require_api_key
-    def list_users(self, substring: Optional[str] = None, uid: bool = False) -> List[Dict[str, Any]]:
+    async def list_users(self, substring: Optional[str] = None, uid: bool = False) -> List[Dict[str, Any]]:
         """
         List users, optionally filtering by a substring.
         
@@ -50,15 +65,19 @@ class UserManagement:
             List of users
             
         Raises:
-            APIError: If the API request fails
+            ListUsersError: If the API request fails
         """
         params = {"uid": str(uid).lower()}
         if substring:
             params['substring'] = substring
-        return self._client._terrakio_request("GET", "admin/users", params=params)
+        response, status = await self._client._terrakio_request("GET", "admin/users", params=params)
+        if status != 200:
+            raise ListUsersError(f"List users failed with status {status}", status_code = status)
+        else:
+            return response
     
     @require_api_key
-    def edit_user(
+    async def edit_user(
         self,
         uid: str,
         email: Optional[str] = None,
@@ -82,7 +101,7 @@ class UserManagement:
             Updated user information
             
         Raises:
-            APIError: If the API request fails
+            EditUserError: If the API request fails
         """
         payload = {"uid": uid}
         payload_mapping = {
@@ -95,10 +114,14 @@ class UserManagement:
         for key, value in payload_mapping.items():
             if value is not None:
                 payload[key] = value
-        return self._client._terrakio_request("PATCH", "admin/users", json=payload)
+        response, status = await self._client._terrakio_request("PATCH", "admin/users", json=payload)
+        if status != 200:
+            raise EditUserError(f"Edit user failed with status {status}", status_code = status)
+        else:
+            return response
     
     @require_api_key
-    def reset_quota(self, email: str, quota: Optional[int] = None) -> Dict[str, Any]:
+    async def reset_quota(self, email: str, quota: Optional[int] = None) -> Dict[str, Any]:
         """
         Reset the quota for a user by email.
         
@@ -106,16 +129,20 @@ class UserManagement:
             email: The user's email (required)
             quota: The new quota value (optional)
             
-        Returns:
-            API response as a dictionary
+        Raises:
+            ResetQuotaError: If the API request fails
         """
         payload = {"email": email}
         if quota is not None:
             payload["quota"] = quota
-        return self._client._terrakio_request("PATCH", f"admin/users/reset_quota/{email}", json=payload)
+        response, status = await self._client._terrakio_request("PATCH", f"admin/users/reset_quota/{email}", json=payload)
+        if status != 200:
+            raise ResetQuotaError(f"Reset quota failed with status {status}", status_code = status)
+        else:
+            return response
     
     @require_api_key
-    def delete_user(self, uid: str) -> Dict[str, Any]:
+    async def delete_user(self, uid: str) -> Dict[str, Any]:
         """
         Delete a user by UID.
 
@@ -123,9 +150,67 @@ class UserManagement:
             uid: The user's UID (required)
             
         Returns:
-            API response as a dictionary
+            Deleted user information
+            
+        Raises:
+            DeleteUserError: If the API request fails
+        """
+        response, status = await self._client._terrakio_request("DELETE", f"admin/users/{uid}")
+        if status != 200:
+            raise DeleteUserError(f"Delete user failed with status {status}", status_code = status)
+        else:
+            return response
+    
+    @require_api_key
+    async def get_users_by_role(self, role: str) -> Dict[str, Any]:
+        """
+        Get users by role.
+
+        Args:
+            role: The user role to filter by (required)
+            
+        Returns:
+            Users with the specified role
+            
+        Raises:
+            GetUsersByRoleError: If the API request fails
+        """
+        response, status = await self._client._terrakio_request("GET", f"admin/users/role?role={role}")
+        if status != 200:
+            if status == 422:
+                raise RoleDoNotExistError(f"Role {role} does not exist", status_code = status)
+            raise GetUsersByRoleError(f"Get users by role failed with status {status}", status_code = status)
+        else:
+            return response
+    
+    @require_api_key
+    async def change_role(self, uid: str, role: str, reset_quota: Optional[bool] = None, limit: Optional[int] = None) -> Dict[str, Any]:
+        """
+        Change user role.
+
+        Args:
+            uid: The user's UID to change role for (required)
+            role: Role to apply (required)
+            reset_quota: Reset user's quota to new role limit (optional)
+            limit: Quota limit if role is custom (optional)
+            
+        Returns:
+            Response from the role change operation
             
         Raises:
-            APIError: If the API request fails
+            ChangeRoleError: If the API request fails
         """
-        return self._client._terrakio_request("DELETE", f"admin/users/{uid}")
+        payload = {"uid": uid, "role": role}
+        if reset_quota is not None:
+            payload["reset_quota"] = reset_quota
+        if limit is not None:
+            payload["limit"] = limit
+        response, status = await self._client._terrakio_request("PATCH", "admin/users/change_role", json=payload)
+        if status != 200:
+            if status == 404:
+                raise UserNotFoundError(f"User {uid} not found", status_code = status)
+            elif status == 422:
+                raise RoleDoNotExistError(f"Role {role} does not exist", status_code = status)
+            raise ChangeRoleError(f"Change role failed with status {status}", status_code = status)
+        else:
+            return response