abacusai 1.4.2__tar.gz → 1.4.4__tar.gz

{abacusai-1.4.2/abacusai.egg-info → abacusai-1.4.4}/PKG-INFO

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

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/__init__.py

@@ -4,4 +4,4 @@ from .prediction_client import PredictionClient
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.4.1"
+__version__ = "1.4.4"

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/api_class/ai_agents.py

@@ -126,18 +126,20 @@ class WorkflowNodeInputMapping(ApiClass):
     def to_dict(self):
         return {
             'name': self.name,
-            'variable_type': self.variable_type,
+            'variable_type': self.variable_type.value,
             'variable_source': self.variable_source,
             'is_required': self.is_required
         }
 
     @classmethod
     def from_dict(cls, mapping: dict):
+        if not isinstance(mapping, dict):
+            raise ValueError('input_mapping', 'Invalid workflow node input mapping. Provided mapping should be a dictionary.')
         if any(field not in mapping for field in ['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'],
+            variable_type=enums.WorkflowNodeInputType(mapping['variable_type']),
             variable_source=mapping.get('variable_source'),
             is_required=mapping.get('is_required', True)
         )
@@ -158,16 +160,18 @@ class WorkflowNodeOutputMapping(ApiClass):
     def to_dict(self):
         return {
             'name': self.name,
-            'variable_type': self.variable_type
+            'variable_type': self.variable_type.value
         }
 
     @classmethod
     def from_dict(cls, mapping: dict):
+        if not isinstance(mapping, dict):
+            raise ValueError('output_mapping', 'Invalid workflow node output mapping. Provided mapping should be a dictionary.')
         if 'name' not in mapping:
             raise ValueError('output_mapping', 'Invalid workflow node output mapping. Must contain keys - name.')
         return cls(
             name=mapping['name'],
-            variable_type=mapping.get('variable_type', enums.WorkflowNodeOutputType.STRING)
+            variable_type=enums.WorkflowNodeOutputType(mapping.get('variable_type', 'STRING'))
         )
 
 

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/api_class/dataset.py

@@ -42,6 +42,7 @@ 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.
+        mask_pii (bool): Whether to mask personally identifiable information (PII) in the document text/tokens. Defaults to False.
     """
     # NOTE: The defaults should match with clouddb.document_processing_results table defaults
     extract_bounding_boxes: bool = False
@@ -50,6 +51,7 @@ class DocumentProcessingConfig(ApiClass):
     remove_header_footer: bool = False
     remove_watermarks: bool = True
     convert_to_markdown: bool = False
+    mask_pii: bool = False
 
 
 @dataclasses.dataclass

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/api_class/enums.py

@@ -422,6 +422,7 @@ class PythonFunctionOutputArgumentType(ApiEnum):
 class VectorStoreTextEncoder(ApiEnum):
     E5 = 'E5'
     OPENAI = 'OPENAI'
+    OPENAI_LARGE = 'OPENAI_LARGE'
     SENTENCE_BERT = 'SENTENCE_BERT'
     E5_SMALL = 'E5_SMALL'
     CODE_BERT = 'CODE_BERT'
@@ -432,27 +433,19 @@ class LLMName(ApiEnum):
     OPENAI_GPT4_32K = 'OPENAI_GPT4_32K'
     OPENAI_GPT4_128K = 'OPENAI_GPT4_128K'
     OPENAI_GPT4_128K_LATEST = 'OPENAI_GPT4_128K_LATEST'
-    OPENAI_GPT4_VISION = 'OPENAI_GPT4_VISION'
     OPENAI_GPT4O = 'OPENAI_GPT4O'
     OPENAI_GPT4O_MINI = 'OPENAI_GPT4O_MINI'
     OPENAI_GPT3_5 = 'OPENAI_GPT3_5'
     OPENAI_GPT3_5_TEXT = 'OPENAI_GPT3_5_TEXT'
+    LLAMA3_1_405B = 'LLAMA3_1_405B'
+    LLAMA3_1_70B = 'LLAMA3_1_70B'
+    LLAMA3_1_8B = 'LLAMA3_1_8B'
     LLAMA3_LARGE_CHAT = 'LLAMA3_LARGE_CHAT'
-    GROQ_LLAMA3_LARGE_CHAT = 'GROQ_LLAMA3_LARGE_CHAT'
-    CLAUDE_V2_1 = 'CLAUDE_V2_1'
     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'
-    PALM = 'PALM'
-    PALM_TEXT = 'PALM_TEXT'
-    GEMINI_PRO = 'GEMINI_PRO'
     GEMINI_1_5_PRO = 'GEMINI_1_5_PRO'
-    MIXTRAL_CHAT = 'MIXTRAL_CHAT'
-    MISTRAL_MEDIUM = 'MISTRAL_MEDIUM'
     ABACUS_SMAUG3 = 'ABACUS_SMAUG3'
     GEMINI_1_5_FLASH = 'GEMINI_1_5_FLASH'
 
@@ -542,6 +535,7 @@ class OcrMode(ApiEnum):
     COMPREHENSIVE_TABLE_MD = 'COMPREHENSIVE_TABLE_MD'
     TESSERACT_FAST = 'TESSERACT_FAST'
     LLM = 'LLM'
+    AUGMENTED_LLM = 'AUGMENTED_LLM'
 
 
 class StdDevThresholdType(ApiEnum):

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/api_class/model.py

@@ -460,8 +460,8 @@ class ChatLLMTrainingConfig(TrainingConfig):
         filter_columns (list): Allow users to filter the document retrievers on these metadata columns.
         include_general_knowledge (bool): Allow the LLM to rely not just on RAG search results, but to fall back on general knowledge. Disabled by default.
         enable_web_search (bool) : Allow the LLM to use Web Search Engines to retrieve information for better results.
-        behavior_instructions (str): Customize the overall role instructions for the LLM.
-        response_instructions (str): Customized instructions for how the LLM should respond.
+        behavior_instructions (str): Customize the overall behaviour of the model. This controls things like - when to execute code (if enabled), write sql query, search web (if enabled), etc.
+        response_instructions (str): Customized instructions for how the model should respond.
         enable_llm_rewrite (bool): If enabled, an LLM will rewrite the RAG queries sent to document retriever. Disabled by default.
         column_filtering_instructions (str): Instructions for a LLM call to automatically generate filter expressions on document metadata to retrieve relevant documents for the conversation.
         keyword_requirement_instructions (str): Instructions for a LLM call to automatically generate keyword requirements to retrieve relevant documents for the conversation.
@@ -469,14 +469,17 @@ class ChatLLMTrainingConfig(TrainingConfig):
         max_search_results (int): Maximum number of search results in the retrieval augmentation step. If we know that the questions are likely to have snippets which are easily matched in the documents, then a lower number will help with accuracy.
         data_feature_group_ids: (List[str]): List of feature group IDs to use to possibly query for the ChatLLM. The created ChatLLM is commonly referred to as DataLLM.
         data_prompt_context (str): Prompt context for the data feature group IDs.
+        data_prompt_table_context (Dict[str, str]): Dict of table name and table context pairs to provide table wise context for each structured data table.
         hide_sql_and_code (bool): When running data queries, this will hide the generated SQL and Code in the response.
         disable_data_summarization (bool): After executing a query summarize the reponse and reply back with only the table and query run.
+        data_columns_to_ignore (List[str]): Columns to ignore while encoding information about structured data tables in context for the LLM. A list of strings of format "<table_name>.<column_name>"
         search_score_cutoff (float): Minimum search score to consider a document as a valid search result.
         include_bm25_retrieval (bool): Combine BM25 search score with vector search using reciprocal rank fusion.
-        database_connector_id (str): Database connector ID to use for the ChatLLM.
+        database_connector_id (str): Database connector ID to use for connecting external database that gives access to structured data to the LLM.
         database_connector_tables (List[str]): List of tables to use from the database connector for the ChatLLM.
         enable_code_execution (bool): Enable python code execution in the ChatLLM. This equips the LLM with a python kernel in which all its code is executed.
         enable_response_caching (bool): Enable caching of LLM responses to speed up response times and improve reproducibility.
+        unknown_answer_phrase (str): Fallback response when the LLM can't find an answer.
     """
     document_retrievers: List[str] = dataclasses.field(default=None)
     num_completion_tokens: int = dataclasses.field(default=None)
