upgini 1.1.242a2__tar.gz → 1.1.242a3__tar.gz

{upgini-1.1.242a2/src/upgini.egg-info → upgini-1.1.242a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: upgini
-Version: 1.1.242a2
+Version: 1.1.242a3
 Summary: Intelligent data search & enrichment for Machine Learning
 Home-page: https://upgini.com/
 Author: Upgini Developers

{upgini-1.1.242a2 → upgini-1.1.242a3}/setup.py

@@ -40,7 +40,7 @@ def send_log(msg: str):
 
 
 here = Path(__file__).parent.resolve()
-version = "1.1.242a2"
+version = "1.1.242a3"
 try:
     send_log(f"Start setup PyLib version {version}")
     setup(

{upgini-1.1.242a2 → upgini-1.1.242a3}/src/upgini/dataset.py

@@ -20,7 +20,7 @@ from pandas.api.types import (
 from pandas.core.dtypes.common import is_period_dtype
 
 from upgini.errors import ValidationError
-from upgini.http import ProgressStage, SearchProgress, get_rest_client
+from upgini.http import ProgressStage, SearchProgress, _RestClient
 from upgini.metadata import (
     EVAL_SET_INDEX,
     SYSTEM_COLUMNS,
@@ -78,8 +78,7 @@ class Dataset:  # (pd.DataFrame):
         search_keys: Optional[List[Tuple[str, ...]]] = None,
         model_task_type: Optional[ModelTaskType] = None,
         random_state: Optional[int] = None,
-        endpoint: Optional[str] = None,
-        api_key: Optional[str] = None,
+        rest_client: Optional[_RestClient] = None,
         logger: Optional[logging.Logger] = None,
         warning_counter: Optional[WarningCounter] = None,
         **kwargs,
@@ -114,8 +113,7 @@ class Dataset:  # (pd.DataFrame):
         self.hierarchical_subgroup_keys = []
         self.file_upload_id: Optional[str] = None
         self.etalon_def: Optional[Dict[str, str]] = None
-        self.endpoint = endpoint
-        self.api_key = api_key
+        self.rest_client = rest_client
         self.random_state = random_state
         self.columns_renaming: Dict[str, str] = {}
         self.imbalanced: bool = False
@@ -983,10 +981,10 @@ class Dataset:  # (pd.DataFrame):
             runtime_parameters=runtime_parameters,
         )
 
-        if self.file_upload_id is not None and get_rest_client(self.endpoint, self.api_key).check_uploaded_file_v2(
+        if self.file_upload_id is not None and self.rest_client.check_uploaded_file_v2(
             trace_id, self.file_upload_id, file_metadata
         ):
-            search_task_response = get_rest_client(self.endpoint, self.api_key).initial_search_without_upload_v2(
+            search_task_response = self.rest_client.initial_search_without_upload_v2(
                 trace_id, self.file_upload_id, file_metadata, file_metrics, search_customization
             )
         else:
@@ -999,7 +997,7 @@ class Dataset:  # (pd.DataFrame):
                     progress_bar.progress = search_progress.to_progress_bar()
                 if progress_callback is not None:
                     progress_callback(search_progress)
-                search_task_response = get_rest_client(self.endpoint, self.api_key).initial_search_v2(
+                search_task_response = self.rest_client.initial_search_v2(
                     trace_id, parquet_file_path, file_metadata, file_metrics, search_customization
                 )
                 # if progress_bar is not None:
@@ -1015,8 +1013,7 @@ class Dataset:  # (pd.DataFrame):
             extract_features,
             accurate_model,
             task_type=self.task_type,
-            endpoint=self.endpoint,
-            api_key=self.api_key,
+            rest_client=self.rest_client,
             logger=self.logger,
         )
 
@@ -1053,10 +1050,10 @@ class Dataset:  # (pd.DataFrame):
             progress_bar.progress = search_progress.to_progress_bar()
         if progress_callback is not None:
             progress_callback(search_progress)
-        if self.file_upload_id is not None and get_rest_client(self.endpoint, self.api_key).check_uploaded_file_v2(
+        if self.file_upload_id is not None and self.rest_client.check_uploaded_file_v2(
             trace_id, self.file_upload_id, file_metadata
         ):
-            search_task_response = get_rest_client(self.endpoint, self.api_key).validation_search_without_upload_v2(
+            search_task_response = self.rest_client.validation_search_without_upload_v2(
                 trace_id, self.file_upload_id, initial_search_task_id, file_metadata, file_metrics, search_customization
             )
         else:
@@ -1065,7 +1062,7 @@ class Dataset:  # (pd.DataFrame):
                 # To avoid rate limit
                 time.sleep(1)
 
-                search_task_response = get_rest_client(self.endpoint, self.api_key).validation_search_v2(
+                search_task_response = self.rest_client.validation_search_v2(
                     trace_id,
                     parquet_file_path,
                     initial_search_task_id,
@@ -1085,8 +1082,7 @@ class Dataset:  # (pd.DataFrame):
             return_scores,
             extract_features,
             initial_search_task_id=initial_search_task_id,
-            endpoint=self.endpoint,
-            api_key=self.api_key,
+            rest_client=self.rest_client,
             logger=self.logger,
         )
 

{upgini-1.1.242a2 → upgini-1.1.242a3}/src/upgini/features_enricher.py

@@ -233,7 +233,7 @@ class FeaturesEnricher(TransformerMixin):
         self.feature_importances_ = []
         self.search_id = search_id
         if search_id:
-            search_task = SearchTask(search_id, endpoint=self.endpoint, api_key=self._api_key, logger=self.logger)
+            search_task = SearchTask(search_id, rest_client=self.rest_client, logger=self.logger)
 
             print(bundle.get("search_by_task_id_start"))
             trace_id = str(uuid.uuid4())
@@ -1817,10 +1817,9 @@ class FeaturesEnricher(TransformerMixin):
 
             dataset = Dataset(
                 "sample_" + str(uuid.uuid4()),
-                df=df_without_features,  # type: ignore
-                endpoint=self.endpoint,  # type: ignore
-                api_key=self.api_key,  # type: ignore
-                date_format=self.date_format,  # type: ignore
+                df=df_without_features,
+                date_format=self.date_format,
+                rest_client=self.rest_client,
                 logger=self.logger,
             )
             dataset.meaning_types = meaning_types
@@ -2150,11 +2149,10 @@ class FeaturesEnricher(TransformerMixin):
         dataset = Dataset(
             "tds_" + str(uuid.uuid4()),
             df=df,  # type: ignore
-            model_task_type=model_task_type,  # type: ignore
-            endpoint=self.endpoint,  # type: ignore
-            api_key=self.api_key,  # type: ignore
-            date_format=self.date_format,  # type: ignore
-            random_state=self.random_state,  # type: ignore
+            model_task_type=model_task_type,
+            date_format=self.date_format,
+            random_state=self.random_state,
+            rest_client=self.rest_client,
             logger=self.logger,
         )
         dataset.meaning_types = meaning_types

{upgini-1.1.242a2 → upgini-1.1.242a3}/src/upgini/search_task.py

@@ -8,6 +8,7 @@ import pandas as pd
 
 from upgini import dataset
 from upgini.http import (
+    _RestClient,
     ProviderTaskSummary,
     SearchProgress,
     SearchTaskSummary,
@@ -41,8 +42,7 @@ class SearchTask:
         accurate_model: bool = False,
         initial_search_task_id: Optional[str] = None,
         task_type: Optional[ModelTaskType] = None,
-        endpoint: Optional[str] = None,
-        api_key: Optional[str] = None,
+        rest_client: Optional[_RestClient] = None,
         logger: Optional[logging.Logger] = None,
     ):
         self.search_task_id = search_task_id
@@ -53,8 +53,7 @@ class SearchTask:
         self.accurate_model = accurate_model
         self.task_type = task_type
         self.summary = None
-        self.endpoint = endpoint
-        self.api_key = api_key
+        self.rest_client = rest_client
         if logger is not None:
             self.logger = logger
         else:
@@ -64,7 +63,7 @@ class SearchTask:
         self.unused_features_for_generation: Optional[List[str]] = None
 
     def get_progress(self, trace_id: str) -> SearchProgress:
-        return get_rest_client(self.endpoint, self.api_key).get_search_progress(trace_id, self.search_task_id)
+        return self.rest_client.get_search_progress(trace_id, self.search_task_id)
 
     def poll_result(self, trace_id: str, quiet: bool = False, check_fit: bool = False) -> "SearchTask":
         completed_statuses = {"COMPLETED", "VALIDATION_COMPLETED"}
@@ -72,7 +71,7 @@ class SearchTask:
         submitted_statuses = {"SUBMITTED", "VALIDATION_SUBMITTED"}
         if not quiet:
             print(bundle.get("polling_search_task").format(self.search_task_id))
-            if is_demo_api_key(self.api_key):
+            if is_demo_api_key(self.rest_client._refresh_token):
                 print(bundle.get("polling_unregister_information"))
         search_task_id = self.initial_search_task_id if self.initial_search_task_id is not None else self.search_task_id
 
@@ -80,14 +79,14 @@ class SearchTask:
             with Spinner():
                 if self.PROTECT_FROM_RATE_LIMIT:
                     time.sleep(1)  # this is neccesary to avoid requests rate limit restrictions
-                self.summary = get_rest_client(self.endpoint, self.api_key).search_task_summary_v2(
+                self.summary = self.rest_client.search_task_summary_v2(
                     trace_id, search_task_id
                 )
                 while self.summary.status not in completed_statuses and (
                     not check_fit or "VALIDATION" not in self.summary.status
                 ):
                     time.sleep(self.POLLING_DELAY_SECONDS)
-                    self.summary = get_rest_client(self.endpoint, self.api_key).search_task_summary_v2(
+                    self.summary = self.rest_client.search_task_summary_v2(
                         trace_id, search_task_id
                     )
                     if self.summary.status in failed_statuses:
@@ -103,7 +102,7 @@ class SearchTask:
         except KeyboardInterrupt as e:
             if not check_fit:
                 print(bundle.get("search_stopping"))
-                get_rest_client(self.endpoint, self.api_key).stop_search_task_v2(trace_id, search_task_id)
+                self.rest_client.stop_search_task_v2(trace_id, search_task_id)
                 self.logger.warning(f"Search {search_task_id} stopped by user")
                 print(bundle.get("search_stopped"))
             raise e
@@ -131,7 +130,7 @@ class SearchTask:
             for provider_summary in self.summary.initial_important_providers:
                 if provider_summary.status == "COMPLETED":
                     self.provider_metadata_v2.append(
-                        get_rest_client(self.endpoint, self.api_key).get_provider_search_metadata_v3(
+                        self.rest_client.get_provider_search_metadata_v3(
                             provider_summary.ads_search_task_id, trace_id
                         )
                     )
@@ -257,8 +256,8 @@ class SearchTask:
         if self.PROTECT_FROM_RATE_LIMIT:
             time.sleep(1)  # this is neccesary to avoid requests rate limit restrictions
         return _get_all_initial_raw_features_cached(
-            self.endpoint,
-            self.api_key,
+            self.rest_client._service_endpoint,
+            self.rest_client._refresh_token,
             trace_id,
             self.search_task_id,
             metrics_calculation,
@@ -268,7 +267,11 @@ class SearchTask:
     def get_target_outliers(self, trace_id: str) -> Optional[pd.DataFrame]:
         self._check_finished_initial_search()
         return _get_target_outliers_cached(
-            self.endpoint, self.api_key, trace_id, self.search_task_id, self.PROTECT_FROM_RATE_LIMIT
+            self.rest_client._service_endpoint,
+            self.rest_client._refresh_token,
+            trace_id,
+            self.search_task_id,
+            self.PROTECT_FROM_RATE_LIMIT
         )
 
     def get_max_initial_eval_set_hit_rate_v2(self) -> Optional[Dict[int, float]]:
@@ -286,8 +289,8 @@ class SearchTask:
     def get_all_validation_raw_features(self, trace_id: str, metrics_calculation=False) -> Optional[pd.DataFrame]:
         self._check_finished_validation_search()
         return _get_all_validation_raw_features_cached(
-            self.endpoint,
-            self.api_key,
+            self.rest_client._service_endpoint,
+            self.rest_client._refresh_token,
             trace_id,
             self.search_task_id,
             metrics_calculation,
@@ -295,7 +298,7 @@ class SearchTask:
         )
 
     def get_file_metadata(self, trace_id: str) -> FileMetadata:
-        return get_rest_client(self.endpoint, self.api_key).get_search_file_metadata(self.search_task_id, trace_id)
+        return self.rest_client.get_search_file_metadata(self.search_task_id, trace_id)
 
 
 @lru_cache()

{upgini-1.1.242a2 → upgini-1.1.242a3/src/upgini.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: upgini
-Version: 1.1.242a2
+Version: 1.1.242a3
 Summary: Intelligent data search & enrichment for Machine Learning
 Home-page: https://upgini.com/
 Author: Upgini Developers

{upgini-1.1.242a2 → upgini-1.1.242a3}/tests/test_etalon_validation.py

@@ -616,7 +616,7 @@ def test_columns_renaming():
     df = pd.concat([df1, df2], axis=1)
 
     dataset = Dataset(
-        "tds", df=df, meaning_types={"date": FileColumnMeaningType.DATE}, search_keys=[("date",)], endpoint="fake.url"
+        "tds", df=df, meaning_types={"date": FileColumnMeaningType.DATE}, search_keys=[("date",)]
     )
     dataset._Dataset__rename_columns()
     print(dataset)
@@ -633,7 +633,7 @@ def test_too_long_columns():
     )
 
     dataset = Dataset(
-        "tds", df=df, meaning_types={"date": FileColumnMeaningType.DATE}, search_keys=[("date",)], endpoint="fake.url"
+        "tds", df=df, meaning_types={"date": FileColumnMeaningType.DATE}, search_keys=[("date",)]
     )
     dataset._Dataset__rename_columns()
     print(dataset)
@@ -665,7 +665,7 @@ def test_downsampling_binary():
         "target": FileColumnMeaningType.TARGET,
         "eval_set_index": FileColumnMeaningType.EVAL_SET_INDEX,
     }
-    dataset = Dataset("tds", df=df, meaning_types=meaning_types, search_keys=[("date",)], endpoint="fake.url")
+    dataset = Dataset("tds", df=df, meaning_types=meaning_types, search_keys=[("date",)])
     dataset.task_type = ModelTaskType.BINARY
 
     old_min_sample_threshold = Dataset.MIN_SAMPLE_THRESHOLD
@@ -708,7 +708,7 @@ def test_downsampling_multiclass():
         "target": FileColumnMeaningType.TARGET,
         "eval_set_index": FileColumnMeaningType.EVAL_SET_INDEX,
     }
-    dataset = Dataset("tds", df=df, meaning_types=meaning_types, search_keys=[("date",)], endpoint="fake.url")
+    dataset = Dataset("tds", df=df, meaning_types=meaning_types, search_keys=[("date",)])
     dataset.task_type = ModelTaskType.MULTICLASS
 
     old_min_sample_threshold = Dataset.MIN_SAMPLE_THRESHOLD

{upgini-1.1.242a2 → upgini-1.1.242a3}/tests/test_features_enricher.py

@@ -1616,7 +1616,7 @@ def test_validation_metrics_calculation(requests_mock: Mocker):
     X = tds[["date"]]
     y = tds.target
 
-    search_task = SearchTask("", endpoint=url)
+    search_task = SearchTask("")
 
     def initial_max_hit_rate() -> Optional[float]:
         return 1.0
@@ -2363,7 +2363,7 @@ def test_search_keys_autodetection(requests_mock: Mocker):
             "email_one_domain_3b0a68",
             "date_0e8763",
         } == {sk for sublist in self.search_keys for sk in sublist}
-        search_task = SearchTask(search_task_id, self, endpoint=url, api_key="fake_api_key")
+        search_task = SearchTask(search_task_id, self, rest_client=enricher.rest_client)
         search_task.provider_metadata_v2 = [
             ProviderTaskMetadataV2(
                 features=[
@@ -2499,7 +2499,7 @@ def test_diff_target_dups(requests_mock: Mocker):
         assert self.data.loc[1, "date_0e8763"] == 1672531200000
         assert self.data.loc[1, "feature_2ad562"] == 13
         assert self.data.loc[1, "target"] == 1
-        return SearchTask("123", self, endpoint=url, api_key="fake_api_key")
+        return SearchTask("123", self, rest_client=enricher.rest_client)
 
     Dataset.search = mock_search
     Dataset.MIN_ROWS_COUNT = 1