abacusai 1.1.5__tar.gz → 1.1.6__tar.gz

{abacusai-1.1.5/abacusai.egg-info → abacusai-1.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: abacusai
-Version: 1.1.5
+Version: 1.1.6
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.1.5"
+__version__ = "1.1.6"

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/api_class/batch_prediction.py

@@ -24,9 +24,11 @@ class AnomalyOutliersBatchPredictionArgs(BatchPredictionArgs):
     Args:
         for_eval (bool): If True, the test fold which was created during training and used for metrics calculation will be used as input data. These predictions are hence, used for model evaluation.
         threshold (float): The threshold for detecting an anomaly. Range: [0.8, 0.99]
+        explain_predictions (bool): If True, calculates explanations for the predicted values along with predictions.
     """
     for_eval: bool = dataclasses.field(default=None)
     threshold: float = dataclasses.field(default=None)
+    explain_predictions: bool = dataclasses.field(default=None)
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.ANOMALY_OUTLIERS
@@ -44,7 +46,8 @@ class ForecastingBatchPredictionArgs(BatchPredictionArgs):
        start_date_offset (int): Sets prediction start date as this offset relative to the prediction start date.
        forecasting_horizon (int): The number of timestamps to predict in the future. Range: [1, 1000].
        item_attributes_to_include_in_the_result (list): List of columns to include in the prediction output.
-       explain_predictions (bool): If True, explain predictions for the forecast.
+       explain_predictions (bool): If True, calculates explanations for the forecasted values along with predictions.
+       automate_monitoring (bool): If True, creates a monitor to calculate the drift for the batch prediction.
     """
     for_eval: bool = dataclasses.field(default=None)
     predictions_start_date: str = dataclasses.field(default=None)
@@ -53,6 +56,7 @@ class ForecastingBatchPredictionArgs(BatchPredictionArgs):
     forecasting_horizon: int = dataclasses.field(default=None)
     item_attributes_to_include_in_the_result: list = dataclasses.field(default=None)
     explain_predictions: bool = dataclasses.field(default=None)
+    automate_monitoring: bool = dataclasses.field(default=None)
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.FORECASTING
@@ -110,6 +114,8 @@ class PredictiveModelingBatchPredictionArgs(BatchPredictionArgs):
        explanation_filter_upper_bound (float): If set explanations will be limited to predictions below this value, Range: [0, 1].
        bound_label (str): For classification problems specifies the label to which the explanation bounds are applied.
        output_columns (list): A list of column names to include in the prediction result.
+       explain_predictions (bool): If True, calculates explanations for the predicted values along with predictions.
+       automate_monitoring (bool): If True, creates a monitor to calculate the drift for the batch prediction.
     """
     for_eval: bool = dataclasses.field(default=None)
     explainer_type: enums.ExplainerType = dataclasses.field(default=None)
@@ -121,6 +127,8 @@ class PredictiveModelingBatchPredictionArgs(BatchPredictionArgs):
     explanation_filter_upper_bound: float = dataclasses.field(default=None)
     explanation_filter_label: str = dataclasses.field(default=None)
     output_columns: list = dataclasses.field(default=None)
+    explain_predictions: bool = dataclasses.field(default=None)
+    automate_monitoring: bool = dataclasses.field(default=None)
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.PREDICTIVE_MODELING
@@ -194,6 +202,21 @@ class ChatLLMBatchPredictionArgs(BatchPredictionArgs):
         self.problem_type = enums.ProblemType.CHAT_LLM
 
 
+@dataclasses.dataclass
+class TrainablePlugAndPlayBatchPredictionArgs(BatchPredictionArgs):
+    """
+    Batch Prediction Config for the TrainablePlugAndPlay problem type
+
+    Args:
+        automate_monitoring (bool): If True, creates a monitor to calculate the drift for the batch prediction.
+    """
+    for_eval: bool = dataclasses.field(default=None)
+    automate_monitoring: bool = dataclasses.field(default=None)
+
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.CUSTOM_ALGORITHM
+
+
 @dataclasses.dataclass
 class _BatchPredictionArgsFactory(_ApiClassFactory):
     config_abstract_class = BatchPredictionArgs
@@ -208,4 +231,5 @@ class _BatchPredictionArgsFactory(_ApiClassFactory):
         enums.ProblemType.SENTENCE_BOUNDARY_DETECTION: SentenceBoundaryDetectionBatchPredictionArgs,
         enums.ProblemType.THEME_ANALYSIS: ThemeAnalysisBatchPredictionArgs,
         enums.ProblemType.CHAT_LLM: ChatLLMBatchPredictionArgs,
+        enums.ProblemType.CUSTOM_ALGORITHM: TrainablePlugAndPlayBatchPredictionArgs,
     }

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/api_class/document_retriever.py

@@ -14,11 +14,13 @@ class VectorStoreConfig(ApiClass):
         chunk_overlap_fraction (float): The fraction of overlap between chunks.
         text_encoder (VectorStoreTextEncoder): Encoder used to index texts from the documents.
         chunk_size_factors (list): Chunking data with multiple sizes. The specified list of factors are used to calculate more sizes, in addition to `chunk_size`.
+        score_multiplier_column (str): If provided, will use the values in this metadata column to modify the relevance score of returned chunks for all queries.
     """
     chunk_size: int = dataclasses.field(default=None)
     chunk_overlap_fraction: float = dataclasses.field(default=None)
     text_encoder: VectorStoreTextEncoder = dataclasses.field(default=None)
     chunk_size_factors: list = dataclasses.field(default=None)
+    score_multiplier_column: str = dataclasses.field(default=None)
 
 
 @dataclasses.dataclass

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/api_class/enums.py

@@ -350,6 +350,7 @@ class ApplicationConnectorType(ApiEnum):
     SHAREPOINT = 'SHAREPOINT'
     TEAMS = 'TEAMS'
     ABACUSUSAGEMETRICS = 'ABACUSUSAGEMETRICS'
+    MICROSOFTAUTH = 'MICROSOFTAUTH'
 
 
 class PythonFunctionArgumentType(ApiEnum):
@@ -476,3 +477,22 @@ class OcrMode(ApiEnum):
     COMPREHENSIVE = 'COMPREHENSIVE'
     COMPREHENSIVE_V2 = 'COMPREHENSIVE_V2'
     COMPREHENSIVE_TABLE_MD = 'COMPREHENSIVE_TABLE_MD'
+
+
+class DataType(ApiEnum):
+    INTEGER = 'integer'
+    FLOAT = 'float'
+    STRING = 'string'
+    DATE = 'date'
+    DATETIME = 'datetime'
+    BOOLEAN = 'boolean'
+    LIST = 'list'
+    STRUCT = 'struct'
+    NULL = 'null'
+
+    @classmethod
+    def from_str(cls, value):
+        if not value:
+            return None
+        default_map = {val.value: val for val in DataType}
+        return default_map[value.lower()]

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/api_class/model.py

@@ -453,6 +453,7 @@ class ChatLLMTrainingConfig(TrainingConfig):
         hide_generated_sql (bool): When running data queries, hides the generated SQL in the response and will just return the table.
         disable_data_summarization (bool): After executing a query summarize the reponse and reply back with only the table and query run.
         search_score_cutoff (float): Minimum search score to consider a document as a valid search result.
+        database_connector_id (str): Database connector id to use for the chatllm.
     """
     document_retrievers: List[str] = None
     num_completion_tokens: int = None
@@ -468,6 +469,7 @@ class ChatLLMTrainingConfig(TrainingConfig):
     hide_generated_sql: bool = None
     disable_data_summarization: bool = None
     search_score_cutoff: float = None
+    database_connector_id: str = None
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.CHAT_LLM

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/batch_prediction.py

@@ -21,11 +21,11 @@ class BatchPrediction(AbstractApiClass):
             fileConnectorOutputLocation (str): Contains information about where the batch predictions are written to.
             databaseConnectorId (str): The database connector to write the results to.
             databaseOutputConfiguration (dict): Contains information about where the batch predictions are written to.
-            explanations (bool): If true, explanations for each prediction will be created.
+            explanations (bool): 
             fileOutputFormat (str): The format of the batch prediction output (CSV or JSON).
             connectorType (str): Null if writing to internal console, else FEATURE_GROUP | FILE_CONNECTOR | DATABASE_CONNECTOR.
             legacyInputLocation (str): The location of the input data.
-            outputFeatureGroupId (str): The BP output feature group id if applicable
+            outputFeatureGroupId (str): The Batch Prediction output feature group ID if applicable
             featureGroupTableName (str): The table name of the Batch Prediction output feature group.
             outputFeatureGroupTableName (str): The table name of the Batch Prediction output feature group.
             summaryFeatureGroupTableName (str): The table name of the metrics summary feature group output by Batch Prediction.
@@ -43,7 +43,7 @@ class BatchPrediction(AbstractApiClass):
             inputFeatureGroups (PredictionFeatureGroup): List of prediction feature groups.
     """
 
-    def __init__(self, client, batchPredictionId=None, createdAt=None, name=None, deploymentId=None, fileConnectorOutputLocation=None, databaseConnectorId=None, databaseOutputConfiguration=None, explanations=None, fileOutputFormat=None, connectorType=None, legacyInputLocation=None, outputFeatureGroupId=None, featureGroupTableName=None, outputFeatureGroupTableName=None, summaryFeatureGroupTableName=None, csvInputPrefix=None, csvPredictionPrefix=None, csvExplanationsPrefix=None, outputIncludesMetadata=None, resultInputColumns=None, modelMonitorId=None, modelVersion=None, algorithm=None, batchInputs={}, latestBatchPredictionVersion={}, refreshSchedules={}, inputFeatureGroups={}, globalPredictionArgs={}):
+    def __init__(self, client, batchPredictionId=None, createdAt=None, name=None, deploymentId=None, fileConnectorOutputLocation=None, databaseConnectorId=None, databaseOutputConfiguration=None, explanations=None, fileOutputFormat=None, connectorType=None, legacyInputLocation=None, outputFeatureGroupId=None, featureGroupTableName=None, outputFeatureGroupTableName=None, summaryFeatureGroupTableName=None, csvInputPrefix=None, csvPredictionPrefix=None, csvExplanationsPrefix=None, outputIncludesMetadata=None, resultInputColumns=None, modelMonitorId=None, modelVersion=None, algorithm=None, batchInputs={}, latestBatchPredictionVersion={}, refreshSchedules={}, inputFeatureGroups={}, globalPredictionArgs={}, batchPredictionArgs={}):
         super().__init__(client, batchPredictionId)
         self.batch_prediction_id = batchPredictionId
         self.created_at = createdAt
@@ -77,11 +77,13 @@ class BatchPrediction(AbstractApiClass):
             PredictionFeatureGroup, inputFeatureGroups)
         self.global_prediction_args = client._build_class(
             BatchPredictionArgs, globalPredictionArgs)
