Version 1.4.31

Author: arvindsun

abacusai/__init__.py

@@ -67,7 +67,6 @@
 from .deployment_statistics import DeploymentStatistics
 from .document_data import DocumentData
 from .document_retriever import DocumentRetriever
-from .document_retriever_config import DocumentRetrieverConfig
 from .document_retriever_lookup_result import DocumentRetrieverLookupResult
 from .document_retriever_version import DocumentRetrieverVersion
 from .drift_distribution import DriftDistribution
@@ -209,6 +208,7 @@
 from .resolved_feature_group_template import ResolvedFeatureGroupTemplate
 from .routing_action import RoutingAction
 from .schema import Schema
+from .sftp_key import SftpKey
 from .streaming_auth_token import StreamingAuthToken
 from .streaming_client import StreamingClient
 from .streaming_connector import StreamingConnector
@@ -236,4 +236,4 @@
 from .workflow_node_template import WorkflowNodeTemplate
 
 
-__version__ = "1.4.30"
+__version__ = "1.4.31"

abacusai/api_class/ai_agents.py

@@ -427,6 +427,9 @@ def _raw_init(cls, name: str, input_mappings: List[WorkflowNodeInputMapping] = N
         elif function_name and (source_code or template_metadata):
             workflow_node.function_name = function_name
             workflow_node.source_code = source_code
+        elif template_metadata and not template_metadata.get('initialized'):
+            workflow_node.function_name = function_name
+            workflow_node.source_code = source_code
         else:
             raise ValueError('workflow_graph_node', 'Either function or function_name and source_code must be provided.')
         workflow_node.input_mappings = input_mappings
@@ -504,6 +507,10 @@ def is_trigger_node(self):
 
     @classmethod
     def from_dict(cls, node: dict):
+        if node.get('template_metadata'):
+            node['function_name'] = node.get('function_name')
+            node['source_code'] = node.get('source_code')
+
         validate_input_dict_param(node, friendly_class_name='workflow_graph_node', must_contain=['name', 'function_name', 'source_code'])
         _cls = cls._raw_init if node.get('__return_filter') else cls
         if node.get('template_metadata') and node.get('template_metadata').get('template_type') == 'trigger':

abacusai/api_class/dataset.py

@@ -45,6 +45,7 @@ class DocumentProcessingConfig(ApiClass):
         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.
         mask_pii (bool): Whether to mask personally identifiable information (PII) in the document text/tokens. Defaults to False.
+        extract_images (bool): Whether to extract images from the document e.g. diagrams in a PDF page. Defaults to False.
     """
     # NOTE: The defaults should match with clouddb.document_processing_results table defaults
     document_type: DocumentType = None
@@ -56,6 +57,7 @@ class DocumentProcessingConfig(ApiClass):
     remove_watermarks: bool = True
     convert_to_markdown: bool = False
     mask_pii: bool = False
+    extract_images: bool = False
 
     def __post_init__(self):
         self.ocr_mode = self._detect_ocr_mode()

abacusai/api_class/dataset_application_connector.py

@@ -50,8 +50,12 @@ class BoxDatasetConfig(ApplicationConnectorDatasetConfig):
     Dataset config for Box Application Connector
     Args:
         location (str): The regex location of the files to fetch
+        csv_delimiter (str): If the file format is CSV, use a specific csv delimiter
+        merge_file_schemas (bool): Signifies if the merge file schema policy is enabled. Not applicable if is_documentset is True
     """
     location: str = dataclasses.field(default=None)
+    csv_delimiter: str = dataclasses.field(default=None)
+    merge_file_schemas: bool = dataclasses.field(default=False)
 
     def __post_init__(self):
         self.application_connector_type = enums.ApplicationConnectorType.BOX

abacusai/api_class/document_retriever.py

@@ -19,6 +19,7 @@ class VectorStoreConfig(ApiClass):
         index_metadata_columns (bool): If True, metadata columns of the FG will also be used for indexing and querying.
         use_document_summary (bool): If True, uses the summary of the document in addition to chunks of the document for indexing and querying.
         summary_instructions (str): Instructions for the LLM to generate the document summary.
+        standalone_deployment (bool): If True, the document retriever will be deployed as a standalone deployment.
     """
     chunk_size: int = dataclasses.field(default=None)
     chunk_overlap_fraction: float = dataclasses.field(default=None)
@@ -29,6 +30,7 @@ class VectorStoreConfig(ApiClass):
     index_metadata_columns: bool = dataclasses.field(default=None)
     use_document_summary: bool = dataclasses.field(default=None)
     summary_instructions: str = dataclasses.field(default=None)
+    standalone_deployment: bool = dataclasses.field(default=False)
 
 
 DocumentRetrieverConfig = VectorStoreConfig

abacusai/api_class/enums.py

@@ -541,6 +541,7 @@ class PythonFunctionType(ApiEnum):
     FEATURE_GROUP = 'FEATURE_GROUP'
     PLOTLY_FIG = 'PLOTLY_FIG'
     STEP_FUNCTION = 'STEP_FUNCTION'
+    USERCODE_TOOL = 'USERCODE_TOOL'
 
 
 class EvalArtifactType(ApiEnum):

abacusai/api_class/python_functions.py

@@ -16,12 +16,16 @@ class PythonFunctionArgument(ApiClass):
         is_required (bool): Whether the argument is required
         value (Any): The value of the argument
         pipeline_variable (str): The name of the pipeline variable to use as the value
+        description (str): The description of the argument
+        item_type (str): Type of items when variable_type is LIST
     """
     variable_type: enums.PythonFunctionArgumentType = dataclasses.field(default=None)
     name: str = dataclasses.field(default=None)
     is_required: bool = dataclasses.field(default=True)
     value: Any = dataclasses.field(default=None)
     pipeline_variable: str = dataclasses.field(default=None)
+    description: str = dataclasses.field(default=None)
+    item_type: str = dataclasses.field(default=None)
 
 
 @dataclasses.dataclass

abacusai/batch_prediction.py

@@ -84,7 +84,7 @@ def __init__(self, client, batchPredictionId=None, createdAt=None, name=None, de
             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/batch_prediction_version.py

@@ -100,7 +100,7 @@ def __init__(self, client, batchPredictionVersion=None, batchPredictionId=None,
             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(

abacusai/client.py

@@ -652,7 +652,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.4.30'
+    client_version = '1.4.31'
 
     def __init__(self, api_key: str = None, server: str = None, client_options: ClientOptions = None, skip_version_check: bool = False, include_tb: bool = False):
         self.api_key = api_key
@@ -3935,13 +3935,102 @@ def get_agent_context_user_info(self):
             raise ValueError(
                 'User information not available. Please use UI interface for this agent to work.')
 
+    def get_agent_runtime_config(self, key: str):
+        """
+        Gets the deployment level runtime config for the agent
+
+        Args:
+            key(str): Key for which the config value is to be fetched
+
+        Returns:
+            str: Config value for the input key
+        """
+        runtime_config = get_object_from_context(
+            self, _request_context, 'deployment_runtime_config', dict) or {}
+        return runtime_config.get(key, None)
+
+    def get_request_user_info(self):
+        """
+        Gets the user information for the current request context.
+
+        Returns:
+            dict: Containing email and name of the end user.
+        """
+        user_info = get_object_from_context(
+            self, _request_context, 'user_info', dict)
+        if user_info:
+            return user_info
+        else:
+            raise ValueError('User information not available')
+
     def clear_agent_context(self):
         """
         Clears the current request context.
         """
         if hasattr(_request_context):
             _request_context.clear()
 
+    def execute_chatllm_computer_streaming(self, computer_id: str, prompt: str, is_transient: bool = False):
+        """
+        Executes a prompt on a remote computer and streams computer responses to the external chat UI in real-time. Must be called from agent execution context only.
+
+        Args:
+            computer_id (str): The ID of the computer to use for the agent.
+            prompt (str): The prompt to do tasks on the computer.
+            is_transient (bool): If True, the message will be marked as transient and will not be persisted on reload in external chatllm UI. Transient messages are useful for streaming interim updates or results.
+
+        Returns:
+            text (str): The text responses from the computer.
+        """
+        request_id = self._get_agent_app_request_id()
+        caller = self._get_agent_async_app_caller()
+        proxy_caller = self._is_proxy_app_caller()
+
+        if not request_id or not caller:
+            raise Exception(
+                'This function can only be called from within an agent execution context')
+
+        if not caller.endswith('/'):
+            caller = caller + '/'
+
+        if proxy_caller:
+            api_endpoint = f'{caller}_executeChatLLMComputerStreaming'
+        else:
+            raise Exception(
+                'This function can only be called from within an agent execution context')
+
+        extra_args = {'stream_type': StreamType.MESSAGE.value,
+                      'response_version': '1.0', 'is_transient': is_transient}
+        if hasattr(_request_context, 'agent_workflow_node_id'):
+            extra_args.update(
+                {'agent_workflow_node_id': _request_context.agent_workflow_node_id})
+
+        computer_use_args = {
+            'computerId': computer_id,
+            'prompt': prompt
+        }
+
+        body = {
+            'requestId': request_id,
+            'computerUseArgs': computer_use_args,
+            'extraArgs': extra_args,
+        }
+        body['connectionId'] = uuid4().hex
+
+        headers = {'APIKEY': self.api_key}
+        self._clean_api_objects(body)
+        for _ in range(3):
+            response = self._request(
+                api_endpoint, method='POST', body=body, headers=headers)
+            if response.status_code == 200:
+                return StreamingHandler(response.json(), _request_context, is_transient=is_transient)
+            elif response.status_code in (502, 503, 504):
+                continue
+            else:
+                break
+        raise Exception(
+            f'Error calling ChatLLM computer streaming endpoint. Status code: {response.status_code}. Response: {response.text}')
+
     def streaming_evaluate_prompt(self, prompt: str = None, system_message: str = None, llm_name: Union[LLMName, str] = None, max_tokens: int = None, temperature: float = 0.0, messages: list = None, response_type: str = None, json_response_schema: dict = None, section_key: str = None):
         """
         Generate response to the prompt using the specified model. This works the same as `evaluate_prompt` but would stream the text to the UI section while generating and returns the streamed text as an object of a `str` subclass.
@@ -7921,7 +8010,7 @@ def get_feature_group_row_process_logs_by_key(self, deployment_id: str, primary_
             FeatureGroupRowProcessLogs: An object representing the logs for the feature group row process"""
         return self._call_api('getFeatureGroupRowProcessLogsByKey', 'POST', query_params={'deploymentId': deployment_id}, body={'primaryKeyValue': primary_key_value}, parse_type=FeatureGroupRowProcessLogs)
 
-    def create_python_function(self, name: str, source_code: str = None, function_name: str = None, function_variable_mappings: List = None, package_requirements: list = None, function_type: str = 'FEATURE_GROUP') -> PythonFunction:
+    def create_python_function(self, name: str, source_code: str = None, function_name: str = None, function_variable_mappings: List = None, package_requirements: list = None, function_type: str = 'FEATURE_GROUP', description: str = None, examples: dict = None) -> PythonFunction:
         """Creates a custom Python function that is reusable.
 
         Args:
@@ -7931,12 +8020,14 @@ def create_python_function(self, name: str, source_code: str = None, function_na
             function_variable_mappings (List): List of Python function arguments.
             package_requirements (list): List of package requirement strings. For example: ['numpy==1.2.3', 'pandas>=1.4.0'].
             function_type (str): Type of Python function to create. Default is FEATURE_GROUP, but can also be PLOTLY_FIG.
+            description (str): Description of the Python function. This should include details about the function's purpose, expected inputs and outputs, and any important usage considerations or limitations.
+            examples (dict): Dictionary containing example use cases and anti-patterns. Should include 'positive_examples' showing recommended usage and 'negative_examples' showing cases to avoid.
 
         Returns:
             PythonFunction: The Python function that can be used (e.g. for feature group transform)."""
-        return self._call_api('createPythonFunction', 'POST', query_params={}, body={'name': name, 'sourceCode': source_code, 'functionName': function_name, 'functionVariableMappings': function_variable_mappings, 'packageRequirements': package_requirements, 'functionType': function_type}, parse_type=PythonFunction)
+        return self._call_api('createPythonFunction', 'POST', query_params={}, body={'name': name, 'sourceCode': source_code, 'functionName': function_name, 'functionVariableMappings': function_variable_mappings, 'packageRequirements': package_requirements, 'functionType': function_type, 'description': description, 'examples': examples}, parse_type=PythonFunction)
 
-    def update_python_function(self, name: str, source_code: str = None, function_name: str = None, function_variable_mappings: List = None, package_requirements: list = None) -> PythonFunction:
+    def update_python_function(self, name: str, source_code: str = None, function_name: str = None, function_variable_mappings: List = None, package_requirements: list = None, description: str = None, examples: dict = None) -> PythonFunction:
         """Update custom python function with user inputs for the given python function.
 
         Args:
@@ -7945,10 +8036,12 @@ def update_python_function(self, name: str, source_code: str = None, function_na
             function_name (str): The name of the Python function within `source_code`.
             function_variable_mappings (List): List of arguments required by `function_name`.
             package_requirements (list): List of package requirement strings. For example: ['numpy==1.2.3', 'pandas>=1.4.0'].
+            description (str): Description of the Python function. This should include details about the function's purpose, expected inputs and outputs, and any important usage considerations or limitations.
+            examples (dict): Dictionary containing example use cases and anti-patterns. Should include 'positive_examples' showing recommended usage and 'negative_examples' showing cases to avoid.
 
         Returns:
             PythonFunction: The Python function object."""
-        return self._call_api('updatePythonFunction', 'PATCH', query_params={}, body={'name': name, 'sourceCode': source_code, 'functionName': function_name, 'functionVariableMappings': function_variable_mappings, 'packageRequirements': package_requirements}, parse_type=PythonFunction)
+        return self._call_api('updatePythonFunction', 'PATCH', query_params={}, body={'name': name, 'sourceCode': source_code, 'functionName': function_name, 'functionVariableMappings': function_variable_mappings, 'packageRequirements': package_requirements, 'description': description, 'examples': examples}, parse_type=PythonFunction)
 
     def delete_python_function(self, name: str):
         """Removes an existing Python function.