@@ -494,8 +497,10 @@ class ChatLLMTrainingConfig(TrainingConfig):
     max_search_results: int = dataclasses.field(default=None)
     data_feature_group_ids: List[str] = dataclasses.field(default=None)
     data_prompt_context: str = dataclasses.field(default=None)
+    data_prompt_table_context: Dict[str, str] = dataclasses.field(default=None)
     hide_sql_and_code: bool = dataclasses.field(default=None)
     disable_data_summarization: bool = dataclasses.field(default=None)
+    data_columns_to_ignore: List[str] = dataclasses.field(default=None)
     search_score_cutoff: float = dataclasses.field(default=None)
     include_bm25_retrieval: bool = dataclasses.field(default=None)
     database_connector_id: str = dataclasses.field(default=None)
@@ -504,6 +509,7 @@ class ChatLLMTrainingConfig(TrainingConfig):
     metadata_columns: list = dataclasses.field(default=None, metadata={'deprecated': True})
     lookup_rewrite_instructions: str = dataclasses.field(default=None, metadata={'deprecated': True})
     enable_response_caching: bool = dataclasses.field(default=None)
+    unknown_answer_phrase: str = dataclasses.field(default=None)
 
     def __post_init__(self):
         self.problem_type = enums.ProblemType.CHAT_LLM

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/api_client_utils.py