-        self.deprecated_keys = {}
+        self.batch_prediction_args = client._build_class(
+            BatchPredictionArgs, batchPredictionArgs)
+        self.deprecated_keys = {'global_prediction_args', 'explanations'}
 
     def __repr__(self):
-        repr_dict = {f'batch_prediction_id': repr(self.batch_prediction_id), f'created_at': repr(self.created_at), f'name': repr(self.name), f'deployment_id': repr(self.deployment_id), f'file_connector_output_location': repr(self.file_connector_output_location), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'output_feature_group_id': repr(self.output_feature_group_id), f'feature_group_table_name': repr(self.feature_group_table_name), f'output_feature_group_table_name': repr(
-            self.output_feature_group_table_name), f'summary_feature_group_table_name': repr(self.summary_feature_group_table_name), f'csv_input_prefix': repr(self.csv_input_prefix), f'csv_prediction_prefix': repr(self.csv_prediction_prefix), f'csv_explanations_prefix': repr(self.csv_explanations_prefix), f'output_includes_metadata': repr(self.output_includes_metadata), f'result_input_columns': repr(self.result_input_columns), f'model_monitor_id': repr(self.model_monitor_id), f'model_version': repr(self.model_version), f'algorithm': repr(self.algorithm), f'batch_inputs': repr(self.batch_inputs), f'latest_batch_prediction_version': repr(self.latest_batch_prediction_version), f'refresh_schedules': repr(self.refresh_schedules), f'input_feature_groups': repr(self.input_feature_groups), f'global_prediction_args': repr(self.global_prediction_args)}
+        repr_dict = {f'batch_prediction_id': repr(self.batch_prediction_id), f'created_at': repr(self.created_at), f'name': repr(self.name), f'deployment_id': repr(self.deployment_id), f'file_connector_output_location': repr(self.file_connector_output_location), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'output_feature_group_id': repr(self.output_feature_group_id), f'feature_group_table_name': repr(self.feature_group_table_name), f'output_feature_group_table_name': repr(self.output_feature_group_table_name), f'summary_feature_group_table_name': repr(
+            self.summary_feature_group_table_name), f'csv_input_prefix': repr(self.csv_input_prefix), f'csv_prediction_prefix': repr(self.csv_prediction_prefix), f'csv_explanations_prefix': repr(self.csv_explanations_prefix), f'output_includes_metadata': repr(self.output_includes_metadata), f'result_input_columns': repr(self.result_input_columns), f'model_monitor_id': repr(self.model_monitor_id), f'model_version': repr(self.model_version), f'algorithm': repr(self.algorithm), f'batch_inputs': repr(self.batch_inputs), f'latest_batch_prediction_version': repr(self.latest_batch_prediction_version), f'refresh_schedules': repr(self.refresh_schedules), f'input_feature_groups': repr(self.input_feature_groups), f'global_prediction_args': repr(self.global_prediction_args), f'batch_prediction_args': repr(self.batch_prediction_args)}
         class_name = "BatchPrediction"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None and key not in self.deprecated_keys])
