atlasai-dstoolkit-client 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

atlasai/toolkit/__init__.py

@@ -0,0 +1,37 @@
+# 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.
+
+# flake8: noqa
+import os
+import os.path as osp
+
+from dotenv import load_dotenv
+
+# do this first!
+from .init import configure_logging
+configure_logging()
+del configure_logging
+
+env_file = osp.join(osp.dirname(__file__), '.env')
+if osp.exists(env_file):
+    load_dotenv(env_file)
+
+del load_dotenv
+del env_file
+del osp
+del os
+
+
+from . import fabric, feature
+from .login import login, logout

atlasai/toolkit/api.py

@@ -0,0 +1,170 @@
+import os
+from http import HTTPStatus
+import json
+import logging
+
+from furl import furl
+
+from .constants import DS_TOOLKIT_URL
+from .requests import get_session
+
+logger = logging.getLogger(__name__)
+
+
+def _get_headers(access_token):
+    return {
+        'Authorization': f'Bearer {access_token}',
+        'Content-Type': 'application/json'
+    }
+
+def _get_access_token():
+    env_name = os.getenv('ATLASAI_TOKEN_NAME', 'ATLASAI_TOKEN')
+    access_token = os.getenv(env_name)
+    if not access_token:
+        raise Exception('No access token found. Call the `login()` method first!')
+    return access_token
+
+def _add_params(f, params=None):
+    if params is None:
+        params = {}
+
+    for k, v in params.items():
+        if v:
+            f.args[k] = v
+    return f
+
+def _process_data(data):
+    if not data:
+        return '{}'
+    return json.dumps({k: v for k, v in data.items() if v is not None})
+
+def _paginate(method, url, access_token=None, params=None):
+    if access_token is None:
+        access_token = _get_access_token()
+    results = []
+    if params is None:
+        params = {}
+    if params.get('limit') is None:
+        params['limit'] = 100
+
+    if params.get('offset') is None:
+        params['offset'] = 0
+
+    def get_results(_url, _params):
+        f = furl(_url)
+        f = _add_params(f, _params)
+        _response = session.request(method, f.url, headers=_get_headers(access_token))
+        _response.raise_for_status()
+        return _response
+
+    session = get_session()
+    while True:
+        response = get_results(url, params)
+        data = response.json()
+        if not data:
+            break
+
+        results.extend(data)
+
+        if len(data) < params['limit']:
+            break
+        params['offset'] = params['offset'] + params['limit']
+    return results
+
+def _get(access_token=None, resource=None, _id=None, method='get', params=None):
+    if access_token is None:
+        access_token = _get_access_token()
+    session = get_session()
+
+    f = furl(DS_TOOLKIT_URL)
+    f.path = f'api/{resource}/{_id}' if _id else f'api/{resource}'
+    f = _add_params(f, params)
+    url = f.url
+
+    response = session.request(method, url, headers=_get_headers(access_token))
+    response.raise_for_status()
+    return response.status_code, response.json()
+
+
+def _list(access_token=None, resource=None, method='get', params=None, is_paginated=True):
+    if access_token is None:
+        access_token = _get_access_token()
+    session = get_session()
+
+    f = furl(DS_TOOLKIT_URL)
+    f.path = f'api/{resource}'
+    url = f.url
+
+    # return all the records if limit not specified
+    if is_paginated and not params.get('limit') and not params.get('offset'):
+        return 200, _paginate(method, url, access_token, params)
+
+    f = _add_params(f, params)
+    url = f.url
+    response = session.request(method, url, headers=_get_headers(access_token))
+    response.raise_for_status()
+    return response.status_code, response.json()
+
+def _post(access_token=None, resource=None, method='post', data=None, params=None, url=None):
+    if access_token is None:
+        access_token = _get_access_token()
+    session = get_session()
+
+    f = furl(url or DS_TOOLKIT_URL)
+    f.path = f'api/{resource}'
+    f = _add_params(f, params)
+
+    url = f.url
+    data = _process_data(data)
+
+    response = session.request(method, url, headers=_get_headers(access_token), data=data)
+    response.raise_for_status()
+    if response.status_code == HTTPStatus.NO_CONTENT:
+        return response.status_code, None
+    elif response.status_code == HTTPStatus.ACCEPTED:
+        return response.status_code, response.headers['Location']
+    else:
+        return response.status_code, response.json()
+
+def _patch(access_token=None, resource=None, method='patch', data=None, params=None):
+    if access_token is None:
+        access_token = _get_access_token()
+    session = get_session()
+
+    f = furl(DS_TOOLKIT_URL)
+    f.path = f'api/{resource}'
+    f = _add_params(f, params)
+
+    url = f.url
+    data = _process_data(data)
+
+    response = session.request(method, url, headers=_get_headers(access_token), data=data)
+    response.raise_for_status()
+    if response.status_code == HTTPStatus.NO_CONTENT:
+        return response.status_code, None
+    elif response.status_code == HTTPStatus.ACCEPTED:
+        return response.status_code, response.headers['Location']
+    else:
+        return response.status_code, response.json()
+
+def _delete(access_token=None, resource=None, method='delete', data=None, params=None):
+    if access_token is None:
+        access_token = _get_access_token()
+    session = get_session()
+
+    f = furl(DS_TOOLKIT_URL)
+    f.path = f'api/{resource}'
+    f = _add_params(f, params)
+
+    url = f.url
+    if data:
+        response = session.request(method, url, headers=_get_headers(access_token), data=_process_data(data))
+    else:
+        response = session.request(method, url, headers=_get_headers(access_token))
+    response.raise_for_status()
+    if response.status_code == HTTPStatus.NO_CONTENT:
+        return response.status_code, None
+    elif response.status_code == HTTPStatus.ACCEPTED:
+        return response.status_code, response.headers['Location']
+    else:
+        return response.status_code, response.json()

atlasai/toolkit/constants.py

@@ -0,0 +1,10 @@
+import os
+
+VINZ_URL = os.getenv('VINZ_URL') or 'https://vinz.atlasai.co'
+DS_TOOLKIT_URL = os.getenv('DS_TOOLKIT_URL') or 'https://dstoolkit.atlasai.co'
+
+DEFAULT_PAGE_SIZE = 20
+DISABLE_SSL_VERIFICATION = 'DISABLE_SSL_VERIFICATION'
+
+TOKEN_ENV_VAR = 'VINZ_ENCRYPTED_TOKEN'
+TOKEN_TIMESTAMP_ENV_VAR = 'VINZ_ENCRYPTED_TOKEN_TIMESTAMP'

atlasai/toolkit/fabric.py

@@ -0,0 +1,67 @@
+# 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 uuid
+import arrow
+
+from . import api, constants, utils
+
+
+def search(*args, **kwargs):
+    fs = FabricSearch(*args, **kwargs)
+    utils.show_page(fs.page)
+    return fs
+
+class FabricSearch:
+
+    def __init__(self, id=None, *args, **kwargs):
+        self.id = id or uuid.uuid4()
+        self._search = None
+
+    def __repr__(self):
+        return f'FabricSearch({self.id})'
+
+    def __str__(self):
+        return f'FabricSearch({self.id})'
+
+    @property
+    def page(self):
+        return f'{constants.DS_TOOLKIT_URL}/fabric/search/{self.id}'
+
+    @property
+    def search(self):
+        return self._search
+
+    def refresh(self):
+        if self._search is None:
+            self._search = self._refresh()
+        else:
+            new_version = self._info()
+            # new updates ? pull them
+            if arrow.get(new_version['update_date']) > arrow.get(self._search['update_date']):
+                self._search = self._refresh()
+        return self._search
+
+    def info(self):
+        return self._info()
+
+    def _info(self):
+        resource = f'/fabric/search/{self.id}/info'
+        _, data = api._get(resource=resource)
+        return data['data']
+
+    def _refresh(self):
+        resource = f'/fabric/search/{self.id}/select'
+        _, data = api._get(resource=resource)
+        return data['data']

atlasai/toolkit/feature.py

@@ -0,0 +1,75 @@
+# 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 uuid
+import pandas as pd
+
+from . import api, constants, utils
+
+
+def export(*args, **kwargs):
+    fe = FeatureExport(*args, **kwargs)
+    utils.show_page(fe.page)
+    return fe
+
+class FeatureExport:
+    def __init__(self, search, id=None, *args, **kwargs):
+        self.id = id or uuid.uuid4()
+        self._search = search
+        self._export = None
+
+    def __repr__(self):
+        return f'FeatureExport({self._search.id})'
+
+    def __str__(self):
+        return f'FeatureExport({self._search.id})'
+
+    @property
+    def page(self):
+        return f'{constants.DS_TOOLKIT_URL}/feature/export/{self.id}?feature_search_id={self._search.id}'
+
+    @property
+    def export(self):
+        return self._export
+
+    @property
+    def search(self):
+        return self._search.search
+
+    def details(self):
+        if self.export is None or self.export.status not in ['completed', 'failed']:
+            self._export = self._details()
+        return self.export
+
+    def refresh(self):
+        self._export = self._details()
+        return self.export
+
+    def results(self, limit=None) -> pd.DataFrame:
+        if self.export.status not in ['completed', 'failed']:
+            raise Exception(f'Export state is: {self.export.status}')
+
+        path = getattr(self.export, 'url_path', getattr(self.export, 'output_path'))
+        if not path:
+            raise Exception('Path not found')
+
+        df = pd.read_parquet(path, engine='pyarrow')
+        if limit:
+            df = df.head(limit)
+        return df
+
+    def _details(self):
+        resource = f'/feature/export/{self._search.id}/details'
+        _, data = api._get(resource=resource)
+        return data['data']

atlasai/toolkit/init.py

@@ -0,0 +1,30 @@
+# 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
+
+def configure_logging():
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter(
+        '%(asctime)s %(levelname)-5.5s [%(name)s][pid:%(process)s tid:%(thread)s] %(message)s'
+    )
+    handler.setFormatter(formatter)
+
+    log_level = (os.getenv('DSTOOLKIT_LOG_LEVEL') or 'INFO').upper()
+    log_level = getattr(logging, log_level, logging.INFO)
+
+    logger = logging.getLogger('atlasai.toolkit')
+    logger.setLevel(log_level)
+    logger.addHandler(handler)

atlasai/toolkit/login.py

@@ -0,0 +1,107 @@
+import hmac
+import logging
+import os
+from urllib.parse import urlparse
+import warnings
+
+import arrow
+from furl import furl
+import requests
+
+from . import constants
+from .utils import clean_token_env_vars
+
+
+logger = logging.getLogger(__name__)
+
+
+def login(env_name='ATLASAI_TOKEN', return_access_token=False):
+    """
+     Authenticate with Vinz
+
+     Returns an OAuth2 Access Token
+
+     If `env_name` provided, the Access Token will be saved
+     to the named environment variable
+
+     #### Usage
+
+     ```python
+     from atlasai.vinz import client
+
+     token = client.authenticate(<OPTIONAL_ENV_VARIABLE_NAME>)
+     ```
+     """
+    os.environ['ATLASAI_TOKEN_NAME'] = env_name
+    f = furl(constants.VINZ_URL)
+    f.path = 'api/token'
+    url = f.url
+    headers = {}
+    include_authorization(url, headers)
+
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    data = response.json()
+    token = data['access_token']
+
+    os.environ[env_name] = token
+
+    user_id = data.get('email') or data.get('sub') or 'AtlasAI Employee'
+    os.environ['LOGNAME'] = user_id
+
+    if return_access_token:
+        return token
+
+def logout():
+    os.environ.pop(os.getenv('ATLASAI_TOKEN_NAME', ''), None)
+    clean_token_env_vars()
+
+def load_credentials(access_key=None, secret_key=None):
+    access_key = access_key or os.getenv('ATLASAI_ACCESS_KEY')
+    secret_key = secret_key or os.getenv('ATLASAI_SECRET_KEY')
+
+    return access_key, secret_key
+
+
+def include_authorization(url, headers, bearer_token=None, access_key=None, secret_key=None):
+    bearer_token = bearer_token or os.getenv('ATLASAI_BEARER_TOKEN')
+    access_key, secret_key = load_credentials(
+        access_key=access_key,
+        secret_key=secret_key,
+    )
+
+    if bearer_token:
+        headers['Authorization'] = f'Bearer {bearer_token}'
+        return
+
+    if not access_key or not secret_key:
+        warnings.warn('No API Keys provided to access Vinz API. Provide the following pair ATLASAI_ACCESS_KEY and ATLASAI_SECRET_KEY')
+        raise ValueError('ATLASAI_ACCESS_KEY and ATLASAI_SECRET_KEY must be provided together')
+
+    product, version = 'atlasai', '1'
+    headers.update({
+        'Host': urlparse(url).netloc,
+        'X-AtlasAI-Date': arrow.utcnow().isoformat(),
+        'X-AtlasAI-Credential': '/'.join([product, version, access_key]),
+        'X-AtlasAI-SignedHeaders': 'x-atlasai-date;x-atlasai-credential;host',
+    })
+
+    sign_request(headers, secret_key)
+
+
+def sign_request(headers, secret_key):
+    product, version, access_key = headers['X-AtlasAI-Credential'].split('/')
+    key = f'{product}{version}{secret_key}'.encode('utf-8')
+    for msg in (
+        headers['X-AtlasAI-Date'],
+        f'{product}_{version}_request',
+    ):
+        obj = hmac.new(key, msg.encode('utf-8'), 'sha256')
+        key = obj.digest()
+
+    msg = '\n'.join([
+        headers['X-AtlasAI-Date'],
+        headers['X-AtlasAI-Credential'],
+        headers['Host']
+    ])
+    headers['X-AtlasAI-Signature'] = hmac.new(key, msg.encode('utf-8'), 'sha256').hexdigest()