@@ -180,6 +180,194 @@ def load_as_pandas_from_avro_files(files: List[str], download_method: Callable,
     return data_df
 
 
+def validate_workflow_node_inputs(nodes_info, agent_workflow_node_id, keyword_arguments: dict, sample_user_inputs: dict, filtered_workflow_vars: dict):
+    from .api_class import WorkflowNodeInputType
+    input_mappings = nodes_info[agent_workflow_node_id].get(
+        'input_mappings', {})
+    input_schema = nodes_info[agent_workflow_node_id].get('input_schema', {})
+    if input_schema.get('runtime_schema', False):
+        keyword_arguments = {input_schema.get(
+            'schema_prop'): keyword_arguments}
+    for input_mapping in input_mappings:
+        input_name = input_mapping['name']
+        variable_type = input_mapping['variable_type']
+        is_required = input_mapping.get('is_required', True)
+        variable_source = input_mapping['variable_source']
+        if variable_type == 'WORKFLOW_VARIABLE':
+            if variable_source not in filtered_workflow_vars:
+                raise ValueError(
+                    f'The stage corresponding to "{agent_workflow_node_id}" requires variables from {variable_source} stage which are not there.')
+            if input_name not in filtered_workflow_vars[variable_source] and is_required:
+                raise ValueError(
+                    f'Missing required input "{input_name}" in workflow vars for workflow node "{agent_workflow_node_id}".')
+            else:
+                keyword_arguments[input_name] = filtered_workflow_vars[variable_source][input_name]
+        elif variable_type == WorkflowNodeInputType.USER_INPUT:
+            if sample_user_inputs and input_name in sample_user_inputs:
+                keyword_arguments[input_name] = sample_user_inputs[input_name]
+            elif variable_source in filtered_workflow_vars and input_name in filtered_workflow_vars[variable_source]:
+                keyword_arguments[input_name] = filtered_workflow_vars[variable_source][input_name]
+            else:
+                if is_required:
+                    raise ValueError(
+                        f'User input for "{input_name}" is required for the "{agent_workflow_node_id}" node.')
+                else:
+                    keyword_arguments[input_name] = None
+
+
+def run(nodes: List[dict], primary_start_node: str, graph_info: dict, sample_user_inputs: dict = None, agent_workflow_node_id: str = None, workflow_vars: dict = {}, topological_dfs_stack: List = []):
+    from .api_class import WorkflowNodeInputType
+    source_code = graph_info['source_code']
+    exec(source_code, globals())
+
+    nodes_info: dict = {node['name']: node for node in nodes}
+    traversal_orders = graph_info['traversal_orders']
+    nodes_ancestors = graph_info['nodes_ancestors']
+    nodes_inedges = graph_info['nodes_inedges']
+    primary_start_node = primary_start_node or graph_info['default_root_node']
+
+    primary_traversal_order = traversal_orders[primary_start_node]
+    run_info = {}
+    workflow_vars = workflow_vars.copy()
+    next_agent_workflow_node_id = None
+
+    if agent_workflow_node_id:
+        next_agent_workflow_node_id = agent_workflow_node_id
+        if next_agent_workflow_node_id not in traversal_orders.keys():
+            if next_agent_workflow_node_id not in nodes_info:
+                raise ValueError(
+                    f'The provided workflow node id "{next_agent_workflow_node_id}" is not part of the workflow. Please provide a valid node id.')
+            else:
+                topological_dfs_stack.append(next_agent_workflow_node_id)
+        else:
+            topological_dfs_stack = [next_agent_workflow_node_id]
+    else:
+        next_agent_workflow_node_id = primary_start_node
+        topological_dfs_stack = [primary_start_node]
+
+    flow_traversal_order = primary_traversal_order
+    for root, traversal_order in traversal_orders.items():
+        if next_agent_workflow_node_id in traversal_order:
+            flow_traversal_order = traversal_order
+            break
+
+    run_history = []
+    workflow_node_outputs = {}
+    while (True):
+        agent_workflow_node_id = next_agent_workflow_node_id
+        node_ancestors = nodes_ancestors[agent_workflow_node_id]
+
+        # To ensure the node takes inputs only from it's ancestors.
+        # workflow_vars must always contain an entry for ancestor, the error is somewhere else if this ever errors out.
+        filtered_workflow_vars = {}
+        for ancestor in node_ancestors:
+            if ancestor not in workflow_vars:
+                raise ValueError(
+                    f'Ancestor "{ancestor}" of node "{agent_workflow_node_id}" is not executed yet. Please make sure the ancestor nodes are executed before the current node.')
+            else:
+                filtered_workflow_vars[ancestor] = workflow_vars[ancestor]
+
+        arguments = []
+        keyword_arguments = {}
+        validate_workflow_node_inputs(nodes_info, agent_workflow_node_id,
+                                      keyword_arguments, sample_user_inputs, filtered_workflow_vars)
+
+        try:
+            func = eval(nodes_info[agent_workflow_node_id]['function_name'])
+            node_response = func(*arguments, **keyword_arguments)
+            workflow_node_outputs[agent_workflow_node_id] = node_response.to_dict(
+            )
+            node_workflow_vars = process_node_response(node_response)
+        except Exception as error:
+            raise ValueError(
+                f'Error in running workflow node {agent_workflow_node_id}: {error}')
+
+        workflow_vars[agent_workflow_node_id] = node_workflow_vars
+        next_agent_workflow_node_id = None
+        needs_user_input = False
+
+        potential_next_index = flow_traversal_order.index(
+            topological_dfs_stack[-1]) + 1
+        potential_next_agent_workflow_node_id = None
+        while (potential_next_index < len(flow_traversal_order)):
+            potential_next_agent_workflow_node_id = flow_traversal_order[potential_next_index]
+            incoming_edges = nodes_inedges[potential_next_agent_workflow_node_id]
+            valid_next_node = True
+            for source, _, details in incoming_edges:
+                if source not in topological_dfs_stack:
+                    valid_next_node = False
+                    potential_next_index += 1
+                    break
+                else:
+                    edge_evaluate_result = evaluate_edge_condition(
+                        source, potential_next_agent_workflow_node_id, details, workflow_vars)
+                    if not edge_evaluate_result:
+                        valid_next_node = False
+                        potential_next_index += 1
+                        break
+            if valid_next_node:
+                next_agent_workflow_node_id = potential_next_agent_workflow_node_id
+                break
+
+        if next_agent_workflow_node_id:
+            next_node_input_mappings = nodes_info[next_agent_workflow_node_id].get(
+                'input_mappings', [])
+            needs_user_input = any([input_mapping['variable_type'] ==
+                                   WorkflowNodeInputType.USER_INPUT for input_mapping in next_node_input_mappings])
+
+        if needs_user_input:
+            run_history.append(
+                f'Workflow node {agent_workflow_node_id} completed with next node {next_agent_workflow_node_id} and needs user_inputs')
+        else:
+            run_history.append(
+                f'Workflow node {agent_workflow_node_id} completed with next node {next_agent_workflow_node_id}')
+            topological_dfs_stack.append(next_agent_workflow_node_id)
+
+        if next_agent_workflow_node_id is None or needs_user_input:
+            break
+
+    run_info['workflow_node_outputs'] = workflow_node_outputs
+    run_info['run_history'] = run_history
+
+    workflow_info = {}
+    workflow_info['workflow_vars'] = workflow_vars
+    workflow_info['topological_dfs_stack'] = topological_dfs_stack
+    workflow_info['run_info'] = run_info
+
+    return workflow_info
+
+
+def evaluate_edge_condition(source, target, details, workflow_vars):
+    try:
+        condition = details.get('EXECUTION_CONDITION')
+        if condition:
+            result = execute_python_source(
+                condition, workflow_vars.get(source, {}))
+            return result
+        return True
+    except Exception as e:
+        raise ValueError(
+            f"Error evaluating edge '{source}'-->'{target}': {str(e)}")
+
+
+def execute_python_source(python_expression, variables):
+    try:
+        # Evaluate the expression using the variables dictionary
+        result = eval(python_expression, {}, variables)
+        return result
+    except Exception as e:
+        # Handle any exceptions that may occur during evaluation
+        raise ValueError(f'Error evaluating expression: {e}')
+
+
+def process_node_response(node_response):
+    output_vars = {}
+    for variable in node_response.section_data_list:
+        for key, value in variable.items():
+            output_vars[key] = value
+    return output_vars
+
+
 class StreamType(Enum):
     MESSAGE = 'message'
     SECTION_OUTPUT = 'section_output'

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/client.py

@@ -47,10 +47,11 @@ from .api_class import (
 from .api_class.abstract import get_clean_function_source_code, get_clean_function_source_code_for_agent, snake_case
 from .api_class.ai_agents import WorkflowGraph
 from .api_class.blob_input import Blob, BlobInput
+from .api_class.enums import AgentInterface
 from .api_class.segments import ResponseSection, Segment
 from .api_client_utils import (
     INVALID_PANDAS_COLUMN_NAME_CHARACTERS, StreamingHandler, StreamType,
-    clean_column_name, get_object_from_context
+    clean_column_name, get_object_from_context, run
 )
 from .api_endpoint import ApiEndpoint
 from .api_key import ApiKey
@@ -272,14 +273,22 @@ class AgentResponse:
                 self.data_list.append(arg)
             else:
                 raise Exception(
-                    'AgentResponse can only contain Blob, ResponseSection or json serializable objects')
+                    'AgentResponse can only contain Blob, ResponseSection or json serializable objects if key is not provided')
 
         for key, value in kwargs.items():
-            if isinstance(value, Blob) or _is_json_serializable(value) or isinstance(value, (ResponseSection, Segment)):
-                self.section_data_list.append({key: value})
-            else:
-                raise Exception(
-                    'AgentResponse can only contain Blob, ResponseSection or json serializable objects')
+            self.section_data_list.append({key: value})
+
+    def to_dict(self):
+        result = {}
+        if self.data_list:
+            result['data_list'] = self.data_list
+        for section_data in self.section_data_list:
+            for k, v in section_data.items():
+                if isinstance(v, ResponseSection):
+                    result[k] = v.to_dict()
+                else:
+                    result[k] = v
+        return result
 
 
 class ClientOptions:
@@ -595,7 +604,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.1'
+    client_version = '1.4.4'
 
     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
@@ -3520,6 +3529,32 @@ class ApiClient(ReadOnlyClient):
             setattr(the_main, name, getattr(module, name))
         return module
 
+    def run_workflow_graph(self, workflow_graph: WorkflowGraph, sample_user_inputs: dict = {}, agent_workflow_node_id: str = None, agent_interface: AgentInterface = AgentInterface.DEFAULT, package_requirements: list = []):
+        """
+        Validates the workflow graph for an AI Agent.
+
+        Args:
+            workflow_graph (WorkflowGraph): The workflow graph to validate.
+            sample_user_inputs (dict): Contains sample values for variables of type user_input for starting node
+            agent_workflow_node_id (str): Node id from which we want to run workflow
+            agent_interface (AgentInterface): The interface that the agent will be deployed with.
+            package_requirements (list): A list of package requirement strings. For example: ['numpy==1.2.3', 'pandas>=1.4.0'].
+
+        Returns:
+            dict: The output variables for every node in workflow which has executed.
+        """
+        graph_info = self.extract_graph_information(
+            workflow_graph=workflow_graph, agent_interface=agent_interface, package_requirements=package_requirements)
+        workflow_vars = get_object_from_context(
+            self, _request_context, 'workflow_vars', dict) or {}
+        topological_dfs_stack = get_object_from_context(
+            self, _request_context, 'topological_dfs_stack', list) or []
+        workflow_info = run(nodes=workflow_graph['nodes'], primary_start_node=workflow_graph['primary_start_node'], graph_info=graph_info,
+                            sample_user_inputs=sample_user_inputs, agent_workflow_node_id=agent_workflow_node_id, workflow_vars=workflow_vars, topological_dfs_stack=topological_dfs_stack)
+        _request_context.workflow_vars = workflow_info['workflow_vars']
+        _request_context.topological_dfs_stack = workflow_info['topological_dfs_stack']
+        return workflow_info['run_info']
+
     def create_agent_from_function(self, project_id: str, agent_function: callable, name: str = None, memory: int = None, package_requirements: list = None, description: str = None, evaluation_feature_group_id: str = None, workflow_graph: WorkflowGraph = None):
         """
         [Deprecated]
@@ -3883,6 +3918,7 @@ class ApiClient(ReadOnlyClient):
             body['extraArgs'] = extra_args
         headers = {'APIKEY': self.api_key}
         body['connectionId'] = uuid4().hex
+        self._clean_api_objects(body)
         for _ in range(3):
             response = self._request(
                 api_endpont, method='POST', body=body, headers=headers)
@@ -5360,7 +5396,7 @@ Creates a new feature group defined as the union of other feature group versions
             DatasetVersion: The new Dataset Version created."""
         return self._call_api('createDatasetVersionFromFileConnector', 'POST', query_params={'datasetId': dataset_id}, body={'location': location, 'fileFormat': file_format, 'csvDelimiter': csv_delimiter, 'mergeFileSchemas': merge_file_schemas, 'parsingConfig': parsing_config}, parse_type=DatasetVersion)
 
-    def create_dataset_from_database_connector(self, table_name: str, database_connector_id: str, object_name: str = None, columns: str = None, query_arguments: str = None, refresh_schedule: str = None, sql_query: str = None, incremental: bool = False, incremental_database_connector_config: Union[dict, IncrementalDatabaseConnectorConfig] = None) -> Dataset:
+    def create_dataset_from_database_connector(self, table_name: str, database_connector_id: str, object_name: str = None, columns: str = None, query_arguments: str = None, refresh_schedule: str = None, sql_query: str = None, incremental: bool = False, incremental_database_connector_config: Union[dict, IncrementalDatabaseConnectorConfig] = None, document_processing_config: Union[dict, DatasetDocumentProcessingConfig] = None) -> Dataset:
         """Creates a dataset from a Database Connector.
 
         Args:
@@ -5373,10 +5409,11 @@ Creates a new feature group defined as the union of other feature group versions
             sql_query (str): The full SQL query to use when fetching data. If present, this parameter will override `object_name`, `columns`, `timestamp_column`, and `query_arguments`.
             incremental (bool): Signifies if the dataset is an incremental dataset.
             incremental_database_connector_config (IncrementalDatabaseConnectorConfig): The config for incremental datasets. Only valid if incremental is True
+            document_processing_config (DatasetDocumentProcessingConfig): The document processing configuration. Only valid when documents are being imported (e.g. importing KnowledgeArticleDescriptions via Salesforce).
 
         Returns:
             Dataset: The created dataset."""
-        return self._call_api('createDatasetFromDatabaseConnector', 'POST', query_params={}, body={'tableName': table_name, 'databaseConnectorId': database_connector_id, 'objectName': object_name, 'columns': columns, 'queryArguments': query_arguments, 'refreshSchedule': refresh_schedule, 'sqlQuery': sql_query, 'incremental': incremental, 'incrementalDatabaseConnectorConfig': incremental_database_connector_config}, parse_type=Dataset)
+        return self._call_api('createDatasetFromDatabaseConnector', 'POST', query_params={}, body={'tableName': table_name, 'databaseConnectorId': database_connector_id, 'objectName': object_name, 'columns': columns, 'queryArguments': query_arguments, 'refreshSchedule': refresh_schedule, 'sqlQuery': sql_query, 'incremental': incremental, 'incrementalDatabaseConnectorConfig': incremental_database_connector_config, 'documentProcessingConfig': document_processing_config}, parse_type=Dataset)
 
     def create_dataset_from_application_connector(self, table_name: str, application_connector_id: str, dataset_config: Union[dict, ApplicationConnectorDatasetConfig] = None, refresh_schedule: str = None) -> Dataset:
         """Creates a dataset from an Application Connector.
@@ -6198,7 +6235,7 @@ Creates a new feature group defined as the union of other feature group versions
         return self._call_api('createPredictionOperator', 'POST', query_params={}, body={'name': name, 'projectId': project_id, 'sourceCode': source_code, 'predictFunctionName': predict_function_name, 'initializeFunctionName': initialize_function_name, 'featureGroupIds': feature_group_ids, 'cpuSize': cpu_size, 'memory': memory, 'packageRequirements': package_requirements, 'useGpu': use_gpu}, parse_type=PredictionOperator)
 
     def update_prediction_operator(self, prediction_operator_id: str, name: str = None, feature_group_ids: List = None, source_code: str = None, initialize_function_name: str = None, predict_function_name: str = None, cpu_size: str = None, memory: int = None, package_requirements: list = None, use_gpu: bool = None) -> PredictionOperator:
-        """Update an existing prediction operator.
+        """Update an existing prediction operator. This does not create a new version.
 
         Args:
             prediction_operator_id (str): The unique ID of the prediction operator.
@@ -8247,6 +8284,14 @@ Creates a new feature group defined as the union of other feature group versions
             Agent: The updated agent."""
         return self._call_api('updateAgent', 'POST', query_params={}, body={'modelId': model_id, 'functionSourceCode': function_source_code, 'agentFunctionName': agent_function_name, 'memory': memory, 'packageRequirements': package_requirements, 'description': description, 'enableBinaryInput': enable_binary_input, 'agentInputSchema': agent_input_schema, 'agentOutputSchema': agent_output_schema, 'workflowGraph': workflow_graph, 'agentInterface': agent_interface, 'includedModules': included_modules}, parse_type=Agent)
 
+    def generate_agent_code(self, project_id: str, prompt: str) -> list:
+        """Generates the code for defining an AI Agent
+
+        Args:
+            project_id (str): The unique ID associated with the project.
+            prompt (str): A natural language prompt which describes agent specification. Describe what the agent will do, what inputs it will expect, and what outputs it will give out"""
+        return self._call_api('generateAgentCode', 'POST', query_params={}, body={'projectId': project_id, 'prompt': prompt})
+
     def 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) -> LlmResponse:
         """Generate response to the prompt using the specified model.
 
@@ -8303,7 +8348,7 @@ Creates a new feature group defined as the union of other feature group versions
 
         Returns:
             ExtractedFields: The response from the document query."""
-        return self._call_api('extractDataUsingLLM', 'POST', query_params={}, body={'fieldDescriptors': field_descriptors, 'documentId': document_id, 'documentText': document_text, 'llmName': llm_name}, parse_type=ExtractedFields)
+        return self._proxy_request('ExtractDataUsingLlm', 'POST', query_params={}, body={'fieldDescriptors': field_descriptors, 'documentId': document_id, 'documentText': document_text, 'llmName': llm_name}, parse_type=ExtractedFields)
 
     def search_web_for_llm(self, queries: List, search_providers: List = None, max_results: int = 1, safe: bool = True, fetch_content: bool = False, max_page_tokens: int = 8192, convert_to_markdown: bool = True) -> WebSearchResponse:
         """Access web search providers to fetch content related to the queries for use in large language model inputs.
@@ -8350,6 +8395,24 @@ Creates a new feature group defined as the union of other feature group versions
             LlmResponse: The response from the LLM App."""
         return self._call_api('getLLMAppResponse', 'POST', query_params={}, body={'llmAppName': llm_app_name, 'prompt': prompt}, parse_type=LlmResponse)
 
+    def validate_workflow_graph(self, workflow_graph: Union[dict, WorkflowGraph], agent_interface: Union[AgentInterface, str] = AgentInterface.DEFAULT, package_requirements: list = []) -> dict:
+        """Validates the workflow graph for an AI Agent.
+
+        Args:
+            workflow_graph (WorkflowGraph): The workflow graph to validate.
+            agent_interface (AgentInterface): The interface that the agent will be deployed with.
+            package_requirements (list): A list of package requirement strings. For example: ['numpy==1.2.3', 'pandas>=1.4.0']."""
+        return self._call_api('validateWorkflowGraph', 'POST', query_params={}, body={'workflowGraph': workflow_graph, 'agentInterface': agent_interface, 'packageRequirements': package_requirements})
+
+    def extract_graph_information(self, workflow_graph: Union[dict, WorkflowGraph], agent_interface: Union[AgentInterface, str] = AgentInterface.DEFAULT, package_requirements: list = []) -> dict:
+        """Extracts information from the workflow.
+
+        Args:
+            workflow_graph (WorkflowGraph): The workflow graph to validate.
+            agent_interface (AgentInterface): The interface that the agent will be deployed with.
+            package_requirements (list): A list of package requirement strings. For example: ['numpy==1.2.3', 'pandas>=1.4.0']."""
+        return self._call_api('extractGraphInformation', 'POST', query_params={}, body={'workflowGraph': workflow_graph, 'agentInterface': agent_interface, 'packageRequirements': package_requirements})
+
     def create_document_retriever(self, project_id: str, name: str, feature_group_id: str, document_retriever_config: Union[dict, VectorStoreConfig] = None) -> DocumentRetriever:
         """Returns a document retriever that stores embeddings for document chunks in a feature group.
 

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/data_upload_result.py

@@ -8,15 +8,18 @@ class DataUploadResult(AbstractApiClass):
         Args:
             client (ApiClient): An authenticated API Client instance
             docInfos (list[agentdatadocumentinfo]): A list of dict for information on the documents uploaded to agent.
+            maxCount (int): The maximum number of documents
     """
 