@@ -94,8 +96,8 @@ class BatchPrediction(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'batch_prediction_id': self.batch_prediction_id, 'created_at': self.created_at, 'name': self.name, 'deployment_id': self.deployment_id, 'file_connector_output_location': self.file_connector_output_location, 'database_connector_id': self.database_connector_id, 'database_output_configuration': self.database_output_configuration, 'explanations': self.explanations, 'file_output_format': self.file_output_format, 'connector_type': self.connector_type, 'legacy_input_location': self.legacy_input_location, 'output_feature_group_id': self.output_feature_group_id, 'feature_group_table_name': self.feature_group_table_name, 'output_feature_group_table_name': self.output_feature_group_table_name, 'summary_feature_group_table_name': self.summary_feature_group_table_name,
-                'csv_input_prefix': self.csv_input_prefix, 'csv_prediction_prefix': self.csv_prediction_prefix, 'csv_explanations_prefix': self.csv_explanations_prefix, 'output_includes_metadata': self.output_includes_metadata, 'result_input_columns': self.result_input_columns, 'model_monitor_id': self.model_monitor_id, 'model_version': self.model_version, 'algorithm': self.algorithm, 'batch_inputs': self._get_attribute_as_dict(self.batch_inputs), 'latest_batch_prediction_version': self._get_attribute_as_dict(self.latest_batch_prediction_version), 'refresh_schedules': self._get_attribute_as_dict(self.refresh_schedules), 'input_feature_groups': self._get_attribute_as_dict(self.input_feature_groups), 'global_prediction_args': self._get_attribute_as_dict(self.global_prediction_args)}
+        resp = {'batch_prediction_id': self.batch_prediction_id, 'created_at': self.created_at, 'name': self.name, 'deployment_id': self.deployment_id, 'file_connector_output_location': self.file_connector_output_location, 'database_connector_id': self.database_connector_id, 'database_output_configuration': self.database_output_configuration, 'explanations': self.explanations, 'file_output_format': self.file_output_format, 'connector_type': self.connector_type, 'legacy_input_location': self.legacy_input_location, 'output_feature_group_id': self.output_feature_group_id, 'feature_group_table_name': self.feature_group_table_name, 'output_feature_group_table_name': self.output_feature_group_table_name, 'summary_feature_group_table_name': self.summary_feature_group_table_name, 'csv_input_prefix': self.csv_input_prefix,
+                'csv_prediction_prefix': self.csv_prediction_prefix, 'csv_explanations_prefix': self.csv_explanations_prefix, 'output_includes_metadata': self.output_includes_metadata, 'result_input_columns': self.result_input_columns, 'model_monitor_id': self.model_monitor_id, 'model_version': self.model_version, 'algorithm': self.algorithm, 'batch_inputs': self._get_attribute_as_dict(self.batch_inputs), 'latest_batch_prediction_version': self._get_attribute_as_dict(self.latest_batch_prediction_version), 'refresh_schedules': self._get_attribute_as_dict(self.refresh_schedules), 'input_feature_groups': self._get_attribute_as_dict(self.input_feature_groups), 'global_prediction_args': self._get_attribute_as_dict(self.global_prediction_args), 'batch_prediction_args': self._get_attribute_as_dict(self.batch_prediction_args)}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}
 
     def start(self):
@@ -145,14 +147,13 @@ class BatchPrediction(AbstractApiClass):
         """
         return self.client.list_batch_prediction_versions(self.batch_prediction_id, limit, start_after_version)
 
-    def update(self, deployment_id: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = None, output_format: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, name: str = None):
+    def update(self, deployment_id: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, batch_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = None, output_format: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, name: str = None):
         """
         Update a batch prediction job description.
 
         Args:
             deployment_id (str): Unique identifier of the deployment.
-            global_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
-            explanations (bool): If True, SHAP explanations for each prediction will be provided, if supported by the use case.
+            batch_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
             output_format (str): If specified, sets the format of the batch prediction output (CSV or JSON).
             csv_input_prefix (str): Prefix to prepend to the input columns, only applies when output format is CSV.
             csv_prediction_prefix (str): Prefix to prepend to the prediction columns, only applies when output format is CSV.
@@ -164,7 +165,7 @@ class BatchPrediction(AbstractApiClass):
         Returns:
             BatchPrediction: The batch prediction.
         """
-        return self.client.update_batch_prediction(self.batch_prediction_id, deployment_id, global_prediction_args, explanations, output_format, csv_input_prefix, csv_prediction_prefix, csv_explanations_prefix, output_includes_metadata, result_input_columns, name)
+        return self.client.update_batch_prediction(self.batch_prediction_id, deployment_id, global_prediction_args, batch_prediction_args, explanations, output_format, csv_input_prefix, csv_prediction_prefix, csv_explanations_prefix, output_includes_metadata, result_input_columns, name)
 
     def set_file_connector_output(self, output_format: str = None, output_location: str = None):
         """

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/batch_prediction_version.py

@@ -24,7 +24,7 @@ class BatchPredictionVersion(AbstractApiClass):
             failedPredictions (int): Number of predictions that failed
             databaseConnectorId (str): The database connector to write the results to
             databaseOutputConfiguration (dict): Contains information about where the batch predictions are written to
-            explanations (bool): If true, explanations for each prediction were created
+            explanations (bool): 
             fileConnectorOutputLocation (str):  Contains information about where the batch predictions are written to
             fileOutputFormat (str): The format of the batch prediction output (CSV or JSON)
             connectorType (str): Null if writing to internal console, else FEATURE_GROUP | FILE_CONNECTOR | DATABASE_CONNECTOR
