atlasai-dstoolkit-client 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

atlasai/toolkit/__init__.py

@@ -33,5 +33,5 @@ del osp
 del os
 
 
-from . import fabric, feature, dataset, model
+from . import fabric, feature, dataset, model, workflows
 from .login import login, logout

atlasai/toolkit/feature.py

@@ -15,6 +15,7 @@
 import logging
 import uuid
 
+import geopandas as gpd
 import pandas as pd
 from shapely import wkb
 
@@ -60,7 +61,7 @@ class FeatureExport:
         self._export = self._details()
         return self.export
 
-    def results(self, limit=None) -> pd.DataFrame:
+    def results(self, limit=None, as_gdf=False) -> pd.DataFrame:
         if not self.export:
             return pd.DataFrame([])
 
@@ -84,6 +85,8 @@ class FeatureExport:
 
         try:
             df['shape'] = wkb.loads(df['shape'])
+            if as_gdf:
+                df = gpd.GeoDataFrame(df, geometry="shape")
         except Exception:
             logger.warning('Shape field is not in WKB format. Skipping conversion.')
 

atlasai/toolkit/output.py

@@ -0,0 +1,18 @@
+from rich.console import Console
+from rich.text import Text
+
+console = Console()
+
+def print_title(text):
+    text = Text(text, style="bold underline", justify="center")
+    console.print(text)
+
+def print_subtitle(text):
+    text = Text(text, style="italic cyan", justify="center")
+    console.print(text)
+
+def print_body(text):
+    if isinstance(text, bytes):
+        text = text.decode("utf-8")
+    text = Text(text, style="white", justify="left")
+    console.print(text)

atlasai/toolkit/workflows.py