-    def __init__(self, client, docInfos=None):
+    def __init__(self, client, docInfos=None, maxCount=None):
         super().__init__(client, None)
         self.doc_infos = docInfos
+        self.max_count = maxCount
         self.deprecated_keys = {}
 
     def __repr__(self):
-        repr_dict = {f'doc_infos': repr(self.doc_infos)}
+        repr_dict = {f'doc_infos': repr(
+            self.doc_infos), f'max_count': repr(self.max_count)}
         class_name = "DataUploadResult"
         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])
@@ -29,5 +32,5 @@ class DataUploadResult(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'doc_infos': self.doc_infos}
+        resp = {'doc_infos': self.doc_infos, 'max_count': self.max_count}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/model_version.py

@@ -21,6 +21,8 @@ class ModelVersion(AbstractApiClass):
             trainingStartedAt (str): The start time and date of the training process in ISO-8601 format.
             trainingCompletedAt (str): The end time and date of the training process in ISO-8601 format.
             featureGroupVersions (list): A list of Feature Group version IDs used for model training.
+            customAlgorithms (list): List of user-defined algorithms used for model training.
+            builtinAlgorithms (list): List of algorithm names builtin algorithms provided by Abacus.AI used for model training.
             error (str): Relevant error if the status is FAILED.
             pendingDeploymentIds (list): List of deployment IDs where deployment is pending.
             failedDeploymentIds (list): List of failed deployment IDs.