@@ -42,15 +42,15 @@ class BatchPredictionVersion(AbstractApiClass):
             modelMonitorVersion (str): The version of the model monitor
             algoName (str): The name of the algorithm used to train the model
             algorithm (str): The algorithm that is currently deployed.
-            outputFeatureGroupId (str): The BP output feature group id if applicable
-            outputFeatureGroupVersion (str): The BP output feature group version if applicable
-            outputFeatureGroupTableName (str): The BP output feature group name if applicable
+            outputFeatureGroupId (str): The Batch Prediction output feature group ID if applicable
+            outputFeatureGroupVersion (str): The Batch Prediction output feature group version if applicable
+            outputFeatureGroupTableName (str): The Batch Prediction output feature group name if applicable
             batchPredictionWarnings (str): Relevant warnings if any issues are found
             batchInputs (PredictionInput): Inputs to the batch prediction
             inputFeatureGroups (PredictionFeatureGroup): List of prediction feature groups
     """
 
-    def __init__(self, client, batchPredictionVersion=None, batchPredictionId=None, status=None, driftMonitorStatus=None, deploymentId=None, modelId=None, modelVersion=None, predictionsStartedAt=None, predictionsCompletedAt=None, databaseOutputError=None, totalPredictions=None, failedPredictions=None, databaseConnectorId=None, databaseOutputConfiguration=None, explanations=None, fileConnectorOutputLocation=None, fileOutputFormat=None, connectorType=None, legacyInputLocation=None, error=None, driftMonitorError=None, monitorWarnings=None, csvInputPrefix=None, csvPredictionPrefix=None, csvExplanationsPrefix=None, databaseOutputTotalWrites=None, databaseOutputFailedWrites=None, outputIncludesMetadata=None, resultInputColumns=None, modelMonitorVersion=None, algoName=None, algorithm=None, outputFeatureGroupId=None, outputFeatureGroupVersion=None, outputFeatureGroupTableName=None, batchPredictionWarnings=None, batchInputs={}, inputFeatureGroups={}, globalPredictionArgs={}):
+    def __init__(self, client, batchPredictionVersion=None, batchPredictionId=None, status=None, driftMonitorStatus=None, deploymentId=None, modelId=None, modelVersion=None, predictionsStartedAt=None, predictionsCompletedAt=None, databaseOutputError=None, totalPredictions=None, failedPredictions=None, databaseConnectorId=None, databaseOutputConfiguration=None, explanations=None, fileConnectorOutputLocation=None, fileOutputFormat=None, connectorType=None, legacyInputLocation=None, error=None, driftMonitorError=None, monitorWarnings=None, csvInputPrefix=None, csvPredictionPrefix=None, csvExplanationsPrefix=None, databaseOutputTotalWrites=None, databaseOutputFailedWrites=None, outputIncludesMetadata=None, resultInputColumns=None, modelMonitorVersion=None, algoName=None, algorithm=None, outputFeatureGroupId=None, outputFeatureGroupVersion=None, outputFeatureGroupTableName=None, batchPredictionWarnings=None, batchInputs={}, inputFeatureGroups={}, globalPredictionArgs={}, batchPredictionArgs={}):
         super().__init__(client, batchPredictionVersion)
         self.batch_prediction_version = batchPredictionVersion
         self.batch_prediction_id = batchPredictionId
@@ -93,11 +93,13 @@ class BatchPredictionVersion(AbstractApiClass):
             PredictionFeatureGroup, inputFeatureGroups)
         self.global_prediction_args = client._build_class(
             BatchPredictionArgs, globalPredictionArgs)
-        self.deprecated_keys = {}
+        self.batch_prediction_args = client._build_class(
+            BatchPredictionArgs, batchPredictionArgs)
+        self.deprecated_keys = {'global_prediction_args', 'explanations'}
 
     def __repr__(self):
-        repr_dict = {f'batch_prediction_version': repr(self.batch_prediction_version), f'batch_prediction_id': repr(self.batch_prediction_id), f'status': repr(self.status), f'drift_monitor_status': repr(self.drift_monitor_status), f'deployment_id': repr(self.deployment_id), f'model_id': repr(self.model_id), f'model_version': repr(self.model_version), f'predictions_started_at': repr(self.predictions_started_at), f'predictions_completed_at': repr(self.predictions_completed_at), f'database_output_error': repr(self.database_output_error), f'total_predictions': repr(self.total_predictions), f'failed_predictions': repr(self.failed_predictions), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_connector_output_location': repr(self.file_connector_output_location), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'error': repr(self.error), f'drift_monitor_error': repr(
-            self.drift_monitor_error), f'monitor_warnings': repr(self.monitor_warnings), f'csv_input_prefix': repr(self.csv_input_prefix), f'csv_prediction_prefix': repr(self.csv_prediction_prefix), f'csv_explanations_prefix': repr(self.csv_explanations_prefix), f'database_output_total_writes': repr(self.database_output_total_writes), f'database_output_failed_writes': repr(self.database_output_failed_writes), f'output_includes_metadata': repr(self.output_includes_metadata), f'result_input_columns': repr(self.result_input_columns), f'model_monitor_version': repr(self.model_monitor_version), f'algo_name': repr(self.algo_name), f'algorithm': repr(self.algorithm), f'output_feature_group_id': repr(self.output_feature_group_id), f'output_feature_group_version': repr(self.output_feature_group_version), f'output_feature_group_table_name': repr(self.output_feature_group_table_name), f'batch_prediction_warnings': repr(self.batch_prediction_warnings), f'batch_inputs': repr(self.batch_inputs), f'input_feature_groups': repr(self.input_feature_groups), f'global_prediction_args': repr(self.global_prediction_args)}
+        repr_dict = {f'batch_prediction_version': repr(self.batch_prediction_version), f'batch_prediction_id': repr(self.batch_prediction_id), f'status': repr(self.status), f'drift_monitor_status': repr(self.drift_monitor_status), f'deployment_id': repr(self.deployment_id), f'model_id': repr(self.model_id), f'model_version': repr(self.model_version), f'predictions_started_at': repr(self.predictions_started_at), f'predictions_completed_at': repr(self.predictions_completed_at), f'database_output_error': repr(self.database_output_error), f'total_predictions': repr(self.total_predictions), f'failed_predictions': repr(self.failed_predictions), f'database_connector_id': repr(self.database_connector_id), f'database_output_configuration': repr(self.database_output_configuration), f'explanations': repr(self.explanations), f'file_connector_output_location': repr(self.file_connector_output_location), f'file_output_format': repr(self.file_output_format), f'connector_type': repr(self.connector_type), f'legacy_input_location': repr(self.legacy_input_location), f'error': repr(self.error), f'drift_monitor_error': repr(self.drift_monitor_error), f'monitor_warnings': repr(
+            self.monitor_warnings), f'csv_input_prefix': repr(self.csv_input_prefix), f'csv_prediction_prefix': repr(self.csv_prediction_prefix), f'csv_explanations_prefix': repr(self.csv_explanations_prefix), f'database_output_total_writes': repr(self.database_output_total_writes), f'database_output_failed_writes': repr(self.database_output_failed_writes), f'output_includes_metadata': repr(self.output_includes_metadata), f'result_input_columns': repr(self.result_input_columns), f'model_monitor_version': repr(self.model_monitor_version), f'algo_name': repr(self.algo_name), f'algorithm': repr(self.algorithm), f'output_feature_group_id': repr(self.output_feature_group_id), f'output_feature_group_version': repr(self.output_feature_group_version), f'output_feature_group_table_name': repr(self.output_feature_group_table_name), f'batch_prediction_warnings': repr(self.batch_prediction_warnings), f'batch_inputs': repr(self.batch_inputs), f'input_feature_groups': repr(self.input_feature_groups), f'global_prediction_args': repr(self.global_prediction_args), f'batch_prediction_args': repr(self.batch_prediction_args)}
         class_name = "BatchPredictionVersion"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None and key not in self.deprecated_keys])
@@ -110,8 +112,8 @@ class BatchPredictionVersion(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'batch_prediction_version': self.batch_prediction_version, 'batch_prediction_id': self.batch_prediction_id, 'status': self.status, 'drift_monitor_status': self.drift_monitor_status, 'deployment_id': self.deployment_id, 'model_id': self.model_id, 'model_version': self.model_version, 'predictions_started_at': self.predictions_started_at, 'predictions_completed_at': self.predictions_completed_at, 'database_output_error': self.database_output_error, 'total_predictions': self.total_predictions, 'failed_predictions': self.failed_predictions, 'database_connector_id': self.database_connector_id, 'database_output_configuration': self.database_output_configuration, 'explanations': self.explanations, 'file_connector_output_location': self.file_connector_output_location, 'file_output_format': self.file_output_format, 'connector_type': self.connector_type, 'legacy_input_location': self.legacy_input_location, 'error': self.error, 'drift_monitor_error': self.drift_monitor_error, 'monitor_warnings': self.monitor_warnings,
-                'csv_input_prefix': self.csv_input_prefix, 'csv_prediction_prefix': self.csv_prediction_prefix, 'csv_explanations_prefix': self.csv_explanations_prefix, 'database_output_total_writes': self.database_output_total_writes, 'database_output_failed_writes': self.database_output_failed_writes, 'output_includes_metadata': self.output_includes_metadata, 'result_input_columns': self.result_input_columns, 'model_monitor_version': self.model_monitor_version, 'algo_name': self.algo_name, 'algorithm': self.algorithm, 'output_feature_group_id': self.output_feature_group_id, 'output_feature_group_version': self.output_feature_group_version, 'output_feature_group_table_name': self.output_feature_group_table_name, 'batch_prediction_warnings': self.batch_prediction_warnings, 'batch_inputs': self._get_attribute_as_dict(self.batch_inputs), 'input_feature_groups': self._get_attribute_as_dict(self.input_feature_groups), 'global_prediction_args': self._get_attribute_as_dict(self.global_prediction_args)}
+        resp = {'batch_prediction_version': self.batch_prediction_version, 'batch_prediction_id': self.batch_prediction_id, 'status': self.status, 'drift_monitor_status': self.drift_monitor_status, 'deployment_id': self.deployment_id, 'model_id': self.model_id, 'model_version': self.model_version, 'predictions_started_at': self.predictions_started_at, 'predictions_completed_at': self.predictions_completed_at, 'database_output_error': self.database_output_error, 'total_predictions': self.total_predictions, 'failed_predictions': self.failed_predictions, 'database_connector_id': self.database_connector_id, 'database_output_configuration': self.database_output_configuration, 'explanations': self.explanations, 'file_connector_output_location': self.file_connector_output_location, 'file_output_format': self.file_output_format, 'connector_type': self.connector_type, 'legacy_input_location': self.legacy_input_location, 'error': self.error, 'drift_monitor_error': self.drift_monitor_error, 'monitor_warnings': self.monitor_warnings, 'csv_input_prefix': self.csv_input_prefix,
+                'csv_prediction_prefix': self.csv_prediction_prefix, 'csv_explanations_prefix': self.csv_explanations_prefix, 'database_output_total_writes': self.database_output_total_writes, 'database_output_failed_writes': self.database_output_failed_writes, 'output_includes_metadata': self.output_includes_metadata, 'result_input_columns': self.result_input_columns, 'model_monitor_version': self.model_monitor_version, 'algo_name': self.algo_name, 'algorithm': self.algorithm, 'output_feature_group_id': self.output_feature_group_id, 'output_feature_group_version': self.output_feature_group_version, 'output_feature_group_table_name': self.output_feature_group_table_name, 'batch_prediction_warnings': self.batch_prediction_warnings, 'batch_inputs': self._get_attribute_as_dict(self.batch_inputs), 'input_feature_groups': self._get_attribute_as_dict(self.input_feature_groups), 'global_prediction_args': self._get_attribute_as_dict(self.global_prediction_args), 'batch_prediction_args': self._get_attribute_as_dict(self.batch_prediction_args)}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}
 
     def download_batch_prediction_result_chunk(self, offset: int = 0, chunk_size: int = 10485760):

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/client.py

@@ -33,11 +33,11 @@ from .annotation_entry import AnnotationEntry
 from .annotations_status import AnnotationsStatus
 from .api_class import (
     AlertActionConfig, AlertConditionConfig, ApiClass, ApiEnum,
-    BatchPredictionArgs, BlobInput, DatasetConfig, DocumentProcessingConfig,
-    DocumentRetrieverConfig, EvalArtifactType, FeatureGroupExportConfig,
-    ForecastingMonitorConfig, IncrementalDatabaseConnectorConfig, LLMName,
-    MergeConfig, ParsingConfig, ProblemType, PythonFunctionType,
-    SamplingConfig, TrainingConfig
+    BatchPredictionArgs, BlobInput, DatasetConfig, DataType,
+    DocumentProcessingConfig, DocumentRetrieverConfig, EvalArtifactType,
+    FeatureGroupExportConfig, ForecastingMonitorConfig,
+    IncrementalDatabaseConnectorConfig, LLMName, MergeConfig, ParsingConfig,
+    ProblemType, PythonFunctionType, SamplingConfig, TrainingConfig
 )
 from .api_client_utils import (
     INVALID_PANDAS_COLUMN_NAME_CHARACTERS, StreamingHandler, clean_column_name,
@@ -575,7 +575,7 @@ class BaseApiClient:
         client_options (ClientOptions): Optional API client configurations
         skip_version_check (bool): If true, will skip checking the server's current API version on initializing the client
     """
