upgini 1.1.280.dev0__py3-none-any.whl → 1.2.31a1__py3-none-any.whl

upgini/http.py

@@ -39,18 +39,6 @@ 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
-
-#     __version__ = version("upgini")
-# except ImportError:
-#     try:
-#         from importlib_metadata import version  # type: ignore
-
-#         __version__ = version("upgini")
-#     except ImportError:
-#         __version__ = "Upgini wasn't installed"
-
 UPGINI_URL: str = "UPGINI_URL"
 UPGINI_API_KEY: str = "UPGINI_API_KEY"
 DEMO_API_KEY: str = "Aa4BPwGFbn1zNEXIkZ-NbhsRk0ricN6puKuga1-O5lM"
@@ -285,12 +273,14 @@ class _RestClient:
     DEACTIVATE_ADS_URI = "private/api/v2/ads/deactivate"
     TOGGLE_ADS_URI_FMT = "private/api/v2/ads/{0}/toggle"
     DELETE_ADS_URI_FMT = "private/api/v2/ads/{0}"
+    REANNOUNCE_ALL_ADS_URI_FMT = "private/api/v2/ads/reannounce-all-ads"
     POLL_ADS_MANAGEMENT_STATUS_URI_FMT = "private/api/v2/ads/management-task/{0}"
     GET_ADS_DESCRIPTION_URI_FMT = "private/api/v2/ads/{0}"
     GET_ALL_ADS_DESCRIPTIONS_URI = "private/api/v2/ads/descriptions"
     GET_ACTIVE_ADS_DEFINITIONS_URI = "private/api/v2/ads/definitions"
     UPLOAD_ONLINE_URI = "private/api/v2/ads/upload-online"
     STOP_ADS_MANAGEMENT_TASK_URI_FMT = "private/api/v2/ads/management-task/{0}/stop"
+    UNION_SEARCH_TASKS_URI_FMT = SERVICE_ROOT_V2 + "search/merge"
 
     ACCESS_TOKEN_HEADER_NAME = "Authorization"
     CONTENT_TYPE_HEADER_NAME = "Content-Type"
@@ -469,7 +459,11 @@ class _RestClient:
                         dumps(track_metrics).encode(),
                         "application/json",
                     ),
-                    "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+                    "metrics": (
+                        "metrics.json",
+                        metrics.json(exclude_none=True).encode(),
+                        "application/json",
+                    ),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
                 }
                 if search_customization is not None:
@@ -553,7 +547,11 @@ class _RestClient:
                         dumps(get_track_metrics(self.client_ip, self.client_visitorid)).encode(),
                         "application/json",
                     ),
-                    "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+                    "metrics": (
+                        "metrics.json",
+                        metrics.json(exclude_none=True).encode(),
+                        "application/json",
+                    ),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
                 }
                 if search_customization is not None:
@@ -649,7 +647,11 @@ class _RestClient:
             with open(file_path, "rb") as file:
                 files = {
                     "file": (metadata.name, file, "application/octet-stream"),
-                    "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
+                    "metadata": (
+                        "metadata.json",
+                        metadata.json(exclude_none=True).encode(),
+                        "application/json",
+                    ),
                 }
 
                 return self._send_post_file_req_v2(api_path, files)
@@ -725,6 +727,11 @@ class _RestClient:
         api_path = self.DEACTIVATE_ADS_URI
         self._with_unauth_retry(lambda: self._send_post_req(api_path, trace_id, request, result_format=None))
 
+    def reannounce_all_ads(self, trace_id: str) -> str:
+        api_path = self.REANNOUNCE_ALL_ADS_URI_FMT
+        response = self._with_unauth_retry(lambda: self._send_post_req(api_path, trace_id))
+        return response["adsManagementTaskId"]
+
     def toggle_ads(self, ads_definition_id: str, trace_id: str):
         api_path = self.TOGGLE_ADS_URI_FMT.format(ads_definition_id)
         return self._with_unauth_retry(lambda: self._send_post_req(api_path, trace_id))
