upgini 1.1.242a2__py3-none-any.whl → 1.1.243__py3-none-any.whl

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/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/http.py

@@ -308,7 +308,6 @@ class _RestClient:
         # self.silent_mode = silent_mode
         self.client_ip = client_ip
         self.client_visitorid = client_visitorid
-        print(f"Created RestClient with {client_ip} and {client_visitorid}")
         self._access_token = self._refresh_access_token()
         # self._access_token: Optional[str] = None  # self._refresh_access_token()
         self.last_refresh_time = time.time()
@@ -442,9 +441,7 @@ class _RestClient:
     ) -> SearchTaskResponse:
         api_path = self.INITIAL_SEARCH_URI_FMT_V2
 
-        print(f"Start initial search with {self.client_ip} and {self.client_visitorid}")
         track_metrics = get_track_metrics(self.client_ip, self.client_visitorid)
-        print(f"Sending track metrics: {track_metrics}")
 
         def open_and_send():
             md5_hash = hashlib.md5()
@@ -486,7 +483,7 @@ class _RestClient:
                     api_path, files, trace_id=trace_id, additional_headers=additional_headers
                 )
 
-        response = self._with_unauth_retry(lambda: open_and_send())
+        response = self._with_unauth_retry(open_and_send)
         return SearchTaskResponse(response)
 
     def check_uploaded_file_v2(self, trace_id: str, file_upload_id: str, metadata: FileMetadata) -> bool:
@@ -571,7 +568,7 @@ class _RestClient:
                     api_path, files, trace_id=trace_id, additional_headers=additional_headers
                 )
 
-        response = self._with_unauth_retry(lambda: open_and_send())
+        response = self._with_unauth_retry(open_and_send)
         return SearchTaskResponse(response)
 
     def validation_search_without_upload_v2(

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.dist-info → upgini-1.1.243.dist-info}/METADATA

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

{upgini-1.1.242a2.dist-info → upgini-1.1.243.dist-info}/RECORD

@@ -1,13 +1,13 @@
 upgini/__init__.py,sha256=asENHgEVHQBIkV-e_0IhE_ZWqkCG6398U3ZLrNzAH6k,407
 upgini/ads.py,sha256=mre6xn44wcC_fg63iLT_kTh4mViZqR9AKRJZAtpQz8Y,2592
-upgini/dataset.py,sha256=qSjv09LKzCYayucb_JlhExw9uSRcscLWTaD8hqATE3s,49676
+upgini/dataset.py,sha256=y9rpNhdLU9QgfFZndrPGK-S6CL67q5ocmB9HMzwHtaA,49395
 upgini/errors.py,sha256=BqpvfhW2jJW5fa5KXj0alhXatGl-WK4xTl309-QNLp8,959
-upgini/features_enricher.py,sha256=kAZqLVG7hxIROp0l_b30n7QG4EUhqt4zbx7wHPHCNug,163196
+upgini/features_enricher.py,sha256=n2L9MWq4WoUQIzoDDECFyiuprwZslFPPhbLfpXsT3sQ,162975
 upgini/fingerprint.js,sha256=VygVIQlN1v4NGZfjHqtRogOw8zjTnnMNJg_f7M5iGQU,33442
-upgini/http.py,sha256=xeSatYNnSBMQfGMXsER_ZvhR5zfDTY8_E1g3YpIOb38,42477
+upgini/http.py,sha256=jsAYvipo6jeN3qsZeSuFkIe4BigGiOM813G5PXWVvb0,42234
 upgini/metadata.py,sha256=FZ5CQluLLWrfrBVThSIes1SW6wcs7n50aNZwzYnHiF0,9584
 upgini/metrics.py,sha256=rteVPPjDFYlL5bBFVpu-YwwXQGNV1IzwT7V7L9JtjaE,23762
-upgini/search_task.py,sha256=C6nqi6eG8pHqrTByC_iNWF1ZCJXLfP0hXgj56V7oEd4,17217
+upgini/search_task.py,sha256=nTVrb3CE4M1zfDkI-W_qVdUhsc90b98w3lo0XxegeKo,17200
 upgini/spinner.py,sha256=yhakBaydMNS8E8TRAwTdCMdnWrHeWT0cR1M8c9hP6jA,1157
 upgini/version_validator.py,sha256=rDIncP6BEko4J2F2hUcMOtKm_vZbI4ICWcNcw8hrwM4,1400
 upgini/ads_management/__init__.py,sha256=qzyisOToVRP-tquAJD1PblZhNtMrOB8FiyF9JvfkvgE,50
@@ -53,8 +53,8 @@ upgini/utils/sklearn_ext.py,sha256=IMx2La70AXAggApVpT7sMEjWqVWon5AMZt4MARDsIMQ,4
 upgini/utils/target_utils.py,sha256=_VjYUm4ECXbgNvxNupr982fuOK_jtkg-8Xw7-zJBz2w,1708
 upgini/utils/track_info.py,sha256=EPcJ13Jqa17_T0JjM37Ac9kWDz5Zk0GVsIZKutOb8aU,5207
 upgini/utils/warning_counter.py,sha256=vnmdFo5-7GBkU2bK9h_uC0K0Y_wtfcYstxOdeRfacO0,228
-upgini-1.1.242a2.dist-info/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
-upgini-1.1.242a2.dist-info/METADATA,sha256=IfEZQIEGTjE4Fze8z-b_kw1RnWCGE1XABgsea-Fiz28,48264
-upgini-1.1.242a2.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-upgini-1.1.242a2.dist-info/top_level.txt,sha256=OFhTGiDIWKl5gFI49qvWq1R9IKflPaE2PekcbDXDtx4,7
-upgini-1.1.242a2.dist-info/RECORD,,
+upgini-1.1.243.dist-info/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
+upgini-1.1.243.dist-info/METADATA,sha256=-4iFDXdoXZA9bcQArbodiC9TcvdKDT9Qc6GFZi7dNYc,48262
+upgini-1.1.243.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+upgini-1.1.243.dist-info/top_level.txt,sha256=OFhTGiDIWKl5gFI49qvWq1R9IKflPaE2PekcbDXDtx4,7
+upgini-1.1.243.dist-info/RECORD,,