-    client_version = '1.1.5'
+    client_version = '1.1.6'
 
     def __init__(self, api_key: str = None, server: str = None, client_options: ClientOptions = None, skip_version_check: bool = False):
         self.api_key = api_key
@@ -5026,13 +5026,13 @@ Creates a new feature group defined as the union of other feature group versions
             DatasetVersion: The new Dataset Version created by taking a snapshot of the current data in the streaming dataset."""
         return self._call_api('snapshotStreamingData', 'POST', query_params={'datasetId': dataset_id}, body={}, parse_type=DatasetVersion)
 
-    def set_dataset_column_data_type(self, dataset_id: str, column: str, data_type: str) -> Dataset:
+    def set_dataset_column_data_type(self, dataset_id: str, column: str, data_type: Union[DataType, str]) -> Dataset:
         """Set a Dataset's column type.
 
         Args:
             dataset_id (str): The unique ID associated with the dataset.
             column (str): The name of the column.
-            data_type (str): The type of the data in the column. Note: Some ColumnMappings may restrict the options or explicitly set the DataType.
+            data_type (DataType): The type of the data in the column. Note: Some ColumnMappings may restrict the options or explicitly set the DataType.
 
         Returns:
             Dataset: The dataset and schema after the data type has been set."""
@@ -6717,15 +6717,14 @@ Creates a new feature group defined as the union of other feature group versions
             deployment_id, deployment_token) if deployment_token else None
         return self._call_api('executeAgentWithBinaryData', 'POST', query_params={'deploymentToken': deployment_token, 'deploymentId': deployment_id}, data={'arguments': json.dumps(arguments) if (arguments is not None and not isinstance(arguments, str)) else arguments, 'keywordArguments': json.dumps(keyword_arguments) if (keyword_arguments is not None and not isinstance(keyword_arguments, str)) else keyword_arguments, 'deploymentConversationId': json.dumps(deployment_conversation_id) if (deployment_conversation_id is not None and not isinstance(deployment_conversation_id, str)) else deployment_conversation_id, 'externalSessionId': json.dumps(external_session_id) if (external_session_id is not None and not isinstance(external_session_id, str)) else external_session_id}, parse_type=AgentDataExecutionResult, files=blobs, server_override=prediction_url, timeout=1500)
 
-    def create_batch_prediction(self, deployment_id: str, table_name: str = None, name: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = False, output_format: str = None, output_location: str = None, database_connector_id: str = None, database_output_config: dict = None, refresh_schedule: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, input_feature_groups: dict = None) -> BatchPrediction:
+    def create_batch_prediction(self, deployment_id: str, table_name: str = None, name: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, batch_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = False, output_format: str = None, output_location: str = None, database_connector_id: str = None, database_output_config: dict = None, refresh_schedule: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, input_feature_groups: dict = None) -> BatchPrediction:
         """Creates a batch prediction job description for the given deployment.
 
         Args:
             deployment_id (str): Unique string identifier for the deployment.
             table_name (str): Name of the feature group table to write the results of the batch prediction. Can only be specified if outputLocation and databaseConnectorId are not specified. If tableName is specified, the outputType will be enforced as CSV.
             name (str): Name of the batch prediction job.