@@ -760,6 +767,10 @@ class _RestClient:
         api_path = self.STOP_ADS_MANAGEMENT_TASK_URI_FMT.format(ads_management_task_id)
         self._with_unauth_retry(lambda: self._send_post_req(api_path, trace_id))
 
+    def union_search_tasks(self, request: dict, trace_id: str):
+        api_path = self.UNION_SEARCH_TASKS_URI_FMT
+        return self._with_unauth_retry(lambda: self._send_post_req(api_path, trace_id, request, result_format=None))
+
     # ---
 
     def _send_get_req(self, api_path: str, trace_id: Optional[str], additional_headers: Optional[dict] = None):
@@ -871,7 +882,7 @@ class _RestClient:
         if content_type:
             headers[_RestClient.CONTENT_TYPE_HEADER_NAME] = content_type
         if trace_id:
-            headers[_RestClient.TRACE_ID_HEADER_NAME] = trace_id
+            headers[_RestClient.TRACE_ID_HEADER_NAME] = str(trace_id)
         for header_key, header_value in additional_headers.items():
             headers[header_key] = header_value
         return headers

upgini/lazy_import.py

@@ -0,0 +1,35 @@
+import importlib
+import importlib.util
+import importlib.machinery
+
+
+class LazyImport:
+    def __init__(self, module_name, class_name):
+        self.module_name = module_name
+        self.class_name = class_name
+        self._module = None
+        self._class = None
+
+    def _load(self):
+        if self._module is None:
+            # Load module and save link to it
+            spec = importlib.util.find_spec(self.module_name)
+            if spec is None:
+                raise ImportError(f"Module {self.module_name} not found")
+
+            # Create module
+            self._module = importlib.util.module_from_spec(spec)
+
+            # Execute module
+            spec.loader.exec_module(self._module)
+
+            # Get class from module
+            self._class = getattr(self._module, self.class_name)
+
+    def __call__(self, *args, **kwargs):
+        self._load()
+        return self._class(*args, **kwargs)
+
+    def __getattr__(self, name):
+        self._load()
+        return getattr(self._class, name)

upgini/metadata.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
 SYSTEM_RECORD_ID = "system_record_id"
+ENTITY_SYSTEM_RECORD_ID = "entity_system_record_id"
+SEARCH_KEY_UNNEST = "search_key_unnest"
 SORT_ID = "sort_id"
 EVAL_SET_INDEX = "eval_set_index"
 TARGET = "target"
@@ -13,7 +15,7 @@ COUNTRY = "country_iso_code"
 RENAMED_INDEX = "index_col"
 DEFAULT_INDEX = "index"
 ORIGINAL_INDEX = "original_index"
-SYSTEM_COLUMNS = {SYSTEM_RECORD_ID, EVAL_SET_INDEX, TARGET, COUNTRY, SORT_ID}
+SYSTEM_COLUMNS = {SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST, EVAL_SET_INDEX, TARGET, COUNTRY}
 
 
 class FileColumnMeaningType(Enum):
@@ -39,6 +41,8 @@ class FileColumnMeaningType(Enum):
     POSTAL_CODE = "POSTAL_CODE"
     SYSTEM_RECORD_ID = "SYSTEM_RECORD_ID"
     EVAL_SET_INDEX = "EVAL_SET_INDEX"
+    ENTITY_SYSTEM_RECORD_ID = "ENTITY_SYSTEM_RECORD_ID"
+    UNNEST_KEY = "UNNEST_KEY"
 
 
 class SearchKey(Enum):
@@ -109,6 +113,21 @@ class SearchKey(Enum):
         if meaning_type == FileColumnMeaningType.MSISDN_RANGE_TO:
             return SearchKey.MSISDN_RANGE_TO
 