@@ -0,0 +1,229 @@
+# Copyright 2025 AtlasAI PBC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import os
+import time
+
+import pandas as pd
+from tqdm import tqdm
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+
+import requests
+
+from . import api, output as o
+
+logger = logging.getLogger('atlasai.toolkit')
+
+@dataclass()
+class WorkflowResponse:
+    id: str
+    uid: str
+    name: str
+    create_date: str
+    update_date: str
+    status: str
+    message: str
+
+
+class Workflow:
+    type = None
+    id = None
+
+    def __init__(self, id=None):
+        # Allow the selection of an existing workflow and check the logs.
+        self.id = id
+
+    @property
+    def status(self):
+        _result = self.get()
+        return _result['data']['status']
+
+    def results(self, *args, **kwargs):
+        return self._results(*args, **kwargs)
+
+    def _set(self, result):
+        self.id = result['data']['id']
+
+    def _config(self):
+        raise NotImplementedError()
+
+    def _validate(self):
+        raise NotImplementedError()
+
+    def _results(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get(self):
+        if not self.id:
+            raise Exception('No valid execution of the workflow. Please use the `run` method first.')
+
+        _, result = api._get(resource=f'workflow/{self.id}')
+        return result
+
+    def logs(self, only_errors=True):
+        if not self.id:
+            raise Exception('No valid execution of the workflow. Please use the `run` method first.')
+
+        workflow = self.get()
+
+        _, result = api._get(resource=f'workflow/{self.id}/logs', params={'only_errors': only_errors})
+
+        o.print_title(f'Workflow: {workflow["data"]["name"]}')
+        o.print_subtitle(f'Status: {workflow["data"]["status"]}')
+        o.print_subtitle(f'Message: {workflow["data"]["message"] or "-"}')
+
+        o.print_body('---------------------------------------------------')
+
+        if not result['data']:
+            if only_errors:
+                o.print_title('No errors detected.')
+            else:
+                o.print_title('No pod found to pull logs from.')
+
+        for pod, data in result['data'].items():
+            o.print_title(f'Pod: {pod}')
+            o.print_subtitle(f'Message: {data["message"] or "-"}')
+            o.print_subtitle(f'Logs: {data["logs"] or "-"}')
+            o.print_body('---------------------------------------------------')
+
+    def run(self, wait_until_complete=False):
+        self._validate()
+        status = 'Running'
+        message = ''
+        data = dict(
+            type=self.type,
+            config=self._config()
+        )
+        _, result = api._post(resource='workflow', data=data)
+        self._set(result)
+        if wait_until_complete:
+            while True:
+                _result = self.get()
+                if _result['data']['status'].lower() != 'running':
+                    status = _result['data']['status']
+                    message = _result['data']['message']
+                    break
+
+                time.sleep(5)
+        return WorkflowResponse(**result['data'], status=status, message=message)
+
+class Test(Workflow):
+    type = 'test'
+
+    def __repr__(self):
+        return 'Test'
+
+    def __str__(self):
+        return 'Test'
+
+    def configure(self):
+        pass
+
+    def _config(self):
+        return dict()
+
+    def _validate(self):
+        pass
+
+
+class Electrification(Workflow):
+    type = 'electrification'
+
+    def __init__(self, id=None):
+        self.aoi_name = None
+        self.aoi_path = None
+        self.aoi = None
+        self.output_path = None
+
+        super().__init__(id)
+
+    def __repr__(self):
+        return f'Electrification({self.aoi_name})'
+
+    def __str__(self):
+        return f'Electrification({self.aoi_name})'
+
+    def configure(self, aoi_name=None, aoi_path=None, aoi=None, output_path=None):
+        self.aoi_name = aoi_name or self.aoi_name
+        self.aoi_path = aoi_path or self.aoi_path
+        self.aoi = aoi or self.aoi
+        self.output_path = output_path or self.output_path
+
+    def _config(self):
+        cfg = {}
+        for name, value in (
+                ('aoi_name', self.aoi_name),
+                ('aoi_geojson', self.aoi),
+                ('aoi_geojson_uri', self.aoi_path),
+                ('output_bucket', self.output_path)
+        ):
+            if value:
+                cfg[name] = value
+        return cfg
+
+    def _validate(self):
+        if not self.aoi_name:
+            raise Exception('`aoi_name` must be specified.')
+
+        if not self.aoi_path and not self.aoi:
+            raise Exception('`aoi_path` or `aoi` must be specified.')
+
+    def _results(self, save_to=None):
+        results = []
+        if not self.id:
+            raise Exception('No valid execution of the workflow. Please use the `run` method first.')
+
+        workflow = self.get()
+        if not workflow['data']['status'] == 'Succeeded':
+            raise Exception(f'Workflow not in `Succeeded` state. Current state: {workflow["data"]["status"]}')
+
+        _, result = api._get(resource=f'workflow/{self.id}/results')
+
+        with ThreadPoolExecutor(max_workers=8) as ex, tqdm(total=len(result['data']), desc="Downloading", unit="file") as pbar:
+            fut_map = {ex.submit(self._download_file, res["name"], res["url"], save_to): i
+                       for i, res in enumerate(result['data'])}
+            for fut in as_completed(fut_map):
+                results.append(fut.result())
+                pbar.update(1)
+
+        return results
+
+    def _download_file(self, name, url, to=None):
+        local_path = os.path.join(to or os.getcwd(), name.split('/')[-1])
+
+        logger.debug(f'Downloading :{name}')
+
+        with requests.get(url, stream=True, timeout=120) as r:
+
+            logger.debug(f'Finished downloading: {name}')
+            r.raise_for_status()
+            with open(local_path, "wb") as f:
+                for chunk in r.iter_content(chunk_size=1024 * 1024):
+                    if chunk:
+                        f.write(chunk)
+
+            logger.debug(f'Finished saving: {name}')
+
+        return local_path
+
+
+def List(search=None, offset=0, limit=100):
+    pd.set_option("display.max_colwidth", None)
+    resource = 'workflows'
+    _, data = api._list(resource=resource, params={'search': search, 'offset': offset, 'limit': limit})
+    df = pd.DataFrame(data['data'])
+    df = df.drop(columns=['update_date', 'uid'])
+    return df

{atlasai_dstoolkit_client-0.0.12.dist-info → atlasai_dstoolkit_client-0.0.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlasai-dstoolkit-client
-Version: 0.0.12
+Version: 0.0.14
 Summary: UNKNOWN
 Home-page: UNKNOWN
 Author: AtlasAI SWE
@@ -12,12 +12,14 @@ License-File: LICENSE.txt
 Requires-Dist: arrow (<2.0.0,>=1.3.0)
 Requires-Dist: atlasai-mlhub-client
 Requires-Dist: furl (<3.0.0,>=2.1.2)
+Requires-Dist: geopandas (~=1.1.1)
 Requires-Dist: pandas (<3.0.0,>=2.2.3)
 Requires-Dist: pyarrow (>=19.0.1)
 Requires-Dist: python-dotenv (<2.0.0,>=1.0.1)
 Requires-Dist: requests (<=3.0.0,>=2.32.3)
 Requires-Dist: semver (<4.0.0,>=3.0.2)
 Requires-Dist: shapely (~=2.1.0)
+Requires-Dist: tqdm (~=4.66.5)
 Provides-Extra: dev
 Requires-Dist: coverage ; extra == 'dev'
 Requires-Dist: flake8 ; extra == 'dev'

{atlasai_dstoolkit_client-0.0.12.dist-info → atlasai_dstoolkit_client-0.0.14.dist-info}/RECORD

@@ -1,16 +1,18 @@
-atlasai/toolkit/__init__.py,sha256=rjt3GfV7vA3RFMIzqPwBebJUXdQbeTLmEYmm7SOoOGw,1002
+atlasai/toolkit/__init__.py,sha256=0d2EYTgFf-RH2z1M83gGfZdrt_PBXbJk6o5ZVAKCPtA,1013
 atlasai/toolkit/api.py,sha256=BvO-gLRmbmkKduwbbADjcLlIkS9blzfM_cbMR4DhQmU,5269
 atlasai/toolkit/constants.py,sha256=Jxozn9tKOvAxyOYOZ6bzFtI9m1YETJF1GURzTl9NNC8,422
 atlasai/toolkit/dataset.py,sha256=fhzhxF9YMzIwEpaJZEPOK2SuLMJhVGI75eahxnH_T2c,3254
 atlasai/toolkit/fabric.py,sha256=6aFR2PGQc9P3Qa07WdBg9eKoUzU8n2y_-gGjYcyMrWY,1921
-atlasai/toolkit/feature.py,sha256=VKwX_yAwwMo1WuLZ_RbKVJEH4L1PAoYEPzWa0lH9ats,2828
+atlasai/toolkit/feature.py,sha256=Q5M9zOGafynYuKaELL1kZemYPfKAZh84TgH7jw9J3ZU,2949
 atlasai/toolkit/init.py,sha256=JkdJ6QGdYWrq65jgz2pn5RYXUeUe2Ez88_-eMf5CNi0,1100
 atlasai/toolkit/login.py,sha256=n4ydfo9qCsmbZq6er1xeljBD76vdTJGjbhYHMmOyDbQ,3061
 atlasai/toolkit/model.py,sha256=RUe0HbDpzvHOV9A4rzG3PgN9boMWDHQ2tR7IKHXzbx8,4126
+atlasai/toolkit/output.py,sha256=FyDjrpVlbrEyfHfwOpxp8H57jx_qXahDjO1qpHIeuYM,473
 atlasai/toolkit/requests.py,sha256=X86nIo07hAjUlilZcZ1lV8RB7KOsTKbTGtcY_SpFEXY,1223
 atlasai/toolkit/utils.py,sha256=lYh3P2XOshRgHCjFeXJ0FOJWQW64sddgx8c2kL6Wqwc,1566
-atlasai_dstoolkit_client-0.0.12.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-atlasai_dstoolkit_client-0.0.12.dist-info/METADATA,sha256=2vdA8hWTJEIvSm1OgPRt9wmsTtZO_5oR_znxLDNuDbI,1405
-atlasai_dstoolkit_client-0.0.12.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-atlasai_dstoolkit_client-0.0.12.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
-atlasai_dstoolkit_client-0.0.12.dist-info/RECORD,,
+atlasai/toolkit/workflows.py,sha256=HlmLNma1P4xyd74gzM5moBXPGjVQqziGOtyRTXBXUhw,6994
+atlasai_dstoolkit_client-0.0.14.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+atlasai_dstoolkit_client-0.0.14.dist-info/METADATA,sha256=vWbApYiDX_XA2mFvwXKRWDCV6KW5OwOYL-jqJN4S56A,1471
+atlasai_dstoolkit_client-0.0.14.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+atlasai_dstoolkit_client-0.0.14.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
+atlasai_dstoolkit_client-0.0.14.dist-info/RECORD,,