@@ -45,7 +47,7 @@ class ModelVersion(AbstractApiClass):
             deployableAlgorithms (DeployableAlgorithm): List of deployable algorithms.
     """
 
-    def __init__(self, client, modelVersion=None, modelConfigType=None, status=None, modelId=None, modelPredictionConfig=None, trainingStartedAt=None, trainingCompletedAt=None, featureGroupVersions=None, error=None, pendingDeploymentIds=None, failedDeploymentIds=None, cpuSize=None, memory=None, automlComplete=None, trainingFeatureGroupIds=None, trainingDocumentRetrieverVersions=None, documentRetrieverMappings=None, bestAlgorithm=None, defaultAlgorithm=None, featureAnalysisStatus=None, dataClusterInfo=None, customAlgorithmConfigs=None, trainedModelTypes=None, useGpu=None, partialComplete=None, modelFeatureGroupSchemaMappings=None, trainingConfigUpdated=None, codeSource={}, modelConfig={}, deployableAlgorithms={}):
+    def __init__(self, client, modelVersion=None, modelConfigType=None, status=None, modelId=None, modelPredictionConfig=None, trainingStartedAt=None, trainingCompletedAt=None, featureGroupVersions=None, customAlgorithms=None, builtinAlgorithms=None, error=None, pendingDeploymentIds=None, failedDeploymentIds=None, cpuSize=None, memory=None, automlComplete=None, trainingFeatureGroupIds=None, trainingDocumentRetrieverVersions=None, documentRetrieverMappings=None, bestAlgorithm=None, defaultAlgorithm=None, featureAnalysisStatus=None, dataClusterInfo=None, customAlgorithmConfigs=None, trainedModelTypes=None, useGpu=None, partialComplete=None, modelFeatureGroupSchemaMappings=None, trainingConfigUpdated=None, codeSource={}, modelConfig={}, deployableAlgorithms={}):
         super().__init__(client, modelVersion)
         self.model_version = modelVersion
         self.model_config_type = modelConfigType
@@ -55,6 +57,8 @@ class ModelVersion(AbstractApiClass):
         self.training_started_at = trainingStartedAt
         self.training_completed_at = trainingCompletedAt
         self.feature_group_versions = featureGroupVersions
+        self.custom_algorithms = customAlgorithms
+        self.builtin_algorithms = builtinAlgorithms
         self.error = error
         self.pending_deployment_ids = pendingDeploymentIds
         self.failed_deployment_ids = failedDeploymentIds
@@ -82,7 +86,7 @@ class ModelVersion(AbstractApiClass):
         self.deprecated_keys = {}
 
     def __repr__(self):
-        repr_dict = {f'model_version': repr(self.model_version), f'model_config_type': repr(self.model_config_type), f'status': repr(self.status), f'model_id': repr(self.model_id), f'model_prediction_config': repr(self.model_prediction_config), f'training_started_at': repr(self.training_started_at), f'training_completed_at': repr(self.training_completed_at), f'feature_group_versions': repr(self.feature_group_versions), f'error': repr(self.error), f'pending_deployment_ids': repr(self.pending_deployment_ids), f'failed_deployment_ids': repr(self.failed_deployment_ids), f'cpu_size': repr(self.cpu_size), f'memory': repr(self.memory), f'automl_complete': repr(self.automl_complete), f'training_feature_group_ids': repr(self.training_feature_group_ids), f'training_document_retriever_versions': repr(
+        repr_dict = {f'model_version': repr(self.model_version), f'model_config_type': repr(self.model_config_type), f'status': repr(self.status), f'model_id': repr(self.model_id), f'model_prediction_config': repr(self.model_prediction_config), f'training_started_at': repr(self.training_started_at), f'training_completed_at': repr(self.training_completed_at), f'feature_group_versions': repr(self.feature_group_versions), f'custom_algorithms': repr(self.custom_algorithms), f'builtin_algorithms': repr(self.builtin_algorithms), f'error': repr(self.error), f'pending_deployment_ids': repr(self.pending_deployment_ids), f'failed_deployment_ids': repr(self.failed_deployment_ids), f'cpu_size': repr(self.cpu_size), f'memory': repr(self.memory), f'automl_complete': repr(self.automl_complete), f'training_feature_group_ids': repr(self.training_feature_group_ids), f'training_document_retriever_versions': repr(
             self.training_document_retriever_versions), f'document_retriever_mappings': repr(self.document_retriever_mappings), f'best_algorithm': repr(self.best_algorithm), f'default_algorithm': repr(self.default_algorithm), f'feature_analysis_status': repr(self.feature_analysis_status), f'data_cluster_info': repr(self.data_cluster_info), f'custom_algorithm_configs': repr(self.custom_algorithm_configs), f'trained_model_types': repr(self.trained_model_types), f'use_gpu': repr(self.use_gpu), f'partial_complete': repr(self.partial_complete), f'model_feature_group_schema_mappings': repr(self.model_feature_group_schema_mappings), f'training_config_updated': repr(self.training_config_updated), f'code_source': repr(self.code_source), f'model_config': repr(self.model_config), f'deployable_algorithms': repr(self.deployable_algorithms)}
         class_name = "ModelVersion"
         repr_str = ',\n  '.join([f'{key}={value}' for key, value in repr_dict.items(
@@ -96,8 +100,8 @@ class ModelVersion(AbstractApiClass):
         Returns:
             dict: The dict value representation of the class parameters
         """
