abacusai 1.3.4__tar.gz → 1.3.6__tar.gz

{abacusai-1.3.4/abacusai.egg-info → abacusai-1.3.6}/PKG-INFO

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

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.3.4"
+__version__ = "1.3.6"

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/__init__.py

@@ -15,3 +15,4 @@ from .monitor_alert import *
 from .project import *
 from .python_functions import *
 from .refresh import *
+from .segments import *

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/ai_agents.py

@@ -100,7 +100,7 @@ class WorkflowNodeInputMapping(ApiClass):
     @classmethod
     def from_dict(cls, mapping: dict):
         if any(field not in mapping for field in ['name', 'variable_type']):
-            raise ValueError('Invalid workflow node input mapping. Must contain keys - name, variable_type')
+            raise ValueError('input_mapping', f'Invalid workflow node input mapping "{mapping}". Must contain keys - name, variable_type')
         return cls(
             name=mapping['name'],
             variable_type=mapping['variable_type'],
@@ -158,7 +158,7 @@ class WorkflowGraphNode(ApiClass):
             self.function_name = function_name
             self.source_code = source_code
         else:
-            raise ValueError('Either function or function_name and source_code must be provided.')
+            raise ValueError('workflow_graph_node', 'Either function or function_name and source_code must be provided.')
 
         self.name = name
         self.input_mappings = input_mappings
@@ -182,7 +182,7 @@ class WorkflowGraphNode(ApiClass):
     @classmethod
     def from_dict(cls, node: dict):
         if any(field not in node for field in ['name', 'function_name', 'source_code', 'input_mappings', 'output_mappings']):
-            raise ValueError('Invalid workflow graph node. Must contain keys - name, function_name, source_code, input_mappings, output_mappings.')
+            raise ValueError('workflow_graph_node', 'Invalid workflow graph node. Must contain keys - name, function_name, source_code, input_mappings, output_mappings.')
         return cls(
             name=node['name'],
             function_name=node['function_name'],
@@ -223,19 +223,23 @@ class WorkflowGraph(ApiClass):
     Args:
         nodes (List[WorkflowGraphNode]): A list of nodes in the workflow graph.
         edges (List[WorkflowGraphEdge]): A list of edges in the workflow graph, where each edge is a tuple of source, target and details.
+        primary_start_node (str): The primary node to start the workflow from.
     """
     nodes: List[WorkflowGraphNode] = dataclasses.field(default_factory=list)
     edges: List[WorkflowGraphEdge] = dataclasses.field(default_factory=list)
+    primary_start_node: str = dataclasses.field(default=None)
 
     def to_dict(self):
         return {
             'nodes': [node.to_dict() for node in self.nodes],
-            'edges': [edge.to_dict() for edge in self.edges]
+            'edges': [edge.to_dict() for edge in self.edges],
+            'primary_start_node': self.primary_start_node
         }
 
     @classmethod
     def from_dict(cls, graph: dict):
         return cls(
             nodes=[WorkflowGraphNode.from_dict(node) for node in graph.get('nodes', [])],
-            edges=[WorkflowGraphEdge.from_dict(edge) for edge in graph.get('edges', [])]
+            edges=[WorkflowGraphEdge.from_dict(edge) for edge in graph.get('edges', [])],
+            primary_start_node=graph.get('primary_start_node', None)
         )

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/batch_prediction.py

@@ -201,6 +201,16 @@ class TrainablePlugAndPlayBatchPredictionArgs(BatchPredictionArgs):
         self.problem_type = enums.ProblemType.CUSTOM_ALGORITHM
 
 
+@dataclasses.dataclass
+class AIAgentBatchPredictionArgs(BatchPredictionArgs):
+    """
+    Batch Prediction Config for the AIAgents problem type
+    """
+
+    def __post_init__(self):
+        self.problem_type = enums.ProblemType.AI_AGENT
+
+
 @dataclasses.dataclass
 class _BatchPredictionArgsFactory(_ApiClassFactory):
     config_abstract_class = BatchPredictionArgs
@@ -215,4 +225,5 @@ class _BatchPredictionArgsFactory(_ApiClassFactory):
         enums.ProblemType.THEME_ANALYSIS: ThemeAnalysisBatchPredictionArgs,
         enums.ProblemType.CHAT_LLM: ChatLLMBatchPredictionArgs,
         enums.ProblemType.CUSTOM_ALGORITHM: TrainablePlugAndPlayBatchPredictionArgs,
+        enums.ProblemType.AI_AGENT: AIAgentBatchPredictionArgs,
     }

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/dataset.py

@@ -42,7 +42,6 @@ class DocumentProcessingConfig(ApiClass):
         remove_header_footer (bool): Whether to remove headers and footers. Defaults to False. This option only takes effect when extract_bounding_boxes is True.
         remove_watermarks (bool): Whether to remove watermarks. By default, it will be decided automatically based on the OCR mode and the document type. This option only takes effect when extract_bounding_boxes is True.
         convert_to_markdown (bool): Whether to convert extracted text to markdown. Defaults to False. This option only takes effect when extract_bounding_boxes is True.
-        return_links (bool): Whether to augment the extracted text with embedded URLs in the PDF. Defaults to False. This option only takes effect when the document is a PDF.
     """
     # NOTE: The defaults should match with clouddb.document_processing_results table defaults
     extract_bounding_boxes: bool = False
@@ -51,7 +50,6 @@ class DocumentProcessingConfig(ApiClass):
     remove_header_footer: bool = False
     remove_watermarks: bool = True
     convert_to_markdown: bool = False
-    return_links: bool = False
 
 
 @dataclasses.dataclass

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/enums.py

@@ -377,6 +377,7 @@ class ApplicationConnectorType(ApiEnum):
     MICROSOFTAUTH = 'MICROSOFTAUTH'
     FRESHSERVICE = 'FRESHSERVICE'
     ZENDESKSUNSHINEMESSAGING = 'ZENDESKSUNSHINEMESSAGING'
+    GOOGLEDRIVEUSER = 'GOOGLEDRIVEUSER'
 
 
 class StreamingConnectorType(ApiEnum):
@@ -438,6 +439,7 @@ class LLMName(ApiEnum):
     CLAUDE_V3_OPUS = 'CLAUDE_V3_OPUS'
     CLAUDE_V3_SONNET = 'CLAUDE_V3_SONNET'
     CLAUDE_V3_HAIKU = 'CLAUDE_V3_HAIKU'
+    CLAUDE_V3_5_SONNET = 'CLAUDE_V3_5_SONNET'
     ABACUS_GIRAFFE = 'ABACUS_GIRAFFE'
     ABACUS_GIRAFFE_LARGE = 'ABACUS_GIRAFFE_LARGE'
     LLAMA2_CHAT = 'LLAMA2_CHAT'
@@ -534,6 +536,7 @@ class OcrMode(ApiEnum):
     COMPREHENSIVE_V2 = 'COMPREHENSIVE_V2'
     COMPREHENSIVE_TABLE_MD = 'COMPREHENSIVE_TABLE_MD'
     TESSERACT_FAST = 'TESSERACT_FAST'
+    LLM = 'LLM'
 
 
 class StdDevThresholdType(ApiEnum):
@@ -586,3 +589,15 @@ class MemorySize(ApiEnum):
             if member.value >= value:
                 return member
         return None
+
+
+class SegmentType(ApiEnum):
+    ATTACHMENTS = 'attachments'
+    AGENT_FLOW_BUTTON = 'agent_flow_button'
+    IMAGE_URL = 'image_url'
+    TEXT = 'text'
+
+
+class CodeSegmentLanguage(ApiEnum):
+    PYTHON = 'python'
+    SQL = 'sql'

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_class/model.py

@@ -65,6 +65,7 @@ class PersonalizationTrainingConfig(TrainingConfig):
         min_item_history (int): Minimum number of interactions an item must have to be included in training.
         query_column (str): Name of column in the interactions table that represents a natural language query, e.g. 'blue t-shirt'.
         item_query_column (str): Name of column in the item catalog that will be matched to the query column in the interactions table.
+        include_item_id_feature (bool): Add Item-Id to the input features of the model. Applicable for Embedding distance and CTR models.
     """
     # top-level params
     objective: enums.PersonalizationObjective = dataclasses.field(default=None)
@@ -112,6 +113,7 @@ class PersonalizationTrainingConfig(TrainingConfig):
     item_query_column: str = dataclasses.field(default=None)
     use_user_id_feature: bool = dataclasses.field(default=None)
     session_dedupe_mins: float = dataclasses.field(default=None)
+    include_item_id_feature: bool = dataclasses.field(default=None)
 
     # outliers
     max_user_history_len_percentile: int = dataclasses.field(default=None)
@@ -171,7 +173,10 @@ class RegressionTrainingConfig(TrainingConfig):
         custom_metrics (List[str]): Registered custom metrics available for selection.
         partial_dependence_analysis (PartialDependenceAnalysis): Specify whether to run partial dependence plots for all features or only some features.
         do_masked_language_model_pretraining (bool): Specify whether to run a masked language model unsupervised pretraining step before supervized training in certain supported algorithms which use BERT-like backbones.
-    """
+        max_tokens_in_sentence (int): Specify the max tokens to be kept in a sentence based on the truncation strategy.
+        truncation_strategy (str): What strategy to use to deal with text rows with more than a given number of tokens (if num of tokens is more than "max_tokens_in_sentence").
+     """
+
     objective: enums.RegressionObjective = dataclasses.field(default=None)
     sort_objective: enums.RegressionObjective = dataclasses.field(default=None)
     tree_hpo_mode: enums.RegressionTreeHPOMode = dataclasses.field(default=None)
@@ -215,6 +220,8 @@ class RegressionTrainingConfig(TrainingConfig):
     pretrained_model_name: str = dataclasses.field(default=None)
     is_multilingual: bool = dataclasses.field(default=None)
     do_masked_language_model_pretraining: bool = dataclasses.field(default=None)
+    max_tokens_in_sentence: int = dataclasses.field(default=None)
+    truncation_strategy: str = dataclasses.field(default=None)
 
     # loss function
     loss_function: enums.RegressionLossFunction = dataclasses.field(default=None)
@@ -299,7 +306,7 @@ class ForecastingTrainingConfig(TrainingConfig):
         use_all_item_totals (bool): Include as input total target across items.
         handle_zeros_as_missing_values (bool): If True, handle zero values in demand as missing data.
         datetime_holiday_calendars (List[HolidayCalendars]): Holiday calendars to augment training with.
-        fill_missing_values (List[dict]): Strategy for filling in missing values.
+        fill_missing_values (List[List[dict]]): Strategy for filling in missing values.
         enable_clustering (bool): Enable clustering in forecasting.
         data_split_feature_group_table_name (str): Specify the table name of the feature group to export training data with the fold column.
         custom_loss_functions (List[str]): Registered custom losses available for selection.
@@ -368,7 +375,7 @@ class ForecastingTrainingConfig(TrainingConfig):
     use_all_item_totals: bool = dataclasses.field(default=None)
     handle_zeros_as_missing_values: bool = dataclasses.field(default=None)
     datetime_holiday_calendars: List[enums.HolidayCalendars] = dataclasses.field(default=None)
-    fill_missing_values: List[dict] = dataclasses.field(default=None)
+    fill_missing_values: List[List[dict]] = dataclasses.field(default=None)
     enable_clustering: bool = dataclasses.field(default=None)
     # Others
     data_split_feature_group_table_name: str = dataclasses.field(default=None)
@@ -656,7 +663,7 @@ class TimeseriesAnomalyTrainingConfig(TrainingConfig):
         type_of_split (TimeseriesAnomalyDataSplitType): Type of data splitting into train/test.
         test_start (str): Limit training data to dates before the given test start.
         test_split (int): Percent of dataset to use for test data. We support using a range between 5 ( i.e. 5% ) to 20 ( i.e. 20% ) of your dataset.
-        fill_missing_values (List[dict]): strategies to fill missing values and missing timestamps
+        fill_missing_values (List[List[dict]]): strategies to fill missing values and missing timestamps
         handle_zeros_as_missing_values (bool): If True, handle zero values in numeric columns as missing data
         timeseries_frequency (str): set this to control frequency of filling missing values
         min_samples_in_normal_region (int): Adjust this to fine-tune the number of anomalies to be identified.
@@ -667,7 +674,7 @@ class TimeseriesAnomalyTrainingConfig(TrainingConfig):
     type_of_split: enums.TimeseriesAnomalyDataSplitType = dataclasses.field(default=None)
     test_start: str = dataclasses.field(default=None)
     test_split: int = dataclasses.field(default=None)
-    fill_missing_values: List[dict] = dataclasses.field(default=None)
+    fill_missing_values: List[List[dict]] = dataclasses.field(default=None)
     handle_zeros_as_missing_values: bool = dataclasses.field(default=None)
     timeseries_frequency: str = dataclasses.field(default=None)
     min_samples_in_normal_region: int = dataclasses.field(default=None)

abacusai-1.3.6/abacusai/api_class/segments.py

@@ -0,0 +1,128 @@
+import dataclasses
+from typing import List
+
+from . import enums
+from .abstract import ApiClass
+
+
+@dataclasses.dataclass
+class Attachment(ApiClass):
+    """
+    An attachment that an agent can return to render attachments.
+
+    Args:
+        filename (str): The name of the file.
+        mime_type (str): The MIME type of the file.
+        attachment_id (str): The ID of the attachment.
+    """
+
+    filename: str
+    mime_type: str
+    attachment_id: str
+
+    def to_dict(self):
+        return {
+            'type': 'attachment',
+            'filename': self.filename,
+            'mime_type': self.mime_type,
+            'attachment_id': self.attachment_id
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(filename=data['filename'], mime_type=data['mime_type'], attachment_id=data['attachment_id'])
+
+
+@dataclasses.dataclass
+class Segment(ApiClass):
+    """
+    A segment that an agent can return to render specific UI elements.
+
+    Args:
+        type (SegmentType): The type of the segment.
+        id (str): The section key of the segment.
+    """
+
+    type: enums.SegmentType
+    id: str
+
+    def to_dict(self):
+        return {k: v.value if isinstance(v, enums.Enum) else v for k, v in dataclasses.asdict(self).items()}
+
+
+@dataclasses.dataclass
+class AttachmentsSegment(Segment):
+    """
+    A segment that an agent can return to render attachments.
+
+    Args:
+        attachments (List[Attachment]): The list of attachments to be displayed.
+    """
+
+    attachments: List[Attachment]
+
+    def __init__(self, attachments: List[Attachment], section_key: str = None):
+        super().__init__(type=enums.SegmentType.ATTACHMENTS, id=section_key)
+        self.attachments = attachments
+
+    def to_dict(self):
+        return {
+            'type': self.type.value,
+            'id': self.id,
+            'attachments': [attachment.to_dict() for attachment in self.attachments]
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(attachments=[Attachment.from_dict(attachment) for attachment in data['attachments']])
+
+
+@dataclasses.dataclass
+class AgentFlowButtonSegment(Segment):
+    """
+    A segment that an AI Agent can return to render a button.
+
+    Args:
+        label (str): The label of the button.
+        agent_workflow_node_name (str): The workflow start node to be executed when the button is clicked.
+    """
+
+    label: str
+    agent_workflow_node_name: str
+
+    def __init__(self, label: str, agent_workflow_node_name: str, section_key: str = None):
+        super().__init__(type=enums.SegmentType.AGENT_FLOW_BUTTON, id=section_key)
+        self.label = label
+        self.agent_workflow_node_name = agent_workflow_node_name
+
+
+@dataclasses.dataclass
+class ImageUrlSegment(Segment):
+    """
+    A segment that an agent can return to render an image.
+
+    Args:
+        segment (str): The url of the image to be displayed.
+    """
+
+    segment: str
+
+    def __init__(self, url: str, section_key: str = None):
+        super().__init__(type=enums.SegmentType.IMAGE_URL, id=section_key)
+        self.segment = url
+
+
+@dataclasses.dataclass
+class TextSegment(Segment):
+    """
+    A segment that an agent can return to render text.
+
+    Args:
+        segment (str): The text to be displayed.
+    """
+
+    segment: str
+
+    def __init__(self, text: str, section_key: str = None):
+        super().__init__(type=enums.SegmentType.TEXT, id=section_key)
+        self.segment = text

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/api_client_utils.py

@@ -57,22 +57,26 @@ class StreamingHandler(str):
 
     @classmethod
     def process_streaming_data(cls, value, context, section_key, data_type):
-        if section_key and hasattr(context, 'streamed_section_response') and isinstance(context.streamed_section_response, list):
-            entry_exists = False
-            for i, item in enumerate(context.streamed_section_response):
-                if item['id'] == section_key:
-                    if ((isinstance(context.streamed_section_response[i]['contents'], str) and isinstance(value, str)) or
-                       (isinstance(context.streamed_section_response[i]['contents'], list) and isinstance(value, list))):
-                        context.streamed_section_response[i]['contents'] += value
-                    else:
-                        context.streamed_section_response[i]['contents'] = value
-                    entry_exists = True
-                    break
-            if not entry_exists:
-                context.streamed_section_response.append(
-                    {'id': section_key, 'type': data_type, 'mime_type': 'text/plain', 'contents': value})
-        elif hasattr(context, 'streamed_response') and isinstance(context.streamed_response, list) and isinstance(value, str):
-            context.streamed_response.append(value)
+        if hasattr(context, 'streamed_section_response') and hasattr(context, 'streamed_response'):
+            if data_type == 'text':
+                if section_key:
+                    entry_exists = False
+                    for i, item in enumerate(context.streamed_section_response):
+                        if item['id'] == section_key:
+                            if ((isinstance(context.streamed_section_response[i]['contents'], str) and isinstance(value, str)) or
+                               (isinstance(context.streamed_section_response[i]['contents'], list) and isinstance(value, list))):
+                                context.streamed_section_response[i]['contents'] += value
+                            else:
+                                context.streamed_section_response[i]['contents'] = value
+                            entry_exists = True
+                            break
+                    if not entry_exists:
+                        context.streamed_section_response.append(
+                            {'id': section_key, 'type': data_type, 'mime_type': 'text/plain', 'contents': value})
+                else:
+                    context.streamed_response.append(value)
+            elif data_type == 'segment':
+                context.streamed_section_response.append(value)
 
 
 def get_object_from_context(client, context, variable_name, return_type):
@@ -178,6 +182,7 @@ def load_as_pandas_from_avro_files(files: List[str], download_method: Callable,
 class StreamType(Enum):
     MESSAGE = 'message'
     SECTION_OUTPUT = 'section_output'
+    SEGMENT = 'segment'
 
 
 class DocstoreUtils:

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/batch_prediction.py

@@ -84,7 +84,7 @@ class BatchPrediction(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(getattr(
             api_class, batchPredictionArgsType, BatchPredictionArgs) if batchPredictionArgsType else BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'explanations', 'global_prediction_args'}
+        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'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(

{abacusai-1.3.4 → abacusai-1.3.6}/abacusai/batch_prediction_version.py

@@ -100,7 +100,7 @@ class BatchPredictionVersion(AbstractApiClass):
             BatchPredictionArgs, globalPredictionArgs)
         self.batch_prediction_args = client._build_class(getattr(
             api_class, batchPredictionArgsType, BatchPredictionArgs) if batchPredictionArgsType else BatchPredictionArgs, batchPredictionArgs)
-        self.deprecated_keys = {'explanations', 'global_prediction_args'}
+        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'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(