-            global_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
-            explanations (bool): If true, SHAP explanations will be provided for each prediction, if supported by the use case.
+            batch_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
             output_format (str): Format of the batch prediction output (CSV or JSON).
             output_location (str): Location to write the prediction results. Otherwise, results will be stored in Abacus.AI.
             database_connector_id (str): Unique identifier of a Database Connection to write predictions to. Cannot be specified in conjunction with outputLocation.
@@ -6740,7 +6739,7 @@ Creates a new feature group defined as the union of other feature group versions
 
         Returns:
             BatchPrediction: The batch prediction description."""
-        return self._call_api('createBatchPrediction', 'POST', query_params={'deploymentId': deployment_id}, body={'tableName': table_name, 'name': name, 'globalPredictionArgs': global_prediction_args, 'explanations': explanations, 'outputFormat': output_format, 'outputLocation': output_location, 'databaseConnectorId': database_connector_id, 'databaseOutputConfig': database_output_config, 'refreshSchedule': refresh_schedule, 'csvInputPrefix': csv_input_prefix, 'csvPredictionPrefix': csv_prediction_prefix, 'csvExplanationsPrefix': csv_explanations_prefix, 'outputIncludesMetadata': output_includes_metadata, 'resultInputColumns': result_input_columns, 'inputFeatureGroups': input_feature_groups}, parse_type=BatchPrediction)
+        return self._call_api('createBatchPrediction', 'POST', query_params={'deploymentId': deployment_id}, body={'tableName': table_name, 'name': name, 'globalPredictionArgs': global_prediction_args, 'batchPredictionArgs': batch_prediction_args, 'explanations': explanations, 'outputFormat': output_format, 'outputLocation': output_location, 'databaseConnectorId': database_connector_id, 'databaseOutputConfig': database_output_config, 'refreshSchedule': refresh_schedule, 'csvInputPrefix': csv_input_prefix, 'csvPredictionPrefix': csv_prediction_prefix, 'csvExplanationsPrefix': csv_explanations_prefix, 'outputIncludesMetadata': output_includes_metadata, 'resultInputColumns': result_input_columns, 'inputFeatureGroups': input_feature_groups}, parse_type=BatchPrediction)
 
     def start_batch_prediction(self, batch_prediction_id: str) -> BatchPredictionVersion:
         """Creates a new batch prediction version job for a given batch prediction job description.
