[DEVX-210]: YAML based workflow creation (#175)

Author: sainivedh

README.md

@@ -195,7 +195,23 @@ all_workflow = app.list_workflow()
 # List all workflow in community filtered by description
 all_face_community_workflows = App().list_workflows(filter_by={"query": "face"}, only_in_app=False) # Get all face related workflows
 ```
+#### Workflow Create
+Create a new workflow specified by a yaml config file.
+```python
+# Note: CLARIFAI_PAT must be set as env variable.
+from clarifai.client.app import App
+app = App(app_id="app_id", user_id="user_id")
+workflow = app.create_workflow(config_filepath="config.yml")
+```
 
+#### Workflow Export
+Export an existing workflow from Clarifai as a local yaml file.
+```python
+# Note: CLARIFAI_PAT must be set as env variable.
+from clarifai.client.workflow import Workflow
+workflow = Workflow("https://clarifai.com/clarifai/main/workflows/Demographics")
+workflow.export('demographics_workflow.yml')
+```
 
 ## More Examples
 See many more code examples in this [repo](https://github.com/Clarifai/examples).

clarifai/client/app.py

@@ -1,5 +1,8 @@
+import os
+import uuid
 from typing import Any, Dict, List
 
+import yaml
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2
 from google.protobuf.json_format import MessageToDict
@@ -13,7 +16,9 @@
 from clarifai.client.workflow import Workflow
 from clarifai.errors import UserError
 from clarifai.urls.helper import ClarifaiUrlHelper
-from clarifai.utils.logging import get_logger
+from clarifai.utils.logging import display_workflow_tree, get_logger
+from clarifai.workflows.utils import get_yaml_output_info_proto, is_same_yaml_model
+from clarifai.workflows.validate import validate
 
 
 class App(Lister, BaseClient):
@@ -236,30 +241,98 @@ def create_model(self, model_id: str, **kwargs) -> Model:
 
     return Model(model_id=model_id, **kwargs)
 
-  def create_workflow(self, workflow_id: str, **kwargs) -> Workflow:
+  def create_workflow(self,
+                      config_filepath: str,
+                      generate_new_id: bool = False,
+                      display: bool = True) -> Workflow:
     """Creates a workflow for the app.
 
     Args:
-        workflow_id (str): The workflow ID for the workflow to create.
-        **kwargs: Additional keyword arguments to be passed to the workflow.
+        config_filepath (str): The path to the yaml workflow config file.
+        generate_new_id (bool): If True, generate a new workflow ID.
+        display (bool): If True, display the workflow nodes tree.
 
     Returns:
-        Workflow: A Workflow object for the specified workflow ID.
+        Workflow: A Workflow object for the specified workflow config.
 
     Example:
         >>> from clarifai.client.app import App
         >>> app = App(app_id="app_id", user_id="user_id")
-        >>> workflow = app.create_workflow(workflow_id="workflow_id")
+        >>> workflow = app.create_workflow(config_filepath="config.yml")
     """
+    if not os.path.exists(config_filepath):
+      raise UserError(f"Workflow config file not found at {config_filepath}")
+
+    with open(config_filepath, 'r') as file:
+      data = yaml.safe_load(file)
+
+    data = validate(data)
+    workflow = data['workflow']
+
+    # Get all model objects from the workflow nodes.
+    all_models = []
+    for node in workflow['nodes']:
+      output_info = get_yaml_output_info_proto(node['model'].get('output_info', None))
+      try:
+        model = self.model(
+            node['model']['model_id'],
+            node['model'].get('model_version_id', ""),
+            user_id=node['model'].get('user_id', ""),
+            app_id=node['model'].get('app_id', ""))
+      except Exception as e:
+        if "Model does not exist" in str(e):
+          model = self.create_model(
+              **{k: v
+                 for k, v in node['model'].items() if k != 'output_info'})
+          model_version = model.create_model_version(
+              model_id=node['model']['model_id'], output_info=output_info)
+          all_models.append(model_version.model_info)
+          continue
+
+      # If the model version ID is specified, or if the yaml model is the same as the one in the api
+      if node["model"].get("model_version_id", "") or is_same_yaml_model(
+          model.model_info, node["model"]):
+        all_models.append(model.model_info)
+      else:  # Create a new model version
+        model = model.create_model_version(
+            model_id=node['model']['model_id'], output_info=output_info)
+        all_models.append(model.model_info)
+
+    # Convert nodes to resources_pb2.WorkflowNodes.
+    nodes = []
+    for i, yml_node in enumerate(workflow['nodes']):
+      node = resources_pb2.WorkflowNode(
+          id=yml_node['id'],
+          model=all_models[i],
+      )
+      # Add node inputs if they exist, i.e. if these nodes do not connect directly to the input.
+      if yml_node.get("node_inputs"):
+        for ni in yml_node.get("node_inputs"):
+          node.node_inputs.append(resources_pb2.NodeInput(node_id=ni['node_id']))
+      nodes.append(node)
+
+    workflow_id = workflow['id']
+    if generate_new_id:
+      workflow_id = str(uuid.uuid4())
+
+    # Create the workflow.
     request = service_pb2.PostWorkflowsRequest(
-        user_app_id=self.user_app_id, workflows=[resources_pb2.Workflow(id=workflow_id, **kwargs)])
+        user_app_id=self.user_app_id,
+        workflows=[resources_pb2.Workflow(id=workflow_id, nodes=nodes)])
+
     response = self._grpc_request(self.STUB.PostWorkflows, request)
     if response.status.code != status_code_pb2.SUCCESS:
       raise Exception(response.status)
     self.logger.info("\nWorkflow created\n%s", response.status)
-    kwargs.update({'app_id': self.id, 'user_id': self.user_id})
 
-    return Workflow(workflow_id=workflow_id, **kwargs)
+    dict_response = MessageToDict(response, preserving_proto_field_name=True)
+    # Display the workflow nodes tree.
+    if display:
+      display_workflow_tree(dict_response["workflows"][0]["nodes"])
+    kwargs = self.process_response_keys(dict_response[list(dict_response.keys())[1]][0],
+                                        "workflow")
+
+    return Workflow(**kwargs)
 
   def create_module(self, module_id: str, description: str, **kwargs) -> Module:
     """Creates a module for the app.
@@ -329,8 +402,16 @@ def model(self, model_id: str, model_version_id: str = "", **kwargs) -> Model:
         >>> app = App(app_id="app_id", user_id="user_id")
         >>> model_v1 = app.model(model_id="model_id", model_version_id="model_version_id")
     """
-    request = service_pb2.GetModelRequest(
-        user_app_id=self.user_app_id, model_id=model_id, version_id=model_version_id)
+    # Change user_app_id based on whether user_id or app_id is specified.
+    if kwargs.get("user_id") or kwargs.get("app_id"):
+      request = service_pb2.GetModelRequest(
+          user_app_id=self.auth_helper.get_user_app_id_proto(
+              kwargs.get("user_id"), kwargs.get("app_id")),
+          model_id=model_id,
+          version_id=model_version_id)
+    else:
+      request = service_pb2.GetModelRequest(
+          user_app_id=self.user_app_id, model_id=model_id, version_id=model_version_id)
     response = self._grpc_request(self.STUB.GetModel, request)
 
     if response.status.code != status_code_pb2.SUCCESS:

clarifai/client/base.py

@@ -2,7 +2,7 @@
 from datetime import datetime
 from typing import Any, Callable
 
-from google.protobuf.json_format import MessageToDict  # noqa
+from google.protobuf import struct_pb2
 from google.protobuf.timestamp_pb2 import Timestamp
 from google.protobuf.wrappers_pb2 import BoolValue
 
@@ -71,7 +71,10 @@ def convert_string_to_timestamp(self, date_str) -> Timestamp:
     try:
       datetime_obj = datetime.strptime(date_str, '%Y-%m-%dT%H:%M:%S.%fZ')
     except ValueError:
-      datetime_obj = datetime.strptime(date_str, '%Y-%m-%dT%H:%M:%SZ')
+      try:
+        datetime_obj = datetime.strptime(date_str, '%Y-%m-%dT%H:%M:%SZ')
+      except ValueError:
+        return Timestamp()
 
     # Convert the datetime object to a Timestamp object
     timestamp_obj = Timestamp()
@@ -99,8 +102,12 @@ def convert_recursive(item):
             value = self.convert_string_to_timestamp(value)
           elif key in ['workflow_recommended']:
             value = BoolValue(value=True)
-          elif key in ['metadata', 'fields_map', 'params']:
-            continue  # TODO Fix "app_duplication",proto struct
+          elif key in ['fields_map', 'params']:
+            value_s = struct_pb2.Struct()
+            value_s.update(value)
+            value = value_s
+          elif key in ['metadata']:
+            continue  # TODO Fix "app_duplication"
           new_item[key] = convert_recursive(value)
         return new_item
       elif isinstance(item, list):

clarifai/client/model.py

@@ -5,6 +5,7 @@
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.resources_pb2 import Input
 from clarifai_grpc.grpc.api.status import status_code_pb2
+from google.protobuf.json_format import MessageToDict
 
 from clarifai.client.base import BaseClient
 from clarifai.client.lister import Lister
@@ -52,6 +53,73 @@ def __init__(self,
     BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id)
     Lister.__init__(self)
 
+  def create_model_version(self, model_id: str, **kwargs) -> 'Model':
+    """Creates a model version for the Model.
+
+    Args:
+        model_id (str): The model ID for the model to create.
+        **kwargs: Additional keyword arguments to be passed to Model Version.
+          - description (str): The description of the model version.
+          - concepts (list[Concept]): The concepts to associate with the model version.
+          - output_info (resources_pb2.OutputInfo(): The output info to associate with the model version.
+
+    Returns:
+        Model: A Model object for the specified model ID.
+
+    Example:
+        >>> from clarifai.client.model import Model
+        >>> model = Model("model_url")
+                    or
+        >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
+        >>> model_version = model.create_model_version(model_id='model_id', description='model_version_description')
+    """
+    request = service_pb2.PostModelVersionsRequest(
+        user_app_id=self.user_app_id,
+        model_id=model_id,
+        model_versions=[resources_pb2.ModelVersion(**kwargs)])
+
+    response = self._grpc_request(self.STUB.PostModelVersions, request)
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(response.status)
+    self.logger.info("\nModel Version created\n%s", response.status)
+
+    kwargs.update({'app_id': self.app_id, 'user_id': self.user_id})
+    dict_response = MessageToDict(response, preserving_proto_field_name=True)
+    kwargs = self.process_response_keys(dict_response['model'], 'model')
+
+    return Model(**kwargs)
+
+  def list_versions(self) -> List['Model']:
+    """Lists all the versions for the model.
+
+    Returns:
+        List[Model]: A list of Model objects for the versions of the model.
+
+    Example:
+        >>> from clarifai.client.model import Model
+        >>> model = Model("model_url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
+                    or
+        >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
+        >>> all_model_versions = model.list_versions()
+    """
+    request_data = dict(
+        user_app_id=self.user_app_id,
+        model_id=self.id,
+        per_page=self.default_page_size,
+    )
+    all_model_versions_info = list(
+        self.list_all_pages_generator(self.STUB.ListModelVersions,
+                                      service_pb2.ListModelVersionsRequest, request_data))
+
+    for model_version_info in all_model_versions_info:
+      model_version_info['id'] = model_version_info['model_version_id']
+      del model_version_info['model_version_id']
+
+    return [
+        Model(model_id=self.id, **dict(self.kwargs, model_version=model_version_info))
+        for model_version_info in all_model_versions_info
+    ]
+
   def predict(self, inputs: List[Input]):
     """Predicts the model based on the given inputs.
 

clarifai/client/workflow.py

@@ -10,6 +10,7 @@
 from clarifai.errors import UserError
 from clarifai.urls.helper import ClarifaiUrlHelper
 from clarifai.utils.logging import get_logger
+from clarifai.workflows.export import Exporter
 
 
 class Workflow(Lister, BaseClient):
@@ -182,6 +183,28 @@ def list_versions(self) -> List['Workflow']:
         for workflow_version_info in all_workflow_versions_info
     ]
 
+  def export(self, out_path: str):
+    """Exports the workflow to a yaml file.
+
+    Args:
+        out_path (str): The path to save the yaml file to.
+
+    Example:
+        >>> from clarifai.client.workflow import Workflow
+        >>> workflow = Workflow("https://clarifai.com/clarifai/main/workflows/Demographics")
+        >>> workflow.export('out_path')
+    """
+    request = service_pb2.GetWorkflowRequest(user_app_id=self.user_app_id, workflow_id=self.id)
+    response = self._grpc_request(self.STUB.GetWorkflow, request)
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(f"Workflow Export failed with response {response.status!r}")
+
+    with Exporter(response) as e:
+      e.parse()
+      e.export(out_path)
+
+    self.logger.info(f"Exported workflow to {out_path}")
+
   def __getattr__(self, name):
     return getattr(self.workflow_info, name)
 

clarifai/utils/logging.py

@@ -1,16 +1,66 @@
 import logging
-from typing import Optional
+from collections import defaultdict
+from typing import Dict, List, Optional
 
+from rich import print as rprint
 from rich.logging import RichHandler
 from rich.table import Table
 from rich.traceback import install
+from rich.tree import Tree
 
 install()
 
 
-def table_from_dict(data, column_names, title="") -> Table:
+def display_workflow_tree(nodes_data: List[Dict]) -> None:
+  """Displays a tree of the workflow nodes."""
+  # Create a mapping of node_id to the list of node_ids that are connected to it.
+  node_adj_mapping = defaultdict(list)
+  # Create a mapping of node_id to the node data info.
+  nodes_data_dict = {}
+  for node in nodes_data:
+    nodes_data_dict[node["id"]] = node
+    if node.get("node_inputs", "") == "":
+      node_adj_mapping["Input"].append(node["id"])
+    else:
+      for node_input in node["node_inputs"]:
+        node_adj_mapping[node_input["node_id"]].append(node["id"])
+
+  # Get all leaf nodes.
+  leaf_node_ids = set()
+  for node_id in list(nodes_data_dict.keys()):
+    if node_adj_mapping.get(node_id, "") == "":
+      leaf_node_ids.add(node_id)
+
+  def build_node_tree(node_id="Input"):
+    """Recursively builds a rich tree of the workflow nodes."""
+    # Set the style of the current node.
+    style_str = "green" if node_id in leaf_node_ids else "white"
+
+    # Create a Tree object for the current node.
+    if node_id != "Input":
+      node_table = table_from_dict(
+          [nodes_data_dict[node_id]["model"]],
+          column_names=["id", "model_type_id", "app_id", "user_id"],
+          title="Node: " + node_id)
+
+      tree = Tree(node_table, style=style_str, guide_style="underline2 white")
+    else:
+      tree = Tree(f"[green] {node_id}", style=style_str, guide_style="underline2 white")
+
+    # Recursively add the child nodes of the current node to the tree.
+    for child in node_adj_mapping.get(node_id, []):
+      tree.add(build_node_tree(child))
+
+    # Return the tree.
+    return tree
+
+  tree = build_node_tree("Input")
+  rprint(tree)
+
+
+def table_from_dict(data: List[Dict], column_names: List[str], title: str = "") -> Table:
   """Use this function for printing tables from a list of dicts."""
-  table = Table(title=title, show_header=True, header_style="bold blue")
+  table = Table(title=title, show_lines=False, show_header=True, header_style="blue")
   for column_name in column_names:
     table.add_column(column_name)
   for row in data: