Add some logger to help usability

src/elemeno_ai_sdk/__init__.py

@@ -1,11 +1,2 @@
 __version__ = "0.6.8"
-
-import logging
-import os
-
 __name__ = "elemeno_ai_sdk"
-
-LOGLEVEL = os.environ.get("LOGLEVEL", "WARNING").upper()
-logger = logging.getLogger()
-logger.setLevel(level=LOGLEVEL)
-print(f"Log level set to {LOGLEVEL}")

src/elemeno_ai_sdk/ml/automl/client.py

@@ -1,21 +1,23 @@
-import asyncio
-import pickle
+import json
 from typing import Dict, Optional
 
+from elemeno_ai_sdk.logger import logger
 from elemeno_ai_sdk.ml.mlhub_client import MLHubRemote
 
 
 class AutoMLClient(MLHubRemote):
     def __init__(self, env: Optional[str] = None) -> None:
         super().__init__(env=env)
 
-    def list_jobs(self):
-        return self.get(url=f"{self.base_url}/automl")
+    async def list_jobs(self):
+        response = await self.get(url=f"{self.base_url}/automl")
+        return response["resources"]
 
-    def get_job(self, job_id: str):
-        return self.get(url=f"{self.base_url}/automl/{job_id}")
+    async def get_job(self, job_id: str):
+        response = await self.get(url=f"{self.base_url}/automl/{job_id}")
+        return response
 
-    def run_job(
+    async def run_job(
         self,
         feature_table_name: str,
         features_selected: str,
@@ -42,10 +44,14 @@ def run_job(
         if features_selected != "":
             body["featuresSelected"] = features_selected
 
-        return self.post(url=f"{self.base_url}/automl", body=body)
+        try:
+            response = await self.post(url=f"{self.base_url}/automl", body=body)
+            return json.loads(response)
+        except Exception:
+            logger.exception("Failed to run job.")
 
-    def get_metadata(self, job_id: str):
-        return self.get(url=f"{self.base_url}/automl/{job_id}/metrics")
+    async def get_metadata(self, job_id: str):
+        return await self.get(url=f"{self.base_url}/automl/{job_id}/metrics")
 
     async def get_model(self, job_id: str):
         response = await self.get(url=f"{self.base_url}/automl/{job_id}/modelfile", is_binary=True)

src/elemeno_ai_sdk/ml/features/feature_table.py

@@ -1,6 +1,7 @@
 import os
 from typing import Any, Dict, List, Optional
 
+from elemeno_ai_sdk.logger import logger
 from elemeno_ai_sdk.ml.features.schema import FeatureTableSchema
 from elemeno_ai_sdk.ml.features.utils import get_feature_server_url_from_api_key
 from elemeno_ai_sdk.ml.mlhub_client import MLHubRemote
@@ -12,37 +13,29 @@ class FeatureTable(MLHubRemote):
     If you're looking to create a new feature table or read data look at ingest_schema of the class FeatureStore.
     """
 
-    def __init__(self, schema_path: str, remote_server: Optional[str] = None):
+    def __init__(self, remote_server: Optional[str] = None):
         if remote_server is None:
             api_key = os.getenv("MLHUB_API_KEY")
             if api_key is None:
                 raise ValueError("Please set the MLHUB_API_KEY environment variable.")
             self._remote_server = get_feature_server_url_from_api_key(api_key)
         else:
             self._remote_server = remote_server
-        self._schema_path = schema_path
 
-    @property
-    def name(self) -> str:
-        return self.table_schema.get("name")
-
-    @property
-    def entities(self) -> List[str]:
-        return self.table_schema.get("entities")
-
-    @property
-    def features(self) -> List[Dict[str, str]]:
-        return self.table_schema.get("schema")
-
-    @property
-    def table_schema(self):
-        return FeatureTableSchema().load_data(self._schema_path)
-
-    async def create(self) -> None:
+    async def create(self, schema_path: str) -> None:
         endpoint = f"{self._remote_server}/feature-view"
 
-        body = {"name": self.name, "entities": self.entities, "schema": self.features}
-        return await self.post(url=endpoint, body=body)
+        table_schema = FeatureTableSchema().load_data(schema_path)
+        name = table_schema["name"]
+        entities = table_schema["entities"]
+        schema = table_schema["schema"]
+
+        body = {"name": name, "entities": entities, "schema": schema}
+        try:
+            await self.post(url=endpoint, body=body)
+            logger.info(f"Feature table {name} created successfully.")
+        except Exception:
+            logger.exception(f"Failed to create feature table {name}.")
 
     async def list(self) -> List[Dict[str, Any]]:
         endpoint = f"{self._remote_server}/list-feature-views"
@@ -51,4 +44,8 @@ async def list(self) -> List[Dict[str, Any]]:
 
     async def delete(self, ft_name: str) -> None:
         endpoint = f"{self._remote_server}/{ft_name}/delete-feature-view"
-        return await self.post(url=endpoint, body={})
+        try:
+            await self.post(url=endpoint, body={})
+            logger.info(f"Deleted feature table {ft_name} successfully.")
+        except Exception:
+            logger.exception(f"Failed to delete feature table {ft_name}")

src/elemeno_ai_sdk/ml/mlhub_client.py

@@ -1,11 +1,11 @@
 import json
-import logging
 import os
 from typing import Any, Dict, Optional
 
 import aiohttp
 from tenacity import AsyncRetrying, RetryError, stop_after_attempt, wait_fixed
 
+from elemeno_ai_sdk.logger import logger
 from elemeno_ai_sdk.utils import mlhub_auth
 
 
@@ -58,7 +58,7 @@ async def post(self, url: str, body: Dict[str, Any], session: Optional[aiohttp.C
                             )
                         return await response.text()
         except RetryError:
-            logging.exception("Max retries reached")
+            logger.exception("Max retries reached")
             return None
 
     @mlhub_auth
@@ -87,5 +87,5 @@ async def get(
 
                         return await response.json(content_type=response.content_type)
         except RetryError:
-            logging.exception("Max retries reached")
+            logger.exception("Max retries reached")
             return None

src/elemeno_ai_sdk/utils.py

@@ -1,10 +1,11 @@
-import logging
 import os
 from functools import wraps
 from typing import Callable, Optional
 
 import aiohttp
 
+from elemeno_ai_sdk.logger import logger
+
 
 def mlhub_auth(func: Callable):
     @wraps(func)
@@ -18,7 +19,7 @@ async def wrapper(self, *args, session: Optional[aiohttp.ClientSession] = None,
                 kwargs["session"] = session
                 return await func(self, *args, **kwargs)
         else:
-            logging.warning(
+            logger.warning(
                 "Calling a method anottated with with_auth and passing a session object is not recommended in production."
             )
             kwargs["session"] = session