-        resp = {'model_version': self.model_version, 'model_config_type': self.model_config_type, 'status': self.status, 'model_id': self.model_id, 'model_prediction_config': self.model_prediction_config, 'training_started_at': self.training_started_at, 'training_completed_at': self.training_completed_at, 'feature_group_versions': self.feature_group_versions, 'error': self.error, 'pending_deployment_ids': self.pending_deployment_ids, 'failed_deployment_ids': self.failed_deployment_ids, 'cpu_size': self.cpu_size, 'memory': self.memory, 'automl_complete': self.automl_complete, 'training_feature_group_ids': self.training_feature_group_ids, 'training_document_retriever_versions': self.training_document_retriever_versions, 'document_retriever_mappings':
-                self.document_retriever_mappings, 'best_algorithm': self.best_algorithm, 'default_algorithm': self.default_algorithm, 'feature_analysis_status': self.feature_analysis_status, 'data_cluster_info': self.data_cluster_info, 'custom_algorithm_configs': self.custom_algorithm_configs, 'trained_model_types': self.trained_model_types, 'use_gpu': self.use_gpu, 'partial_complete': self.partial_complete, 'model_feature_group_schema_mappings': self.model_feature_group_schema_mappings, 'training_config_updated': self.training_config_updated, 'code_source': self._get_attribute_as_dict(self.code_source), 'model_config': self._get_attribute_as_dict(self.model_config), 'deployable_algorithms': self._get_attribute_as_dict(self.deployable_algorithms)}
+        resp = {'model_version': self.model_version, 'model_config_type': self.model_config_type, 'status': self.status, 'model_id': self.model_id, 'model_prediction_config': self.model_prediction_config, 'training_started_at': self.training_started_at, 'training_completed_at': self.training_completed_at, 'feature_group_versions': self.feature_group_versions, 'custom_algorithms': self.custom_algorithms, 'builtin_algorithms': self.builtin_algorithms, 'error': self.error, 'pending_deployment_ids': self.pending_deployment_ids, 'failed_deployment_ids': self.failed_deployment_ids, 'cpu_size': self.cpu_size, 'memory': self.memory, 'automl_complete': self.automl_complete, 'training_feature_group_ids': self.training_feature_group_ids, 'training_document_retriever_versions': self.training_document_retriever_versions,
+                'document_retriever_mappings': self.document_retriever_mappings, 'best_algorithm': self.best_algorithm, 'default_algorithm': self.default_algorithm, 'feature_analysis_status': self.feature_analysis_status, 'data_cluster_info': self.data_cluster_info, 'custom_algorithm_configs': self.custom_algorithm_configs, 'trained_model_types': self.trained_model_types, 'use_gpu': self.use_gpu, 'partial_complete': self.partial_complete, 'model_feature_group_schema_mappings': self.model_feature_group_schema_mappings, 'training_config_updated': self.training_config_updated, 'code_source': self._get_attribute_as_dict(self.code_source), 'model_config': self._get_attribute_as_dict(self.model_config), 'deployable_algorithms': self._get_attribute_as_dict(self.deployable_algorithms)}
         return {key: value for key, value in resp.items() if value is not None and key not in self.deprecated_keys}
 
     def describe_train_test_data_split_feature_group_version(self):

{abacusai-1.4.2 → abacusai-1.4.4}/abacusai/prediction_operator.py

@@ -95,7 +95,7 @@ class PredictionOperator(AbstractApiClass):
 
     def update(self, name: str = None, feature_group_ids: List = None, source_code: str = None, initialize_function_name: str = None, predict_function_name: str = None, cpu_size: str = None, memory: int = None, package_requirements: list = None, use_gpu: bool = None):
         """
-        Update an existing prediction operator.
+        Update an existing prediction operator. This does not create a new version.
 
         Args:
             name (str): Name of the prediction operator.