atlasai/toolkit/requests.py

@@ -0,0 +1,52 @@
+import os
+
+import requests
+from requests.adapters import HTTPAdapter, Retry
+from .constants import DISABLE_SSL_VERIFICATION
+
+
+STATUS_FORCELIST = tuple([429, 500, 502, 503, 504])
+
+def mount_retry(
+    session,
+    total=10,
+    backoff_factor=0.2,
+    allowed_methods=None,
+    status_forcelist=STATUS_FORCELIST,
+):
+    """
+    Attach retry handlers to HTTP and HTTPS endpoints of a Requests Session
+    """
+
+    retries = Retry(
+        total=total,
+        backoff_factor=backoff_factor,
+        allowed_methods=allowed_methods,
+        status_forcelist=status_forcelist,
+    )
+
+    session.mount('http://', HTTPAdapter(max_retries=retries))
+    session.mount('https://', HTTPAdapter(max_retries=retries))
+
+def get_session(
+    total=3,
+    backoff_factor=0.2,
+    allowed_methods=None,
+    status_forcelist=STATUS_FORCELIST,
+):
+    """
+    Get a Requests Session with retry handlers for HTTP and HTTPS endpoints
+    """
+
+    sess = requests.Session()
+    if os.getenv(DISABLE_SSL_VERIFICATION):
+        sess.verify = False
+    mount_retry(
+        sess,
+        total=total,
+        backoff_factor=backoff_factor,
+        allowed_methods=allowed_methods,
+        status_forcelist=status_forcelist,
+    )
+
+    return sess

atlasai/toolkit/utils.py

@@ -0,0 +1,58 @@
+import arrow
+import os
+
+from furl import furl
+import webbrowser
+
+from . import api, constants
+
+def is_notebook() -> bool:
+    try:
+        from IPython import get_ipython
+        shell = get_ipython().__class__.__name__
+        if shell == 'ZMQInteractiveShell':
+            return True
+        elif shell == 'TerminalInteractiveShell':
+            return False
+        else:
+            return False
+    except Exception:
+        return False
+
+
+def show_page(page):
+    f = furl(page)
+
+    if is_first_page():
+        token = retrieve_one_time_token()
+        set_token_env_vars(token)
+        f.args['token'] = token
+
+    if is_notebook():
+        from IPython.core.display import display
+        from IPython.display import IFrame
+        display(IFrame(f.url, width=800, height=600))
+    else:
+        webbrowser.open(f.url)
+
+def is_first_page():
+    if os.getenv(constants.TOKEN_ENV_VAR) is None:
+        return True
+    last_set = os.getenv(constants.TOKEN_TIMESTAMP_ENV_VAR)
+    if last_set:
+        last_set = arrow.get(last_set)
+        # if had passed more than 24h since last iframe. pass the token again in case it needs it
+        if arrow.utcnow() >= last_set.shift(hours=+24):
+            return True
+
+def set_token_env_vars(token):
+    os.environ[constants.TOKEN_ENV_VAR] = token
+    os.environ[constants.TOKEN_TIMESTAMP_ENV_VAR] = arrow.now().isoformat()
+
+def clean_token_env_vars():
+    os.environ.pop(constants.TOKEN_ENV_VAR, None)
+    os.environ.pop(constants.TOKEN_TIMESTAMP_ENV_VAR, None)
+
+def retrieve_one_time_token():
+    _, data = api._post(resource='token/wrap', url=constants.VINZ_URL)
+    return data['token']

