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

upgini/data_source/data_source_publisher.py

@@ -3,7 +3,7 @@ import time
 import uuid
 from datetime import datetime
 from enum import Enum
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 from upgini.errors import HttpError, ValidationError
 from upgini.http import LoggerFactory, get_rest_client
@@ -47,7 +47,9 @@ class DataSourcePublisher:
         self,
         data_table_uri: str,
         search_keys: Dict[str, SearchKey],
-        update_frequency: str,
+        update_frequency: Union[
+            Literal["Daily"], Literal["Weekly"], Literal["Monthly"], Literal["Quarterly"], Literal["Annually"]
+        ],
         exclude_from_autofe_generation: Optional[List[str]],
         secondary_search_keys: Optional[Dict[str, SearchKey]] = None,
         sort_column: Optional[str] = None,
@@ -58,9 +60,43 @@ class DataSourcePublisher:
         join_date_abs_limit_days: Optional[int] = None,
         features_for_embeddings: Optional[List[str]] = DEFAULT_GENERATE_EMBEDDINGS,
         data_table_id_to_replace: Optional[str] = None,
+        keep_features: Optional[List[str]] = None,
+        date_features: Optional[List[str]] = None,
+        date_vector_features: Optional[List[str]] = None,
+        generate_runtime_embeddings: Optional[List[str]] = None,
+        exclude_raw: Optional[List[str]] = None,
         _force_generation=False,
         _silent=False,
     ) -> str:
+        """Register new ADS
+
+        Parameters
+        ----------
+        data_table_uri - str - table name in format {project_id}.{datasource_name}.{table_name}
+
+        search_keys - dict with column names as keys and SearchKey as value
+
+        update_frequency - str - (Monthly, Weekly, Daily, Annually, Quarterly)
+
+        exclude_from_autofe_generation - optional list of features that should be excluded from AutoFE
+
+        secondary_search_keys - optional dict of secondary search keys
+
+        sort_column - optional str - name of unique column that could be used for sort
+
+        date_format - optional str - format of date if it is present in search keys
+
+        features_for_embeddings - optional list of str - list of features that should be used for GPT features
+            generation
+
+        exclude_raw - optional list of str - list of features that should NOT be used as raw features
+
+        ...
+
+        data_table_id_to_replace - optional str - id of registered ADS that should be replaced by new table
+
+        keep_features - optional list - features that should not be removed from ADS (even if they are personal)
+        """
         trace_id = str(uuid.uuid4())
 
         with MDC(trace_id=trace_id):
@@ -94,6 +130,11 @@ class DataSourcePublisher:
                         "With MSISDN and DATE keys one of the snapshot_frequency_days or"
                         " join_date_abs_limit_days parameters is required"
                     )
+                if (
+                    set(search_keys.values()) == {SearchKey.PHONE, SearchKey.DATE}
+                    or set(search_keys.values()) == {SearchKey.HEM, SearchKey.DATE}
+                ) and not date_format:
+                    raise ValidationError("date_format argument is required for PHONE+DATE and HEM+DATE search keys")
 
                 request = {
                     "dataTableUri": data_table_uri,
@@ -116,6 +157,20 @@ class DataSourcePublisher:
                     request["adsDefinitionIdToReplace"] = data_table_id_to_replace
                 if exclude_from_autofe_generation is not None:
                     request["excludeFromGeneration"] = exclude_from_autofe_generation
+                if keep_features is not None:
+                    request["keepFeatures"] = keep_features
+                if date_features is not None:
+                    if date_format is None:
+                        raise ValidationError("date_format should be presented if you use date features")
+                    request["dateFeatures"] = date_features
+                if date_vector_features is not None:
+                    if date_format is None:
+                        raise ValidationError("date_format should be presented if you use date vector features")
+                    request["dateVectorFeatures"] = date_vector_features
+                if generate_runtime_embeddings is not None:
+                    request["generateRuntimeEmbeddingsFeatures"] = generate_runtime_embeddings
+                if exclude_raw is not None:
+                    request["excludeRaw"] = exclude_raw
                 self.logger.info(f"Start registering data table {request}")
 
                 task_id = self._rest_client.register_ads(request, trace_id)
@@ -173,6 +228,9 @@ class DataSourcePublisher:
                 msg = f"Data table successfully registered with id: {data_table_id}"
                 self.logger.info(msg)
                 print(msg)
+                if "warnings" in status_response and status_response["warnings"]:
+                    self.logger.warning(status_response["warnings"])
+                    print(status_response["warnings"])
                 return data_table_id
             except KeyboardInterrupt:
                 if task_id is not None:
@@ -185,11 +243,17 @@ class DataSourcePublisher:
                 self.logger.exception("Failed to register data table")
                 raise
 
-    def remove(self, data_table_ids: List[str]):
+    def remove(self, data_table_ids: Union[List[str], str]):
         trace_id = str(uuid.uuid4())
         with MDC(trace_id=trace_id):
             try:
-                if data_table_ids is None or len(data_table_ids) == 0:
+                if not data_table_ids:
+                    raise ValidationError("Empty data table ids")
+                if isinstance(data_table_ids, str):
+                    data_table_ids = [data_table_ids]
+                if not isinstance(data_table_ids, list):
+                    raise ValidationError("Invalid format of data_table_ids argument")
+                if len(data_table_ids) == 0:
                     raise ValidationError("Empty data table ids")
 
                 for data_table_id in data_table_ids:
@@ -218,16 +282,23 @@ class DataSourcePublisher:
         source_link: Optional[str] = None,
         update_frequency: Optional[str] = None,
         client_emails: Optional[List[str]] = None,
+        date_features: Optional[List[str]] = None,
+        date_vector_features: Optional[List[str]] = None,
+        exclude_from_autofe_generation: Optional[List[str]] = None,
+        generate_runtime_embeddings: Optional[List[str]] = None,
+        exclude_raw: Optional[List[str]] = None,
     ):
         trace_id = str(uuid.uuid4())
         with MDC(trace_id=trace_id):
             try:
-                if data_table_ids is None or len(data_table_ids) == 0:
+                if data_table_ids is None:
                     raise ValidationError("Empty data table ids")
                 if isinstance(data_table_ids, str):
                     data_table_ids = [data_table_ids]
                 if not isinstance(data_table_ids, list):
                     raise ValidationError("data_table_ids should be string or list of strings")
+                if len(data_table_ids) == 0:
+                    raise ValidationError("Empty data table ids")
                 if update_frequency is not None and update_frequency not in self.ACCEPTABLE_UPDATE_FREQUENCIES:
                     raise ValidationError(
                         f"Invalid update frequency: {update_frequency}. "
@@ -263,6 +334,16 @@ class DataSourcePublisher:
                     request["updateFrequency"] = update_frequency
                 if client_emails is not None:
                     request["clientEmails"] = client_emails
+                if date_features is not None:
+                    request["dateFeatures"] = date_features
+                if date_vector_features is not None:
+                    request["dateVectorFeatures"] = date_vector_features
+                if exclude_from_autofe_generation is not None:
+                    request["excludeFromGenerationFeatures"] = exclude_from_autofe_generation
+                if generate_runtime_embeddings is not None:
+                    request["generateRuntimeEmbeddingsFeatures"] = generate_runtime_embeddings
+                if exclude_raw is not None:
+                    request["excludeRaw"] = exclude_raw
                 self.logger.info(f"Activating data tables with request {request}")
 
                 self._rest_client.activate_datatables(request, trace_id)
@@ -344,3 +425,45 @@ class DataSourcePublisher:
             self.upload_online(search_keys=keys)
 
         print("All ADS-es successfully uploaded")
+
+    def union_search_tasks(
+        self,
+        search_ids: List[str],
+        target_user_email: str,
+        selected_features: Optional[List[str]] = None,
+        exclude_features: Optional[List[str]] = None,
+    ) -> str:
+        if not search_ids:
+            raise Exception("Empty search ids list")
+
+        if not target_user_email:
+            raise Exception("Empty target user email")
+
+        request = {
+            "search_task_ids": search_ids,
+            "target_user_email": target_user_email,
+        }
+        if selected_features:
+            request["selected_features"] = selected_features
+        if exclude_features:
+            request["exclude_features"] = exclude_features
+
+        response = self._rest_client.union_search_tasks(request, "trace_id")
+        print(response)
+        return response
+
+    def reannounce_all_ads(self):
+        trace_id = str(uuid.uuid4())
+        with MDC(trace_id=trace_id):
+            try:
+                task_id = self._rest_client.reannounce_all_ads(trace_id)
+                with Spinner():
+                    status_response = self._rest_client.poll_ads_management_task_status(task_id, trace_id)
+                    while status_response["status"] not in self.FINAL_STATUSES:
+                        time.sleep(5)
+                        status_response = self._rest_client.poll_ads_management_task_status(task_id, trace_id)
+
+                if status_response["status"] != "COMPLETED":
+                    raise Exception("Failed to reannounce all ADS: " + status_response["errorMessage"])
+            except Exception:
+                self.logger.exception("Failed to reannounce all ADS-es")