+    @staticmethod
+    def find_key(search_keys: Dict[str, SearchKey], keys: Union[SearchKey, List[SearchKey]]) -> Optional[SearchKey]:
+        if isinstance(keys, SearchKey):
+            keys = [keys]
+        for col, key_type in search_keys.items():
+            if key_type in keys:
+                return col
+        return None
+
+    @staticmethod
+    def find_all_keys(search_keys: Dict[str, SearchKey], keys: Union[SearchKey, List[SearchKey]]) -> List[SearchKey]:
+        if isinstance(keys, SearchKey):
+            keys = [keys]
+        return [col for col, key_type in search_keys.items() if key_type in keys]
+
 
 class DataType(Enum):
     INT = "INT"
@@ -153,23 +172,23 @@ class FileMetricsInterval(BaseModel):
     date_cut: float
     count: float
     valid_count: float
-    avg_target: Optional[float]  # not for multiclass
-    avg_score_etalon: Optional[float]
+    avg_target: Optional[float] = None  # not for multiclass
+    avg_score_etalon: Optional[float] = None
 
 
 class FileMetrics(BaseModel):
     # etalon metadata
-    task_type: Optional[ModelTaskType]
-    label: Optional[ModelLabelType]
-    count: Optional[int]
-    valid_count: Optional[int]
-    valid_rate: Optional[float]
-    avg_target: Optional[float]
-    metrics_binary_etalon: Optional[BinaryTask]
-    metrics_regression_etalon: Optional[RegressionTask]
-    metrics_multiclass_etalon: Optional[MulticlassTask]
-    cuts: Optional[List[float]]
-    interval: Optional[List[FileMetricsInterval]]
+    task_type: Optional[ModelTaskType] = None
+    label: Optional[ModelLabelType] = None
+    count: Optional[int] = None
+    valid_count: Optional[int] = None
+    valid_rate: Optional[float] = None
+    avg_target: Optional[float] = None
+    metrics_binary_etalon: Optional[BinaryTask] = None
+    metrics_regression_etalon: Optional[RegressionTask] = None
+    metrics_multiclass_etalon: Optional[MulticlassTask] = None
+    cuts: Optional[List[float]] = None
+    interval: Optional[List[FileMetricsInterval]] = None
 
 
 class NumericInterval(BaseModel):
@@ -183,21 +202,25 @@ class FileColumnMetadata(BaseModel):
     dataType: DataType
     meaningType: FileColumnMeaningType
     minMaxValues: Optional[NumericInterval] = None
-    originalName: Optional[str]
+    originalName: Optional[str] = None
+    # is this column contains keys from multiple key columns like msisdn1, msisdn2
+    isUnnest: bool = False
+    # list of original etalon key column names like msisdn1, msisdn2
+    unnestKeyNames: Optional[List[str]] = None
 
 
 class FileMetadata(BaseModel):
     name: str
-    description: Optional[str]
+    description: Optional[str] = None
     columns: List[FileColumnMetadata]
     searchKeys: List[List[str]]
-    excludeFeaturesSources: Optional[List[str]]
-    hierarchicalGroupKeys: Optional[List[str]]
-    hierarchicalSubgroupKeys: Optional[List[str]]
-    taskType: Optional[ModelTaskType]
-    rowsCount: Optional[int]
-    checksumMD5: Optional[str]
-    digest: Optional[str]
+    excludeFeaturesSources: Optional[List[str]] = None
+    hierarchicalGroupKeys: Optional[List[str]] = None
+    hierarchicalSubgroupKeys: Optional[List[str]] = None
+    taskType: Optional[ModelTaskType] = None
+    rowsCount: Optional[int] = None
+    checksumMD5: Optional[str] = None
+    digest: Optional[str] = None
 
     def column_by_name(self, name: str) -> Optional[FileColumnMetadata]:
         for c in self.columns:
@@ -205,13 +228,13 @@ class FileMetadata(BaseModel):
                 return c
         return None
 
-    def search_types(self) -> Set[SearchKey]:
-        search_keys = set()
+    def search_types(self) -> Dict[SearchKey, str]:
+        search_keys = dict()
         for keys_group in self.searchKeys:
             for key in keys_group:
                 column = self.column_by_name(key)
                 if column:
-                    search_keys.add(SearchKey.from_meaning_type(column.meaningType))
+                    search_keys[SearchKey.from_meaning_type(column.meaningType)] = column.name
         return search_keys
 
 