{atlasai_dstoolkit_client-0.0.1.dist-info → atlasai_dstoolkit_client-0.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlasai-dstoolkit-client
-Version: 0.0.1
+Version: 0.0.2
 Summary: UNKNOWN
 Home-page: UNKNOWN
 Author: AtlasAI SWE

atlasai_dstoolkit_client-0.0.2.dist-info/RECORD

@@ -0,0 +1,14 @@
+atlasai/toolkit/__init__.py,sha256=szu5LHZ41ccWztCI2LvQ3QnZRybosgTbU3yd6tS7cUw,986
+atlasai/toolkit/api.py,sha256=BvO-gLRmbmkKduwbbADjcLlIkS9blzfM_cbMR4DhQmU,5269
+atlasai/toolkit/constants.py,sha256=sE0PeFa9_htCPVFADHbkyPIz3QOai0tAIA5IiiZI8wA,329
+atlasai/toolkit/fabric.py,sha256=vjc8PJ30bUZ9z-ZqovpfqDkOo_OkMH80GMX8YoOPdfI,1923
+atlasai/toolkit/feature.py,sha256=ijoljwdtlI2djeALY14ctL97K5wbzpJkRMtwq-y-PIg,2276
+atlasai/toolkit/init.py,sha256=JkdJ6QGdYWrq65jgz2pn5RYXUeUe2Ez88_-eMf5CNi0,1100
+atlasai/toolkit/login.py,sha256=n4ydfo9qCsmbZq6er1xeljBD76vdTJGjbhYHMmOyDbQ,3061
+atlasai/toolkit/requests.py,sha256=X86nIo07hAjUlilZcZ1lV8RB7KOsTKbTGtcY_SpFEXY,1223
+atlasai/toolkit/utils.py,sha256=kEoNz3BbSgKbq6X2dDlM4Odcbw2BfoNyls2RKRAL5wQ,1639
+atlasai_dstoolkit_client-0.0.2.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+atlasai_dstoolkit_client-0.0.2.dist-info/METADATA,sha256=2c-eLkb3z0P2uuQ1E-3J3B4Bi-rQIEYHslmGZ1H3LSU,1301
+atlasai_dstoolkit_client-0.0.2.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+atlasai_dstoolkit_client-0.0.2.dist-info/top_level.txt,sha256=HRTbErU8nmHFDaJJ5R_XYbwpt21dqdjDpSva8xyy_0k,8
+atlasai_dstoolkit_client-0.0.2.dist-info/RECORD,,

atlasai_dstoolkit_client-0.0.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+atlasai

atlasai_dstoolkit_client-0.0.1.dist-info/RECORD

@@ -1,5 +0,0 @@
-atlasai_dstoolkit_client-0.0.1.dist-info/LICENSE.txt,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-atlasai_dstoolkit_client-0.0.1.dist-info/METADATA,sha256=Kqg9p64yglqrHsSkqwxYOcZGUDC-2D9QG1RCSS4oa0k,1301
-atlasai_dstoolkit_client-0.0.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-atlasai_dstoolkit_client-0.0.1.dist-info/top_level.txt,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-atlasai_dstoolkit_client-0.0.1.dist-info/RECORD,,

atlasai_dstoolkit_client-0.0.1.dist-info/top_level.txt

@@ -1 +0,0 @@
-