@@ -6752,14 +6751,13 @@ Creates a new feature group defined as the union of other feature group versions
             BatchPredictionVersion: The batch prediction version started by this method call."""
         return self._call_api('startBatchPrediction', 'POST', query_params={}, body={'batchPredictionId': batch_prediction_id}, parse_type=BatchPredictionVersion)
 
-    def update_batch_prediction(self, batch_prediction_id: str, deployment_id: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = None, output_format: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, name: str = None) -> BatchPrediction:
+    def update_batch_prediction(self, batch_prediction_id: str, deployment_id: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, batch_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = None, output_format: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, name: str = None) -> BatchPrediction:
         """Update a batch prediction job description.
 
         Args:
             batch_prediction_id (str): Unique identifier of the batch prediction.
             deployment_id (str): Unique identifier of the deployment.
-            global_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
-            explanations (bool): If True, SHAP explanations for each prediction will be provided, if supported by the use case.
+            batch_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
             output_format (str): If specified, sets the format of the batch prediction output (CSV or JSON).
             csv_input_prefix (str): Prefix to prepend to the input columns, only applies when output format is CSV.
             csv_prediction_prefix (str): Prefix to prepend to the prediction columns, only applies when output format is CSV.
@@ -6770,7 +6768,7 @@ Creates a new feature group defined as the union of other feature group versions
 
         Returns:
             BatchPrediction: The batch prediction."""
-        return self._call_api('updateBatchPrediction', 'POST', query_params={'deploymentId': deployment_id}, body={'batchPredictionId': batch_prediction_id, 'globalPredictionArgs': global_prediction_args, 'explanations': explanations, 'outputFormat': output_format, 'csvInputPrefix': csv_input_prefix, 'csvPredictionPrefix': csv_prediction_prefix, 'csvExplanationsPrefix': csv_explanations_prefix, 'outputIncludesMetadata': output_includes_metadata, 'resultInputColumns': result_input_columns, 'name': name}, parse_type=BatchPrediction)
+        return self._call_api('updateBatchPrediction', 'POST', query_params={'deploymentId': deployment_id}, body={'batchPredictionId': batch_prediction_id, 'globalPredictionArgs': global_prediction_args, 'batchPredictionArgs': batch_prediction_args, 'explanations': explanations, 'outputFormat': output_format, 'csvInputPrefix': csv_input_prefix, 'csvPredictionPrefix': csv_prediction_prefix, 'csvExplanationsPrefix': csv_explanations_prefix, 'outputIncludesMetadata': output_includes_metadata, 'resultInputColumns': result_input_columns, 'name': name}, parse_type=BatchPrediction)
 
     def set_batch_prediction_file_connector_output(self, batch_prediction_id: str, output_format: str = None, output_location: str = None) -> BatchPrediction:
         """Updates the file connector output configuration of the batch prediction
@@ -6983,19 +6981,7 @@ Creates a new feature group defined as the union of other feature group versions
             streaming_token, feature_group_id=feature_group_id)
         return self._call_api('appendMultipleData', 'POST', query_params={'streamingToken': streaming_token}, body={'featureGroupId': feature_group_id, 'data': data}, server_override=prediction_url)
 
-    def upsert_data(self, feature_group_id: str, streaming_token: str = None, data: dict = None) -> FeatureGroupRow:
-        """Update new data into the feature group for a given lookup key record ID if the record ID is found; otherwise, insert new data into the feature group.
-
-        Args:
-            feature_group_id (str): A unique string identifier of the streaming feature group to record data to.
-            streaming_token (str): Optional streaming token for authenticating requests if upserting to streaming FG.
-            data (dict): The data to record, in JSON format.
-
-        Returns:
-            FeatureGroupRow: The feature group row that was upserted."""
-        return self._proxy_request('upsertData', 'POST', query_params={'streamingToken': streaming_token}, body={'featureGroupId': feature_group_id, 'data': data}, parse_type=FeatureGroupRow, is_sync=True)
-
-    def upsert_online_data(self, feature_group_id: str, data: dict, streaming_token: str = None) -> FeatureGroupRow:
+    def upsert_data(self, feature_group_id: str, data: dict, streaming_token: str = None) -> FeatureGroupRow:
         """Update new data into the feature group for a given lookup key record ID if the record ID is found; otherwise, insert new data into the feature group.
 
         Args:
@@ -7005,7 +6991,7 @@ Creates a new feature group defined as the union of other feature group versions
 
         Returns:
             FeatureGroupRow: The feature group row that was upserted."""
-        return self._proxy_request('upsertOnlineData', 'POST', query_params={'streamingToken': streaming_token}, body={'featureGroupId': feature_group_id, 'data': data}, parse_type=FeatureGroupRow, is_sync=True)
+        return self._proxy_request('upsertData', 'POST', query_params={'streamingToken': streaming_token}, body={'featureGroupId': feature_group_id, 'data': data}, parse_type=FeatureGroupRow, is_sync=True)
 
     def delete_data(self, feature_group_id: str, primary_key: str):
         """Deletes a row from the feature group given the primary key
@@ -7764,7 +7750,7 @@ Creates a new feature group defined as the union of other feature group versions
             ExternalApplication: The newly created External Application."""
         return self._call_api('createExternalApplication', 'POST', query_params={'deploymentId': deployment_id}, body={'name': name, 'logo': logo, 'theme': theme}, parse_type=ExternalApplication)
 
-    def update_external_application(self, external_application_id: str, name: str = None, theme: dict = None, deployment_id: str = None) -> ExternalApplication:
+    def update_external_application(self, external_application_id: str, name: str = None, theme: dict = None, deployment_id: str = None, deployment_conversation_retention_hours: int = None, reset_retention_policy: bool = False) -> ExternalApplication:
         """Updates an External Application.
 
         Args:
@@ -7772,10 +7758,12 @@ Creates a new feature group defined as the union of other feature group versions
             name (str): The name of the External Application.
             theme (dict): The visual theme of the External Application.
             deployment_id (str): The ID of the deployment to use.
+            deployment_conversation_retention_hours (int): The number of hours to retain the conversations for.
+            reset_retention_policy (bool): If true, the retention policy will be removed.
 
         Returns:
             ExternalApplication: The updated External Application."""