@@ -221,17 +244,17 @@ class FeaturesMetadataV2(BaseModel):
     source: str
     hit_rate: float
     shap_value: float
-    commercial_schema: Optional[str]
-    data_provider: Optional[str]
-    data_providers: Optional[List[str]]
-    data_provider_link: Optional[str]
-    data_provider_links: Optional[List[str]]
-    data_source: Optional[str]
-    data_sources: Optional[List[str]]
-    data_source_link: Optional[str]
-    data_source_links: Optional[List[str]]
-    doc_link: Optional[str]
-    update_frequency: Optional[str]
+    commercial_schema: Optional[str] = None
+    data_provider: Optional[str] = None
+    data_providers: Optional[List[str]] = None
+    data_provider_link: Optional[str] = None
+    data_provider_links: Optional[List[str]] = None
+    data_source: Optional[str] = None
+    data_sources: Optional[List[str]] = None
+    data_source_link: Optional[str] = None
+    data_source_links: Optional[List[str]] = None
+    doc_link: Optional[str] = None
+    update_frequency: Optional[str] = None
 
 
 class HitRateMetrics(BaseModel):
@@ -251,46 +274,48 @@ class ModelEvalSet(BaseModel):
 class BaseColumnMetadata(BaseModel):
     original_name: str
     hashed_name: str
-    ads_definition_id: Optional[str]
+    ads_definition_id: Optional[str] = None
     is_augmented: bool
 
 
 class GeneratedFeatureMetadata(BaseModel):
-    formula: str  # on hashed names
+    alias: Optional[str] = None
+    formula: str
     display_index: str
     base_columns: List[BaseColumnMetadata]
+    operator_params: Optional[Dict[str, str]] = None
 
 
 class ProviderTaskMetadataV2(BaseModel):
     features: List[FeaturesMetadataV2]
-    hit_rate_metrics: Optional[HitRateMetrics]
-    eval_set_metrics: Optional[List[ModelEvalSet]]
-    zero_hit_rate_search_keys: Optional[List[str]]
-    features_used_for_embeddings: Optional[List[str]]
-    shuffle_kfold: Optional[bool]
-    generated_features: Optional[List[GeneratedFeatureMetadata]]
+    hit_rate_metrics: Optional[HitRateMetrics] = None
+    eval_set_metrics: Optional[List[ModelEvalSet]] = None
+    zero_hit_rate_search_keys: Optional[List[str]] = None
+    features_used_for_embeddings: Optional[List[str]] = None
+    shuffle_kfold: Optional[bool] = None
+    generated_features: Optional[List[GeneratedFeatureMetadata]] = None
 
 
 class FeaturesFilter(BaseModel):
-    minImportance: Optional[float]
-    maxPSI: Optional[float]
-    maxCount: Optional[int]
-    selectedFeatures: Optional[List[str]]
+    minImportance: Optional[float] = None
+    maxPSI: Optional[float] = None
+    maxCount: Optional[int] = None
+    selectedFeatures: Optional[List[str]] = None
 
 
 class RuntimeParameters(BaseModel):
-    properties: Dict[str, str] = dict()
+    properties: Dict[str, Any] = {}
 
 
 class SearchCustomization(BaseModel):
-    featuresFilter: Optional[FeaturesFilter]
-    extractFeatures: Optional[bool]
-    accurateModel: Optional[bool]
-    importanceThreshold: Optional[float]
-    maxFeatures: Optional[int]
-    returnScores: Optional[bool]
-    runtimeParameters: Optional[RuntimeParameters]
-    metricsCalculation: Optional[bool]
+    featuresFilter: Optional[FeaturesFilter] = None
+    extractFeatures: Optional[bool] = None
+    accurateModel: Optional[bool] = None
+    importanceThreshold: Optional[float] = None
+    maxFeatures: Optional[int] = None
+    returnScores: Optional[bool] = None
+    runtimeParameters: Optional[RuntimeParameters] = None
+    metricsCalculation: Optional[bool] = None
 
     def __repr__(self):
         return (