cognite-neat 0.104.0__py3-none-any.whl → 0.105.1__py3-none-any.whl

cognite/neat/_workflows/examples/Validate_Solution_Model/workflow.yaml

@@ -1,64 +0,0 @@
-configs: []
-description: null
-implementation_module: null
-name: Validate Solution Model
-steps:
--   complex_configs: {}
-    configs: {}
-    description: null
-    enabled: true
-    id: step_769298
-    label: Upload File
-    max_retries: 0
-    method: null
-    params:
-        file_type: rules
-    retry_delay: 3
-    stype: file_uploader
-    system_component_id: null
-    transition_to:
-    - step_399494
-    trigger: true
-    ui_config:
-        pos_x: 627
-        pos_y: -19
--   complex_configs: {}
-    configs:
-        File name: ''
-        Report formatter: BasicHTML
-        Role: infer
-    description: null
-    enabled: true
-    id: step_399494
-    label: Import Excel
-    max_retries: 0
-    method: ExcelToRules
-    params: {}
-    retry_delay: 3
-    stype: stdstep
-    system_component_id: null
-    transition_to:
-    - step_273233
-    trigger: false
-    ui_config:
-        pos_x: 627
-        pos_y: 66
--   complex_configs: {}
-    configs:
-        Report Formatter: BasicHTML
-    description: null
-    enabled: true
-    id: step_273233
-    label: Validate Against CDF
-    max_retries: 0
-    method: ValidateRulesAgainstCDF
-    params: {}
-    retry_delay: 3
-    stype: stdstep
-    system_component_id: null
-    transition_to: []
-    trigger: false
-    ui_config:
-        pos_x: 627
-        pos_y: 154
-system_components: []

cognite/neat/_workflows/manager.py

