mmrech
diff --git a/‎abacusai/__init__.py
Lines changed: 1 addition & 1 deletion b/‎abacusai/__init__.py
Lines changed: 1 addition & 1 deletion
diff --git a/‎abacusai/api_class/ai_agents.py
Lines changed: 4 additions & 4 deletions b/‎abacusai/api_class/ai_agents.py
Lines changed: 4 additions & 4 deletions
diff --git a/‎abacusai/api_class/enums.py
Lines changed: 3 additions & 11 deletions b/‎abacusai/api_class/enums.py
Lines changed: 3 additions & 11 deletions
diff --git a/‎abacusai/api_class/model.py
Lines changed: 6 additions & 2 deletions b/‎abacusai/api_class/model.py
Lines changed: 6 additions & 2 deletions
diff --git a/‎abacusai/api_client_utils.py
Lines changed: 188 additions & 0 deletions b/‎abacusai/api_client_utils.py
Lines changed: 188 additions & 0 deletions
@@ -4,4 +4,4 @@
 from .streaming_client import StreamingClient
 
 
-__version__ = "1.4.3"
+__version__ = "1.4.4"
@@ -126,7 +126,7 @@ 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
         }
@@ -139,7 +139,7 @@ def from_dict(cls, mapping: dict):
             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)
         )
@@ -160,7 +160,7 @@ class WorkflowNodeOutputMapping(ApiClass):
     def to_dict(self):
         return {
             'name': self.name,
-            'variable_type': self.variable_type
+            'variable_type': self.variable_type.value
         }
 
     @classmethod
@@ -171,7 +171,7 @@ def from_dict(cls, mapping: dict):
             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'))
         )
 
 
 
@@ -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,28 +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'
 
 
@@ -460,15 +460,16 @@ 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.
         query_rewrite_instructions (str): Special instructions for the LLM which rewrites the RAG query.
         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>"
@@ -478,6 +479,7 @@ class ChatLLMTrainingConfig(TrainingConfig):
         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)
@@ -495,6 +497,7 @@ 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)
@@ -506,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
 
@@ -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'
Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,4 @@`
`4`	`4`	`from .streaming_client import StreamingClient`
`5`	`5`
`6`	`6`
`7`		`-__version__ = "1.4.3"`
	`7`	`+__version__ = "1.4.4"`