-        return self._call_api('updateExternalApplication', 'POST', query_params={'deploymentId': deployment_id}, body={'externalApplicationId': external_application_id, 'name': name, 'theme': theme}, parse_type=ExternalApplication)
+        return self._call_api('updateExternalApplication', 'POST', query_params={'deploymentId': deployment_id}, body={'externalApplicationId': external_application_id, 'name': name, 'theme': theme, 'deploymentConversationRetentionHours': deployment_conversation_retention_hours, 'resetRetentionPolicy': reset_retention_policy}, parse_type=ExternalApplication)
 
     def delete_external_application(self, external_application_id: str):
         """Deletes an External Application.

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/dataset.py

@@ -1,6 +1,6 @@
 from typing import Union
 
-from .api_class import DatasetConfig, DocumentProcessingConfig, ParsingConfig
+from .api_class import DatasetConfig, DataType, DocumentProcessingConfig, ParsingConfig
 from .dataset_column import DatasetColumn
 from .dataset_version import DatasetVersion
 from .refresh_schedule import RefreshSchedule
@@ -152,13 +152,13 @@ class Dataset(AbstractApiClass):
         """
         return self.client.snapshot_streaming_data(self.dataset_id)
 
-    def set_column_data_type(self, column: str, data_type: str):
+    def set_column_data_type(self, column: str, data_type: Union[dict, DataType]):
         """
         Set a Dataset's column type.
 
         Args:
             column (str): The name of the column.
-            data_type (str): The type of the data in the column. Note: Some ColumnMappings may restrict the options or explicitly set the DataType.
+            data_type (DataType): The type of the data in the column. Note: Some ColumnMappings may restrict the options or explicitly set the DataType.
 
         Returns:
             Dataset: The dataset and schema after the data type has been set.

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/deployment.py

@@ -334,15 +334,14 @@ class Deployment(AbstractApiClass):
         """
         return self.client.get_conversation_response_with_binary_data(self.deployment_id, deployment_token, message, deployment_conversation_id, external_session_id, llm_name, num_completion_tokens, system_message, temperature, filter_key_values, search_score_cutoff, chat_config, ignore_documents, attachments)
 
-    def create_batch_prediction(self, table_name: str = None, name: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = False, output_format: str = None, output_location: str = None, database_connector_id: str = None, database_output_config: dict = None, refresh_schedule: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, input_feature_groups: dict = None):
+    def create_batch_prediction(self, table_name: str = None, name: str = None, global_prediction_args: Union[dict, BatchPredictionArgs] = None, batch_prediction_args: Union[dict, BatchPredictionArgs] = None, explanations: bool = False, output_format: str = None, output_location: str = None, database_connector_id: str = None, database_output_config: dict = None, refresh_schedule: str = None, csv_input_prefix: str = None, csv_prediction_prefix: str = None, csv_explanations_prefix: str = None, output_includes_metadata: bool = None, result_input_columns: list = None, input_feature_groups: dict = None):
         """
         Creates a batch prediction job description for the given deployment.
 
         Args:
             table_name (str): Name of the feature group table to write the results of the batch prediction. Can only be specified if outputLocation and databaseConnectorId are not specified. If tableName is specified, the outputType will be enforced as CSV.
             name (str): Name of the batch prediction job.
-            global_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
-            explanations (bool): If true, SHAP explanations will be provided for each prediction, if supported by the use case.
+            batch_prediction_args (BatchPredictionArgs): Batch Prediction args specific to problem type.
             output_format (str): Format of the batch prediction output (CSV or JSON).
             output_location (str): Location to write the prediction results. Otherwise, results will be stored in Abacus.AI.
             database_connector_id (str): Unique identifier of a Database Connection to write predictions to. Cannot be specified in conjunction with outputLocation.
@@ -358,7 +357,7 @@ class Deployment(AbstractApiClass):
         Returns:
             BatchPrediction: The batch prediction description.
         """
-        return self.client.create_batch_prediction(self.deployment_id, table_name, name, global_prediction_args, explanations, output_format, output_location, database_connector_id, database_output_config, refresh_schedule, csv_input_prefix, csv_prediction_prefix, csv_explanations_prefix, output_includes_metadata, result_input_columns, input_feature_groups)
+        return self.client.create_batch_prediction(self.deployment_id, table_name, name, global_prediction_args, batch_prediction_args, explanations, output_format, output_location, database_connector_id, database_output_config, refresh_schedule, csv_input_prefix, csv_prediction_prefix, csv_explanations_prefix, output_includes_metadata, result_input_columns, input_feature_groups)
 
     def describe_feature_group_row_process_by_key(self, primary_key_value: str):
         """

{abacusai-1.1.5 → abacusai-1.1.6}/abacusai/document_retriever_config.py

@@ -11,18 +11,20 @@ class DocumentRetrieverConfig(AbstractApiClass):
             chunkSize (int): The size of chunks for vector store, i.e., maximum number of words in the chunk.
             chunkOverlapFraction (float): The fraction of overlap between two consecutive chunks.
             textEncoder (str): The text encoder used to encode texts in the vector store.
+            scoreMultiplierColumn (str): The values in this metadata column are used to modify the relevance scores of returned chunks.
     """
 
-    def __init__(self, client, chunkSize=None, chunkOverlapFraction=None, textEncoder=None):
+    def __init__(self, client, chunkSize=None, chunkOverlapFraction=None, textEncoder=None, scoreMultiplierColumn=None):
         super().__init__(client, None)
         self.chunk_size = chunkSize
         self.chunk_overlap_fraction = chunkOverlapFraction
         self.text_encoder = textEncoder
+        self.score_multiplier_column = scoreMultiplierColumn
         self.deprecated_keys = {}
 
     def __repr__(self):
-        repr_dict = {f'chunk_size': repr(self.chunk_size), f'chunk_overlap_fraction': repr(
-            self.chunk_overlap_fraction), f'text_encoder': repr(self.text_encoder)}
+        repr_dict = {f'chunk_size': repr(self.chunk_size), f'chunk_overlap_fraction': repr(self.chunk_overlap_fraction), f'text_encoder': repr(
+            self.text_encoder), f'score_multiplier_column': repr(self.score_multiplier_column)}
         class_name = "DocumentRetrieverConfig"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
         ) if getattr(self, key, None) is not None and key not in self.deprecated_keys])
@@ -35,6 +37,6 @@ class DocumentRetrieverConfig(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'chunk_size': self.chunk_size, 'chunk_overlap_fraction':
-                self.chunk_overlap_fraction, 'text_encoder': self.text_encoder}
+        resp = {'chunk_size': self.chunk_size, 'chunk_overlap_fraction': self.chunk_overlap_fraction,
+                'text_encoder': self.text_encoder, 'score_multiplier_column': self.score_multiplier_column}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}