@@ -1,292 +0,0 @@
-import logging
-import os
-import shutil
-import sys
-import time
-import traceback
-
-from cognite.client import CogniteClient
-from prometheus_client import Gauge
-from pydantic import BaseModel
-
-from cognite.neat._config import Config
-from cognite.neat._workflows import BaseWorkflow
-from cognite.neat._workflows.base import WorkflowDefinition
-from cognite.neat._workflows.model import FlowMessage, InstanceStartMethod, WorkflowState, WorkflowStepDefinition
-from cognite.neat._workflows.steps_registry import StepsRegistry
-from cognite.neat._workflows.tasks import WorkflowTaskBuilder
-
-live_workflow_instances = Gauge("neat_workflow_live_instances", "Count of live workflow instances", ["itype"])
-
-
-class WorkflowStartStatus(BaseModel, arbitrary_types_allowed=True):
-    workflow_instance: BaseWorkflow | None = None
-    is_success: bool = True
-    status_text: str | None = None
-
-
-class WorkflowManager:
-    """Workflow manager is responsible for loading, saving and managing workflows
-    client: CogniteClient
-    config: Config
-    """
-
-    def __init__(self, client: CogniteClient, config: Config):
-        self.client = client
-        self.data_set_id = config.cdf_default_dataset_id
-        self.data_store_path = config.data_store_path
-        self.workflow_registry: dict[str, BaseWorkflow] = {}
-        self.ephemeral_instance_registry: dict[str, BaseWorkflow] = {}
-        self.workflows_storage_type = config.workflows_store_type
-        self.config = config
-        self.workflows_storage_path = config.workflows_store_path
-        self.rules_storage_path = config.rules_store_path
-        self.task_builder = WorkflowTaskBuilder(client, self)
-        self.steps_registry = StepsRegistry(self.config)
-        self.steps_registry.load_step_classes()
-
-    def update_cdf_client(self, client: CogniteClient):
-        self.client = client
-        self.task_builder = WorkflowTaskBuilder(client, self)
-        self.workflow_registry = {}
-        self.load_workflows_from_storage()
-
-    def get_list_of_workflows(self):
-        return list(self.workflow_registry.keys())
-
-    def get_steps_registry(self):
-        return self.steps_registry
-
-    def get_workflow(self, name: str) -> BaseWorkflow | None:
-        try:
-            return self.workflow_registry[name]
-        except KeyError:
-            return None
-
-    def full_reset(self):
-        self.workflow_registry = {}
-        self.ephemeral_instance_registry = {}
-
-    def start_workflow(self, name: str, sync=False, **kwargs):
-        workflow = self.get_workflow(name)
-        if workflow is None:
-            raise ValueError(f"Workflow {name} not found")
-        workflow.start(sync, kwargs=kwargs)
-        return workflow
-
-    def delete_workflow(self, name: str):
-        workflow = self.get_workflow(name)
-        if workflow is not None:
-            workflow.cleanup_workflow_context()
-        del self.workflow_registry[name]
-        full_path = self.workflows_storage_path / name
-        shutil.rmtree(full_path)
-
-    def update_workflow(self, name: str, workflow: WorkflowDefinition):
-        self.workflow_registry[name].workflow_steps = workflow.steps
-        self.workflow_registry[name].configs = workflow.configs
-        self.workflow_registry[name].workflow_system_components = workflow.system_components
-
-    def save_workflow_to_storage(self, name: str, custom_implementation_module: str | None = None):
-        """Save workflow from memory to storage"""
-        if self.workflows_storage_type == "file":
-            full_path = self.workflows_storage_path / name / "workflow.yaml"
-            full_path.parent.mkdir(parents=True, exist_ok=True)
-            wf = self.workflow_registry[name]
-            full_path.write_text(
-                wf.serialize_workflow(output_format="yaml", custom_implementation_module=custom_implementation_module)
-            )
-
-    def create_new_workflow(self, name: str, description=None, mode: str = "manifest") -> WorkflowDefinition:
-        """Create new workflow in memory"""
-        workflow_definition = WorkflowDefinition(
-            name=name, description=description, steps=[], system_components=[], configs=[]
-        )
-        self.register_workflow(BaseWorkflow, name, workflow_definition)
-        self.save_workflow_to_storage(name)
-        return workflow_definition
-
-    def load_single_workflow_from_storage(self, workflow_name: str):
-        """Load single workflow from storage into memory"""
-
-        # check if workflow is already loaded. If so, perform context cleanup first
-        if workflow_name in self.workflow_registry:
-            self.workflow_registry[workflow_name].cleanup_workflow_context()
-            del self.workflow_registry[workflow_name]
-
-        workflow_path = self.workflows_storage_path / workflow_name
-        # THIS IS WHERE WE LOAD THE WORKFLOW MODULES
-        if workflow_path.is_dir():
-            logging.info(f"Loading workflow {workflow_name} from {workflow_path}")
-            try:
-                if (workflow_definition_path := workflow_path / "workflow.yaml").exists():
-                    with workflow_definition_path.open() as workflow_definition_file:
-                        workflow_definition: WorkflowDefinition = BaseWorkflow.deserialize_definition(
-                            workflow_definition_file.read(), output_format="yaml"
-                        )
-                else:
-                    logging.info(f"Definition file {workflow_definition_path} not found, skipping")
-                    return
-
-                # This is convinience feature when user wants to share self contained workflow
-                # folder as single zip file.
-                if (workflow_path / "rules").exists():
-                    logging.info(f"Copying rules from {workflow_path / 'rules'} to {self.rules_storage_path} ")
-                    for rule_file in os.listdir(workflow_path / "rules"):
-                        shutil.copy(workflow_path / "rules" / rule_file, self.rules_storage_path)
-
-                self.steps_registry.load_custom_step_classes(workflow_path, workflow_name)
-                self.register_workflow(BaseWorkflow, workflow_name, workflow_definition)
-
-            except Exception as e:
-                trace = traceback.format_exc()
-                logging.error(f"Error loading workflow {workflow_name}: error: {e} trace : {trace}")
-
-    def load_workflows_from_storage(self):
-        """Loads workflows from disk/storage into memory, initializes and register them in the workflow registry"""
-
-        # set workflow storage path
-        workflows_storage_path = self.workflows_storage_path
-
-        # set system path to be used when importing individual workflows as python modules
-        # via importlib.import_module(...)
-        sys.path.append(str(workflows_storage_path))
-
-        for workflow_name in os.listdir(workflows_storage_path):
-            self.load_single_workflow_from_storage(workflow_name)
-
-    def register_workflow(self, obj, workflow_name, workflow_definition):
-        """Register workflow in the workflow registry
-
-        Parameters
-        ----------
-        obj : BaseWorkflow
-            Class of the workflow to be registered
-        workflow_name : _type_
-            Name of the workflow to be registered
-        workflow_definition : _type_
-            Definition of the workflow to be registered originating from workflow.yaml
-        """
-        self.workflow_registry[workflow_name] = obj(workflow_name, self.client, steps_registry=self.steps_registry)
-        self.workflow_registry[workflow_name].set_definition(workflow_definition)
-        # Comment: Not entirely sure what is task_builder meant for
-        # as at first glance it looks like circular import ???
-        self.workflow_registry[workflow_name].set_task_builder(self.task_builder)
-        self.workflow_registry[workflow_name].set_default_dataset_id(self.data_set_id)
-        self.workflow_registry[workflow_name].set_storage_path("transformation_rules", self.rules_storage_path)
-        self.workflow_registry[workflow_name].set_storage_path("data_store", self.data_store_path)
-
-    def create_workflow_instance(self, template_name: str, add_to_registry: bool = True) -> BaseWorkflow:
-        new_instance = self.workflow_registry[template_name].copy()
-        if add_to_registry:
-            self.ephemeral_instance_registry[new_instance.instance_id] = new_instance
-        live_workflow_instances.labels(itype="ephemeral").set(len(self.ephemeral_instance_registry))
-        return new_instance
-
-    def get_workflow_instance(self, instance_id: str) -> BaseWorkflow:
-        return self.ephemeral_instance_registry[instance_id]
-
-    def delete_workflow_instance(self, instance_id: str):
-        del self.ephemeral_instance_registry[instance_id]
-        live_workflow_instances.labels(itype="ephemeral").set(len(self.ephemeral_instance_registry))
-        return
-
-    def start_workflow_instance(
-        self, workflow_name: str, step_id: str = "", flow_msg: FlowMessage | None = None, sync: bool | None = None
-    ) -> WorkflowStartStatus:
-        retrieved = self.get_workflow(workflow_name)
-
-        if retrieved is None:
-            return WorkflowStartStatus(
-                workflow_instance=None, is_success=False, status_text="Workflow not found in registry"
-            )
-
-        if self._is_workflow_made_of_mixed_steps(retrieved.workflow_steps):
-            retrieved.state = WorkflowState.FAILED
-            return WorkflowStartStatus(
-                workflow_instance=None,
-                is_success=False,
-                status_text="Workflow consists of both legacy and current steps. "
-                "Please update the workflow to use only current steps.",
-            )
-
-        workflow = retrieved
-        retrieved_step = workflow.get_trigger_step(step_id)
-        if retrieved_step is None:
-            return WorkflowStartStatus(
-                workflow_instance=None, is_success=False, status_text="Step not found in workflow"
-            )
-        trigger_step = retrieved_step
-        if not trigger_step.trigger:
-            logging.info(f"Step {step_id} is not a trigger step")
-            return WorkflowStartStatus(
-                workflow_instance=None, is_success=False, status_text="Step is not a trigger step"
-            )
-        if sync is None:
-            sync = trigger_step.params.get("sync", "true").lower() == "true"
-
-        max_wait_time = int(trigger_step.params.get("max_wait_time", "30"))
-        instance_start_method = trigger_step.params.get(
-            "workflow_start_method", InstanceStartMethod.PERSISTENT_INSTANCE_BLOCKING
-        )
-
-        logging.info(
-            f"""----- Starting workflow {workflow_name} , step_id = {step_id} , sync = {sync},
-              max_wait_time = {max_wait_time}, instance_start_method = {instance_start_method} -----"""
-        )
-
-        if instance_start_method == InstanceStartMethod.PERSISTENT_INSTANCE_BLOCKING:
-            live_workflow_instances.labels(itype="persistent").set(len(self.workflow_registry))
-            # wait until workflow transition to RUNNING state and then start , set max wait time to 30 seconds
-            start_time = time.perf_counter()
-            # wait until workflow transition to RUNNING state and then start , set max wait time to 30 seconds.
-            # The operation is executed in callers thread
-            while workflow.state == WorkflowState.RUNNING:
-                logging.info("Existing workflow instance already running , waiting for RUNNING state")
-                elapsed_time = time.perf_counter() - start_time
-                if elapsed_time > max_wait_time:
-                    logging.info(
-                        f"Workflow {workflow_name} wait time exceeded . "
-                        f"elapsed time = {elapsed_time}, max wait time = {max_wait_time}"
-                    )
-                    return WorkflowStartStatus(
-                        workflow_instance=None,
-                        is_success=False,
-                        status_text="Workflow instance already running.Wait time exceeded",
-                    )
-                time.sleep(0.5)
-            workflow.start(sync=sync, flow_message=flow_msg, start_step_id=step_id)
-            return WorkflowStartStatus(workflow_instance=workflow, is_success=True, status_text="")
-
-        elif instance_start_method == InstanceStartMethod.PERSISTENT_INSTANCE_NON_BLOCKING:
-            live_workflow_instances.labels(itype="persistent").set(len(self.workflow_registry))
-            # start workflow if not already running, skip if already running
-            if workflow.state == WorkflowState.RUNNING:
-                return WorkflowStartStatus(
-                    workflow_instance=None, is_success=False, status_text="Workflow instance already running"
-                )
-
-            workflow.start(sync=sync, flow_message=flow_msg, start_step_id=step_id)
-            return WorkflowStartStatus(workflow_instance=workflow, is_success=True, status_text="")
-
-        elif instance_start_method == InstanceStartMethod.EPHEMERAL_INSTANCE:
-            # start new workflow instance in new thread
-            workflow = self.create_workflow_instance(workflow_name, add_to_registry=True)
-            workflow.start(sync=sync, delete_after_completion=True, flow_message=flow_msg, start_step_id=step_id)
-            if sync:
-                self.delete_workflow_instance(workflow.instance_id)
-            return WorkflowStartStatus(workflow_instance=workflow, is_success=True, status_text="")
-
-        return WorkflowStartStatus(
-            workflow_instance=None, is_success=False, status_text="Unsupported workflow start method"
-        )
-
-    def _is_workflow_made_of_mixed_steps(self, steps: list[WorkflowStepDefinition]):
-        legacy_steps = 0
-        current_steps = 0
-        for step in steps:
-            if step.method in self.steps_registry.categorized_steps["legacy"]:
-                legacy_steps += 1
-            if step.method in self.steps_registry.categorized_steps["current"]:
-                current_steps += 1
-        return legacy_steps > 0 and current_steps > 0

cognite/neat/_workflows/model.py

@@ -1,203 +0,0 @@
-import sys
-from typing import Any
-
-if sys.version_info >= (3, 11):
-    from enum import StrEnum
-else:
-    from backports.strenum import StrEnum
-
-from pydantic import BaseModel, Field, field_validator
-
-
-class WorkflowState(StrEnum):
-    CREATED = "CREATED"
-    RUNNING = "RUNNING"
-    RUNNING_WAITING = "RUNNING_WAITING"
-    COMPLETED = "COMPLETED"
-    FAILED = "FAILED"
-    EXPIRED = "EXPIRED"
-    UNKNOWN = "UNKNOWN"
-
-
-class StepExecutionStatus(StrEnum):
-    SUCCESS = "COMPLETED"
-    FAILED = "FAILED"
-    SKIPPED = "SKIPPED"
-    STARTED = "STARTED"
-    UNKNOWN = "UNKNOWN"
-    ABORT_AND_FAIL = "ABORT_AND_FAIL"
-
-
-class InstanceStartMethod(StrEnum):
-    PERSISTENT_INSTANCE_NON_BLOCKING = "persistent_non_blocking"
-    PERSISTENT_INSTANCE_BLOCKING = "persistent_blocking"
-    EPHEMERAL_INSTANCE = "ephemeral_instance"
-
-
-class WorkflowStartException(Exception):
-    pass
-
-
-class FlowMessage(BaseModel):
-    """A message that can be sent between steps in a workflow.It's the only parameter step takes as input."""
-
-    payload: Any = None  # The payload of the message
-    headers: dict[str, str] | None = None  # The headers of the message
-    output_text: str | None = None  # The output text of the step that is captured in the execution log
-    error_text: str | None = None  # The error text of the step that is captured in the execution log
-    next_step_ids: list[str] | None = (
-        None  # If set, the workflow will skip default route and go to the next step in the list
-    )
-    step_execution_status: StepExecutionStatus = StepExecutionStatus.UNKNOWN  # The status of the step execution
-
-
-class StepType(StrEnum):
-    PYSTEP = "pystep"
-    STD_STEP = "stdstep"
-    START_WORKFLOW_TASK_STEP = "start_workflow_task_step"  # Call a workflow from another workflow
-    CONTAINER_TASK_STEP = "container_step"  # TODO: implement
-    CLI_TASK_STEP = "cli_step"  # TODO: implement
-    HTTP_TRIGGER = "http_trigger"
-    TIME_TRIGGER = "time_trigger"
-    WAIT_FOR_EVENT = "wait_for_event"  # TODO: implement
-    EVENT_TRIGGER = "event_trigger"  # TODO: implement
-
-
-class UIConfig(BaseModel):
-    pos_x: int = 0
-    pos_y: int = 0
-
-
-class WorkflowConfigItem(BaseModel):
-    name: str
-    value: str | None = None
-    label: str | None = None
-    type: str | None = None
-    required: bool = False
-    options: list[str] | None = None
-    group: str | None = None
-
-
-class WorkflowStepDefinition(BaseModel):
-    id: str
-    label: str | None = None
-    stype: str = StepType.PYSTEP
-    description: str | None = None
-    method: str | None = None
-    enabled: bool = True
-    system_component_id: str | None = None
-    trigger: bool = False
-    transition_to: list[str] | None = None
-    ui_config: UIConfig = UIConfig()
-    params: dict[str, str] = Field(default_factory=dict)  # System parameters
-    configs: dict[str, Any] = Field(default_factory=dict)  # Step configurations
-    complex_configs: dict[str, Any] = Field(default_factory=dict)  # Complex step configurations
-    max_retries: int = 0
-    retry_delay: int = 3
-    auto_workflow_cleanup: bool = False
-
-    @field_validator("configs", "params", mode="before")
-    def none_as_empty_dict(cls, value):
-        if value is None:
-            return {}
-        return value
-
-
-class WorkflowSystemComponent(BaseModel):
-    # Container for steps
-    id: str
-    label: str
-    transition_to: list[str] | None = None
-    description: str | None = None
-    ui_config: UIConfig = UIConfig()
-
-
-class WorkflowDefinition(BaseModel):
-    """Workflow definition"""
-
-    name: str
-    description: str | None = None
-    implementation_module: str | None = None
-    steps: list[WorkflowStepDefinition] = Field(default_factory=list)
-    system_components: list[WorkflowSystemComponent] = Field(default_factory=list)
-    configs: list[WorkflowConfigItem] = Field(default_factory=list)
-
-    def get_config_item(self, name: str) -> WorkflowConfigItem | None:
-        for config in self.configs:
-            if config.name == name:
-                return config
-        return None
-
-    def upsert_config_item(self, config_item: WorkflowConfigItem) -> None:
-        for config in self.configs:
-            if config.name == config_item.name:
-                config.value = config_item.value
-                config.label = config_item.label
-                config.type = config_item.type
-                config.required = config_item.required
-                config.options = config_item.options
-                config.group = config_item.group
-                return
-        self.configs.append(config_item)
-
-
-class WorkflowStepEvent(BaseModel):
-    """Workflow step event represent a single step execution"""
-
-    id: str
-    system_component_id: str | None = None
-    state: StepExecutionStatus
-    elapsed_time: float
-    timestamp: str
-    error: str | None = None
-    output_text: str | None = None
-    data: Any = None
-
-
-class WorkflowFullStateReport(BaseModel):
-    """Workflow state report is complete log of workflow execution"""
-
-    workflow_name: str | None = None
-    workflow_version: str | None = None
-    run_id: str | None = None
-    state: WorkflowState
-    start_time: int | None = None
-    end_time: int | None = None
-    elapsed_time: float = 0
-    last_error: str | None = None
-    execution_log: list[WorkflowStepEvent]
-    last_updated_time: int | None = None
-
-    @field_validator("start_time", "end_time", mode="before")
-    def float_to_int(cls, value):
-        if isinstance(value, float):
-            return int(value)
-        return value
-
-
-class WorkflowConfigs(BaseModel):
-    """Workflow configs"""
-
-    configs: list[WorkflowConfigItem] = []
-
-    def get_config_item(self, config_name: str) -> WorkflowConfigItem | None:
-        return next((item for item in self.configs if item.name == config_name), None)
-
-    def set_config_item(self, config_item: WorkflowConfigItem):
-        for item in self.configs:
-            if item.name == config_item.name:
-                item.value = config_item.value
-                return
-        self.configs.append(config_item)
-
-    def get_config_group_values_by_name(
-        self, group_name: str, remove_group_prefix: bool = True
-    ) -> dict[str, str | None]:
-        return {
-            (item.name.removeprefix(item.group) if remove_group_prefix else item.name): item.value
-            for item in self.configs
-            if item.group == group_name
-        }
-
-    def get_config_item_value(self, config_name: str, default_value=None) -> str | None:
-        return config.value if (config := self.get_config_item(config_name)) else default_value

cognite/neat/_workflows/steps/data_contracts.py

@@ -1,109 +0,0 @@
-from pathlib import Path
-
-from cognite.client import CogniteClient
-from cognite.client.data_classes import (
-    Asset,
-    AssetUpdate,
-    Relationship,
-    RelationshipUpdate,
-)
-from cognite.client.data_classes.data_modeling import EdgeApply, NodeApply
-
-from cognite.neat._rules.models import (
-    DMSRules,
-    InformationRules,
-)
-from cognite.neat._store import NeatGraphStore
-from cognite.neat._workflows.steps.step_model import DataContract
-
-
-class MultiRuleData(DataContract):
-    information: InformationRules | None = None
-    dms: DMSRules | None = None
-
-    @classmethod
-    def from_rules(cls, rules: InformationRules | DMSRules):
-        if isinstance(rules, InformationRules):
-            return cls(information=rules)
-        elif isinstance(rules, DMSRules):
-            return cls(dms=rules)
-        else:
-            raise ValueError(f"Unsupported rules type {type(rules)}")
-
-
-class PathData(DataContract):
-    """
-    This represents the path to the excel file.
-
-    Args:
-        excel_file_path: The path to the excel file.
-    """
-
-    excel_file_path: Path
-
-
-class NeatGraph(DataContract):
-    """
-    This represents the neat graph.
-
-    Args:
-        graph: The neat graph store.
-    """
-
-    graph: NeatGraphStore
-
-
-class ClientData(DataContract):
-    """
-    This represents an instantiated CogniteClient object.
-
-    Args:
-        client: The CogniteClient object.
-    """
-
-    client: CogniteClient
-
-
-class CategorizedAssets(DataContract):
-    """ "
-    This represents the categorized assets.
-
-    Args:
-        assets: The categorized assets.
-    """
-
-    assets: dict[str, list[Asset | AssetUpdate]]
-
-
-class CategorizedRelationships(DataContract):
-    """
-    This represents the categorized relationships.
-
-    Args:
-        relationships: The categorized relationships.
-
-    """
-
-    relationships: dict[str, list[Relationship] | list[RelationshipUpdate]]
-
-
-class Nodes(DataContract):
-    """
-    This represents nodes.
-
-    Args:
-        nodes: list of nodes.
-    """
-
-    nodes: list[NodeApply]
-
-
-class Edges(DataContract):
-    """
-    This represents edges.
-
-    Args:
-        edges: list of edges.
-    """
-
-    edges: list[EdgeApply]

cognite/neat/_workflows/steps/lib/current/__init__.py

@@ -1,6 +0,0 @@
-from .rules_importer import *  # noqa
-from .rules_validator import *  # noqa
-from .rules_exporter import *  # noqa
-from .graph_store import *  # noqa
-from .graph_loader import *  # noqa
-from .graph_extractor import *  # noqa