upgini 1.1.262a3250.post4__py3-none-any.whl → 1.1.280a3418.post2__py3-none-any.whl

upgini/features_enricher.py

@@ -1,4 +1,5 @@
 import dataclasses
+import datetime
 import gc
 import hashlib
 import itertools
@@ -20,6 +21,7 @@ from pandas.api.types import (
     is_bool,
     is_datetime64_any_dtype,
     is_numeric_dtype,
+    is_object_dtype,
     is_period_dtype,
     is_string_dtype,
 )
@@ -70,6 +72,7 @@ from upgini.utils.datetime_utils import (
     DateTimeSearchKeyConverter,
     is_blocked_time_series,
     is_time_series,
+    validate_dates_distribution,
 )
 from upgini.utils.deduplicate_utils import (
     clean_full_duplicates,
@@ -93,7 +96,7 @@ try:
 except Exception:
     from upgini.utils.fallback_progress_bar import CustomFallbackProgressBar as ProgressBar
 
-from upgini.utils.target_utils import define_task
+from upgini.utils.target_utils import calculate_psi, define_task
 from upgini.utils.warning_counter import WarningCounter
 from upgini.version_validator import validate_version
 
@@ -145,6 +148,7 @@ class FeaturesEnricher(TransformerMixin):
     """
 
     TARGET_NAME = "target"
+    CURRENT_DATE = "current_date"
     RANDOM_STATE = 42
     CALCULATE_METRICS_THRESHOLD = 50_000_000
     CALCULATE_METRICS_MIN_THRESHOLD = 500
@@ -206,6 +210,7 @@ class FeaturesEnricher(TransformerMixin):
         client_ip: Optional[str] = None,
         client_visitorid: Optional[str] = None,
         custom_bundle_config: Optional[str] = None,
+        add_date_if_missing: bool = True,
         **kwargs,
     ):
         self.bundle = get_custom_bundle(custom_bundle_config)
@@ -316,6 +321,7 @@ class FeaturesEnricher(TransformerMixin):
         self.raise_validation_error = raise_validation_error
         self.exclude_columns = exclude_columns
         self.baseline_score_column = baseline_score_column
+        self.add_date_if_missing = add_date_if_missing
 
     def _get_api_key(self):
         return self._api_key
@@ -419,11 +425,14 @@ class FeaturesEnricher(TransformerMixin):
 
             self.__validate_search_keys(self.search_keys, self.search_id)
 
+            # Validate client estimator params
+            self._get_client_cat_features(estimator, X, self.search_keys)
+
             try:
                 self.X = X
                 self.y = y
                 self.eval_set = self._check_eval_set(eval_set, X, self.bundle)
-                self.dump_input(trace_id, X, y, eval_set)
+                self.dump_input(trace_id, X, y, self.eval_set)
                 self.__inner_fit(
                     trace_id,
                     X,
@@ -562,7 +571,7 @@ class FeaturesEnricher(TransformerMixin):
                 self.X = X
                 self.y = y
                 self.eval_set = self._check_eval_set(eval_set, X, self.bundle)
-                self.dump_input(trace_id, X, y, eval_set)
+                self.dump_input(trace_id, X, y, self.eval_set)
 
                 if _num_samples(drop_duplicates(X)) > Dataset.MAX_ROWS:
                     raise ValidationError(self.bundle.get("dataset_too_many_rows_registered").format(Dataset.MAX_ROWS))
@@ -812,6 +821,7 @@ class FeaturesEnricher(TransformerMixin):
         trace_id = trace_id or str(uuid.uuid4())
         start_time = time.time()
         with MDC(trace_id=trace_id):
+            self.logger.info("Start calculate metrics")
             if len(args) > 0:
                 msg = f"WARNING: Unsupported positional arguments for calculate_metrics: {args}"
                 self.logger.warning(msg)
@@ -822,12 +832,16 @@ class FeaturesEnricher(TransformerMixin):
                 print(msg)
 
             self.__validate_search_keys(self.search_keys, self.search_id)
+            effective_X = X if X is not None else self.X
+            effective_y = y if y is not None else self.y
+            effective_eval_set = eval_set if eval_set is not None else self.eval_set
+            effective_eval_set = self._check_eval_set(effective_eval_set, effective_X, self.bundle)
 
             try:
                 self.__log_debug_information(
-                    X if X is not None else self.X,
-                    y if y is not None else self.y,
-                    eval_set if eval_set is not None else self.eval_set,
+                    effective_X,
+                    effective_y,
+                    effective_eval_set,
                     exclude_features_sources=exclude_features_sources,
                     cv=cv if cv is not None else self.cv,
                     importance_threshold=importance_threshold,
@@ -841,17 +855,14 @@ class FeaturesEnricher(TransformerMixin):
                     self._search_task is None
                     or self._search_task.provider_metadata_v2 is None
                     or len(self._search_task.provider_metadata_v2) == 0
-                    or (self.X is None and X is None)
-                    or (self.y is None and y is None)
+                    or effective_X is None
+                    or effective_y is None
                 ):
                     raise ValidationError(self.bundle.get("metrics_unfitted_enricher"))
 
                 if X is not None and y is None:
                     raise ValidationError("X passed without y")
 
-                effective_X = X if X is not None else self.X
-                effective_eval_set = eval_set if eval_set is not None else self.eval_set
-
                 validate_scoring_argument(scoring)
 
                 self._validate_baseline_score(effective_X, effective_eval_set)
@@ -862,29 +873,15 @@ class FeaturesEnricher(TransformerMixin):
                     self.__display_support_link(msg)
                     return None
 
-                cat_features = None
-                search_keys_for_metrics = []
-                if (
-                    estimator is not None
-                    and hasattr(estimator, "get_param")
-                    and estimator.get_param("cat_features") is not None
-                ):
-                    cat_features = estimator.get_param("cat_features")
-                    if len(cat_features) > 0 and isinstance(cat_features[0], int):
-                        effectiveX = X or self.X
-                        cat_features = [effectiveX.columns[i] for i in cat_features]
-                        for cat_feature in cat_features:
-                            if cat_feature in self.search_keys:
-                                if self.search_keys[cat_feature] in [SearchKey.COUNTRY, SearchKey.POSTAL_CODE]:
-                                    search_keys_for_metrics.append(cat_feature)
-                                else:
-                                    raise ValidationError(self.bundle.get("cat_feature_search_key").format(cat_feature))
+                cat_features, search_keys_for_metrics = self._get_client_cat_features(
+                    estimator, effective_X, self.search_keys
+                )
 
                 prepared_data = self._prepare_data_for_metrics(
                     trace_id=trace_id,
-                    X=X,
-                    y=y,
-                    eval_set=eval_set,
+                    X=effective_X,
+                    y=effective_y,
+                    eval_set=effective_eval_set,
                     exclude_features_sources=exclude_features_sources,
                     importance_threshold=importance_threshold,
                     max_features=max_features,
@@ -893,6 +890,7 @@ class FeaturesEnricher(TransformerMixin):
                     search_keys_for_metrics=search_keys_for_metrics,
                     progress_bar=progress_bar,
                     progress_callback=progress_callback,
+                    cat_features=cat_features,
                 )
                 if prepared_data is None:
                     return None
@@ -994,8 +992,6 @@ class FeaturesEnricher(TransformerMixin):
                         enriched_metric = None
                         uplift = None
 
-                    effective_X = X if X is not None else self.X
-                    effective_y = y if y is not None else self.y
                     train_metrics = {
                         self.bundle.get("quality_metrics_segment_header"): self.bundle.get(
                             "quality_metrics_train_segment"
@@ -1256,6 +1252,7 @@ class FeaturesEnricher(TransformerMixin):
             ).get_cv_and_groups(X)
         else:
             from sklearn import __version__ as sklearn_version
+
             try:
                 from sklearn.model_selection._split import GroupsConsumerMixin
 
@@ -1269,6 +1266,29 @@ class FeaturesEnricher(TransformerMixin):
 
         return _cv, groups
 
+    def _get_client_cat_features(
+        self, estimator: Optional[Any], X: pd.DataFrame, search_keys: Dict[str, SearchKey]
+    ) -> Optional[List[str]]:
+        cat_features = None
+        search_keys_for_metrics = []
+        if (
+            estimator is not None
+            and hasattr(estimator, "get_param")
+            and estimator.get_param("cat_features") is not None
+        ):
+            cat_features = estimator.get_param("cat_features")
+            if len(cat_features) > 0:
+                if all([isinstance(f, int) for f in cat_features]):
+                    cat_features = [X.columns[i] for i in cat_features]
+                self.logger.info(f"Collected categorical features {cat_features} from user estimator")
+                for cat_feature in cat_features:
+                    if cat_feature in search_keys:
+                        if search_keys[cat_feature] in [SearchKey.COUNTRY, SearchKey.POSTAL_CODE]:
+                            search_keys_for_metrics.append(cat_feature)
+                        else:
+                            raise ValidationError(self.bundle.get("cat_feature_search_key").format(cat_feature))
+        return cat_features, search_keys_for_metrics
+
     def _prepare_data_for_metrics(
         self,
         trace_id: str,
@@ -1283,6 +1303,7 @@ class FeaturesEnricher(TransformerMixin):
         search_keys_for_metrics: Optional[List[str]] = None,
         progress_bar: Optional[ProgressBar] = None,
         progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        cat_features: Optional[List[str]] = None,
     ):
         is_input_same_as_fit, X, y, eval_set = self._is_input_same_as_fit(X, y, eval_set)
         is_demo_dataset = hash_input(X, y, eval_set) in DEMO_DATASET_HASHES
@@ -1340,9 +1361,8 @@ class FeaturesEnricher(TransformerMixin):
 
         # Detect and drop high cardinality columns in train
         columns_with_high_cardinality = FeaturesValidator.find_high_cardinality(fitting_X)
-        columns_with_high_cardinality = [
-            c for c in columns_with_high_cardinality if c not in (self.generate_features or [])
-        ]
+        non_excluding_columns = (self.generate_features or []) + (cat_features or [])
+        columns_with_high_cardinality = [c for c in columns_with_high_cardinality if c not in non_excluding_columns]
         if len(columns_with_high_cardinality) > 0:
             self.logger.warning(
                 f"High cardinality columns {columns_with_high_cardinality} will be dropped for metrics calculation"
@@ -1684,6 +1704,9 @@ class FeaturesEnricher(TransformerMixin):
             df = validated_X.copy()
 
             df[TARGET] = validated_y
+
+            df = clean_full_duplicates(df, logger=self.logger, silent=True, bundle=self.bundle)
+
             num_samples = _num_samples(df)
             if num_samples > Dataset.FIT_SAMPLE_THRESHOLD:
                 self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_ROWS}")
@@ -1801,10 +1824,11 @@ class FeaturesEnricher(TransformerMixin):
         else:
             features_section = ""
 
-        api_example = f"""curl 'https://inference-upgini.azurewebsites.net/api/http_inference_trigger' \\
+        search_id = self._search_task.search_task_id
+        api_example = f"""curl 'https://search.upgini.com/online/api/http_inference_trigger?search_id={search_id}' \\
     -H 'Authorization: {self.api_key}' \\
     -H 'Content-Type: application/json' \\
-    -d '{{"search_id": "{self._search_task.search_task_id}", "search_keys": {keys}{features_section}}}'"""
+    -d '{{"search_keys": {keys}{features_section}}}'"""
         return api_example
 
     def _get_copy_of_runtime_parameters(self) -> RuntimeParameters:
@@ -1899,6 +1923,8 @@ class FeaturesEnricher(TransformerMixin):
                 generated_features.extend(converter.generated_features)
             else:
                 self.logger.info("Input dataset hasn't date column")
+                if self.add_date_if_missing:
+                    df = self._add_current_date_as_key(df, search_keys, self.logger, self.bundle)
             email_column = self._get_email_column(search_keys)
             hem_column = self._get_hem_column(search_keys)
             email_converted_to_hem = False
@@ -1918,6 +1944,7 @@ class FeaturesEnricher(TransformerMixin):
 
             meaning_types = {col: key.value for col, key in search_keys.items()}
             non_keys_columns = [column for column in df.columns if column not in search_keys.keys()]
+
             if email_converted_to_hem:
                 non_keys_columns.append(email_column)
 
@@ -1939,6 +1966,7 @@ class FeaturesEnricher(TransformerMixin):
             if add_fit_system_record_id:
                 df = self.__add_fit_system_record_id(df, dict(), search_keys)
                 df = df.rename(columns={SYSTEM_RECORD_ID: SORT_ID})
+                non_keys_columns.append(SORT_ID)
 
             columns_for_system_record_id = sorted(list(search_keys.keys()) + (original_features_for_transform or []))
 
@@ -2215,14 +2243,13 @@ class FeaturesEnricher(TransformerMixin):
         self.fit_search_keys = self.search_keys.copy()
         self.fit_search_keys = self.__prepare_search_keys(validated_X, self.fit_search_keys, is_demo_dataset)
 
-        has_date = self._get_date_column(self.fit_search_keys) is not None
+        validate_dates_distribution(validated_X, self.fit_search_keys, self.logger, self.bundle, self.warning_counter)
+
+        maybe_date_column = self._get_date_column(self.fit_search_keys)
+        has_date = maybe_date_column is not None
         model_task_type = self.model_task_type or define_task(validated_y, has_date, self.logger)
         self._validate_binary_observations(validated_y, model_task_type)
 
-        df = self.__handle_index_search_keys(df, self.fit_search_keys)
-
-        df = self.__correct_target(df)
-
         self.runtime_parameters = get_runtime_params_custom_loss(
             self.loss, model_task_type, self.runtime_parameters, self.logger
         )
@@ -2234,6 +2261,13 @@ class FeaturesEnricher(TransformerMixin):
                 eval_df[EVAL_SET_INDEX] = idx + 1
                 df = pd.concat([df, eval_df])
 
+        df = self.__correct_target(df)
+
+        df = self.__handle_index_search_keys(df, self.fit_search_keys)
+
+        if is_numeric_dtype(df[self.TARGET_NAME]) and has_date:
+            self._validate_PSI(df.sort_values(by=maybe_date_column))
+
         if DEFAULT_INDEX in df.columns:
             msg = self.bundle.get("unsupported_index_column")
             self.logger.info(msg)
@@ -2260,6 +2294,8 @@ class FeaturesEnricher(TransformerMixin):
             self.fit_generated_features.extend(converter.generated_features)
         else:
             self.logger.info("Input dataset hasn't date column")
+            if self.add_date_if_missing:
+                df = self._add_current_date_as_key(df, self.fit_search_keys, self.logger, self.bundle)
         email_column = self._get_email_column(self.fit_search_keys)
         hem_column = self._get_hem_column(self.fit_search_keys)
         email_converted_to_hem = False
@@ -2512,7 +2548,7 @@ class FeaturesEnricher(TransformerMixin):
             validated_X = X.copy()
         elif isinstance(X, pd.Series):
             validated_X = X.to_frame()
-        elif isinstance(X, np.ndarray) or isinstance(X, list):
+        elif isinstance(X, (list, np.ndarray)):
             validated_X = pd.DataFrame(X)
             renaming = {c: str(c) for c in validated_X.columns}
             validated_X = validated_X.rename(columns=renaming)
@@ -2601,7 +2637,7 @@ class FeaturesEnricher(TransformerMixin):
             validated_eval_X = eval_X.copy()
         elif isinstance(eval_X, pd.Series):
             validated_eval_X = eval_X.to_frame()
-        elif isinstance(eval_X, np.ndarray) or isinstance(eval_X, list):
+        elif isinstance(eval_X, (list, np.ndarray)):
             validated_eval_X = pd.DataFrame(eval_X)
             renaming = {c: str(c) for c in validated_eval_X.columns}
             validated_eval_X = validated_eval_X.rename(columns=renaming)
@@ -2783,7 +2819,7 @@ class FeaturesEnricher(TransformerMixin):
             )
 
             def sample(df):
-                if isinstance(df, pd.Series) or isinstance(df, pd.DataFrame):
+                if isinstance(df, (pd.DataFrame, pd.Series)):
                     return df.head(10)
                 else:
                     return df[:10]
@@ -2808,6 +2844,7 @@ class FeaturesEnricher(TransformerMixin):
 
             maybe_date_col = self._get_date_column(self.search_keys)
             if X is not None and maybe_date_col is not None and maybe_date_col in X.columns:
+                # TODO cast date column to single dtype
                 min_date = X[maybe_date_col].min()
                 max_date = X[maybe_date_col].max()
                 self.logger.info(f"Dates interval is ({min_date}, {max_date})")
@@ -2839,6 +2876,25 @@ class FeaturesEnricher(TransformerMixin):
             if t in [SearchKey.DATE, SearchKey.DATETIME]:
                 return col
 
+    @staticmethod
+    def _add_current_date_as_key(
+        df: pd.DataFrame, search_keys: Dict[str, SearchKey], logger: logging.Logger, bundle: ResourceBundle
+    ) -> pd.DataFrame:
+        if (
+            set(search_keys.values()) == {SearchKey.PHONE}
+            or set(search_keys.values()) == {SearchKey.EMAIL}
+            or set(search_keys.values()) == {SearchKey.HEM}
+            or set(search_keys.values()) == {SearchKey.COUNTRY, SearchKey.POSTAL_CODE}
+        ):
+            msg = bundle.get("current_date_added")
+            print(msg)
+            logger.warning(msg)
+            df[FeaturesEnricher.CURRENT_DATE] = datetime.date.today()
+            search_keys[FeaturesEnricher.CURRENT_DATE] = SearchKey.DATE
+            converter = DateTimeSearchKeyConverter(FeaturesEnricher.CURRENT_DATE, None, logger, bundle)
+            df = converter.convert(df)
+        return df
+
     @staticmethod
     def _get_group_columns(df: pd.DataFrame, search_keys: Dict[str, SearchKey]) -> List[str]:
         return [
@@ -2877,26 +2933,33 @@ class FeaturesEnricher(TransformerMixin):
 
         # order by date and idempotent order by other keys
         if self.cv not in [CVType.time_series, CVType.blocked_time_series]:
+            sort_exclude_columns = [original_order_name, ORIGINAL_INDEX, EVAL_SET_INDEX, TARGET, "__target"]
             if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
+                sort_exclude_columns.append(self._get_date_column(search_keys))
             else:
                 date_column = self._get_date_column(search_keys)
             sort_columns = [date_column] if date_column is not None else []
 
-            other_search_keys = sorted(
+            other_columns = sorted(
                 [
-                    sk
-                    for sk, key_type in search_keys.items()
-                    if key_type not in [SearchKey.DATE, SearchKey.DATETIME]
-                    and sk in df.columns
-                    and df[sk].nunique() > 1  # don't use constant keys for hash
+                    c
+                    for c in df.columns
+                    if c not in sort_columns and c not in sort_exclude_columns and df[c].nunique() > 1
                 ]
+                # [
+                #     sk
+                #     for sk, key_type in search_keys.items()
+                #     if key_type not in [SearchKey.DATE, SearchKey.DATETIME]
+                #     and sk in df.columns
+                #     and df[sk].nunique() > 1  # don't use constant keys for hash
+                # ]
             )
 
             search_keys_hash = "search_keys_hash"
-            if len(other_search_keys) > 0:
+            if len(other_columns) > 0:
                 sort_columns.append(search_keys_hash)
-                df[search_keys_hash] = pd.util.hash_pandas_object(df[sorted(other_search_keys)], index=False)
+                df[search_keys_hash] = pd.util.hash_pandas_object(df[other_columns], index=False)
 
             df = df.sort_values(by=sort_columns)
 
@@ -2920,7 +2983,7 @@ class FeaturesEnricher(TransformerMixin):
 
     def __correct_target(self, df: pd.DataFrame) -> pd.DataFrame:
         target = df[self.TARGET_NAME]
-        if is_string_dtype(target):
+        if is_string_dtype(target) or is_object_dtype(target):
             maybe_numeric_target = pd.to_numeric(target, errors="coerce")
             # If less than 5% is non numeric then leave this rows with NaN target and later it will be dropped
             if maybe_numeric_target.isna().sum() <= _num_samples(df) * 0.05:
@@ -3185,22 +3248,23 @@ class FeaturesEnricher(TransformerMixin):
                 return None
             features_meta = self._search_task.get_all_features_metadata_v2()
 
-            def get_feature_by_display_index(idx, op):
+            def get_feature_by_name(name: str):
                 for m in features_meta:
-                    if m.name.endswith(f"_{op}_{idx}"):
+                    if m.name == name:
                         return m
 
             descriptions = []
             for m in autofe_meta:
                 autofe_feature = Feature.from_formula(m.formula)
+                orig_to_hashed = {base_column.original_name: base_column.hashed_name for base_column in m.base_columns}
+                autofe_feature.rename_columns(orig_to_hashed)
+                autofe_feature.set_display_index(m.display_index)
                 if autofe_feature.op.is_vector:
                     continue
 
                 description = dict()
 
-                feature_meta = get_feature_by_display_index(
-                    m.display_index, autofe_feature.op.alias or autofe_feature.op.name
-                )
+                feature_meta = get_feature_by_name(autofe_feature.get_display_name(shorten=True))
                 if feature_meta is None:
                     self.logger.warning(f"Feature meta for display index {m.display_index} not found")
                     continue
@@ -3321,7 +3385,8 @@ class FeaturesEnricher(TransformerMixin):
                 valid_search_keys[column_name] = SearchKey.CUSTOM_KEY
             else:
                 if x[column_name].isnull().all() or (
-                    is_string_dtype(x[column_name]) and (x[column_name].astype("string").str.strip() == "").all()
+                    (is_string_dtype(x[column_name]) or is_object_dtype(x[column_name]))
+                    and (x[column_name].astype("string").str.strip() == "").all()
                 ):
                     raise ValidationError(self.bundle.get("empty_search_key").format(column_name))
 
@@ -3547,6 +3612,34 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg)
             print(msg)
 
+    def _validate_PSI(self, df: pd.DataFrame):
+        if EVAL_SET_INDEX in df.columns:
+            train = df.query(f"{EVAL_SET_INDEX} == 0")
+            eval1 = df.query(f"{EVAL_SET_INDEX} == 1")
+        else:
+            train = df
+            eval1 = None
+
+        # 1. Check train PSI
+        half_train = round(len(train) / 2)
+        part1 = train[:half_train]
+        part2 = train[half_train:]
+        train_psi = calculate_psi(part1[self.TARGET_NAME], part2[self.TARGET_NAME])
+        if train_psi > 0.2:
+            self.warning_counter.increment()
+            msg = self.bundle.get("train_unstable_target").format(train_psi)
+            print(msg)
+            self.logger.warning(msg)
+
+        # 2. Check train-test PSI
+        if eval1 is not None:
+            train_test_psi = calculate_psi(train[self.TARGET_NAME], eval1[self.TARGET_NAME])
+            if train_test_psi > 0.2:
+                self.warning_counter.increment()
+                msg = self.bundle.get("eval_unstable_target").format(train_test_psi)
+                print(msg)
+                self.logger.warning(msg)
+
     def _dump_python_libs(self):
         try:
             from pip._internal.operations.freeze import freeze
@@ -3600,7 +3693,7 @@ class FeaturesEnricher(TransformerMixin):
                 def sample(inp, sample_index):
                     if _num_samples(inp) <= 1000:
                         return inp
-                    if isinstance(inp, pd.DataFrame) or isinstance(inp, pd.Series):
+                    if isinstance(inp, (pd.DataFrame, pd.Series)):
                         return inp.sample(n=1000, random_state=random_state)
                     if isinstance(inp, np.ndarray):
                         return inp[sample_index]
@@ -3613,7 +3706,7 @@ class FeaturesEnricher(TransformerMixin):
                     if y is not None:
                         with open(f"{tmp_dir}/y.pickle", "wb") as y_file:
                             pickle.dump(sample(y, xy_sample_index), y_file)
-                        if eval_set is not None:
+                        if eval_set:
                             eval_xy_sample_index = rnd.randint(0, _num_samples(eval_set[0][0]), size=1000)
                             with open(f"{tmp_dir}/eval_x.pickle", "wb") as eval_x_file:
                                 pickle.dump(sample(eval_set[0][0], eval_xy_sample_index), eval_x_file)

upgini/http.py

@@ -22,6 +22,7 @@ from pydantic import BaseModel
 from pythonjsonlogger import jsonlogger
 from requests.exceptions import RequestException
 
+from upgini.__about__ import __version__
 from upgini.errors import (
     HttpError,
     UnauthorizedError,
@@ -38,17 +39,17 @@ from upgini.metadata import (
 from upgini.resource_bundle import bundle
 from upgini.utils.track_info import get_track_metrics
 
-try:
-    from importlib_metadata import version  # type: ignore
+# try:
+#     from importlib.metadata import version  # type: ignore
 
-    __version__ = version("upgini")
-except ImportError:
-    try:
-        from importlib.metadata import version  # type: ignore
+#     __version__ = version("upgini")
+# except ImportError:
+#     try:
+#         from importlib_metadata import version  # type: ignore
 
-        __version__ = version("upgini")
-    except ImportError:
-        __version__ = "Upgini wasn't installed"
+#         __version__ = version("upgini")
+#     except ImportError:
+#         __version__ = "Upgini wasn't installed"
 
 UPGINI_URL: str = "UPGINI_URL"
 UPGINI_API_KEY: str = "UPGINI_API_KEY"
@@ -925,7 +926,7 @@ def is_demo_api_key(api_token: Optional[str]) -> bool:
     return api_token is None or api_token == "" or api_token == DEMO_API_KEY
 
 
-@lru_cache()
+@lru_cache
 def _get_rest_client(
     backend_url: str, api_token: str, client_ip: Optional[str] = None, client_visitorid: Optional[str] = None
 ) -> _RestClient:

upgini/mdc/__init__.py

@@ -1,15 +1,13 @@
-# -*- coding: utf-8 -*-
 """
 .. module: mdc
 .. moduleauthor:: Aljosha Friemann a.friemann@automate.wtf
 """
-from __future__ import absolute_import, division, print_function, unicode_literals
 
 import logging
 
-from upgini.mdc.context import new_log_context, get_mdc_fields
 from pythonjsonlogger import jsonlogger
 
+from upgini.mdc.context import get_mdc_fields, new_log_context
 
 MDContext = new_log_context
 MDC = new_log_context

upgini/mdc/context.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 .. module: TODO
     :platform: TODO
@@ -7,12 +6,11 @@
 .. moduleauthor:: Aljosha Friemann a.friemann@automate.wtf
 """
 
-import time
-import uuid
+import collections
 import logging
 import threading
-import collections
-
+import time
+import uuid
 from contextlib import contextmanager
 
 LOGGER = logging.getLogger(__name__)
@@ -32,7 +30,7 @@ def get_mdc_fields():
 
 @contextmanager
 def new_log_context(**kwargs):
-    context_id = "mdc-{thread}-{context}".format(thread=threading.current_thread().ident, context=uuid.uuid4())
+    context_id = f"mdc-{threading.current_thread().ident}-{uuid.uuid4()}"
 
     LOGGER.debug("creating context %s", context_id)
 

upgini/metadata.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from enum import Enum
 from typing import Dict, List, Optional, Set
 
@@ -201,6 +203,7 @@ class FileMetadata(BaseModel):
         for c in self.columns:
             if c.name == name:
                 return c
+        return None
 
     def search_types(self) -> Set[SearchKey]:
         search_keys = set()