feat(workflows): Add support for simulation integration

CHANGELOG.md

@@ -17,6 +17,10 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [7.70.0] - 2024-11-29
+### Added
+- Workflow support for "simulation" task type. 
+
 ## [7.69.2] - 2024-11-28
 ### Improved
 - Handle conversion of instance lists like NodeList to pandas DataFrame in scenarios where: a) properties are expanded

cognite/client/_version.py

@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.69.2"
+__version__ = "7.70.0"
 __api_subversion__ = "20230101"

cognite/client/data_classes/__init__.py

@@ -291,6 +291,8 @@
     DynamicTaskParameters,
     FunctionTaskOutput,
     FunctionTaskParameters,
+    SimulationTaskOutput,
+    SimulationTaskParameters,
     SubworkflowTaskParameters,
     TransformationTaskOutput,
     TransformationTaskParameters,
@@ -473,6 +475,8 @@
     "DatapointSubscriptionWriteList",
     "OidcCredentials",
     "RawTable",
+    "SimulationTaskParameters",
+    "SimulationTaskOutput",
     "Transformation",
     "TransformationWrite",
     "TransformationBlockedInfo",

cognite/client/data_classes/simulators.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import Self
+
+from cognite.client.data_classes._base import (
+    CogniteObject,
+)
+
+if TYPE_CHECKING:
+    from cognite.client import CogniteClient
+
+
+@dataclass
+class SimulationValueUnitName(CogniteObject):
+    name: str
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            name=resource["name"],
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        return super().dump(camel_case=camel_case)
+
+
+@dataclass
+class SimulationInputOverride(CogniteObject):
+    reference_id: str
+    value: str | int | float
+    unit: SimulationValueUnitName | None = None
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            reference_id=resource["referenceId"],
+            value=resource["value"],
+            unit=SimulationValueUnitName._load(resource["unit"], cognite_client) if "unit" in resource else None,
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        output = super().dump(camel_case=camel_case)
+        if self.unit is not None:
+            output["unit"] = self.unit.dump(camel_case=camel_case)
+
+        return output

cognite/client/data_classes/workflows.py

@@ -20,6 +20,9 @@
     WriteableCogniteResourceList,
 )
 from cognite.client.data_classes.data_modeling.query import Query, ResultSetExpression, Select
+from cognite.client.data_classes.simulators import (
+    SimulationInputOverride,
+)
 from cognite.client.utils._text import convert_all_keys_to_camel_case, to_snake_case
 
 if TYPE_CHECKING:
@@ -131,7 +134,7 @@ def as_write(self) -> WorkflowUpsertList:
         return WorkflowUpsertList([workflow.as_write() for workflow in self.data])
 
 
-ValidTaskType = Literal["function", "transformation", "cdf", "dynamic", "subworkflow"]
+ValidTaskType = Literal["function", "transformation", "cdf", "dynamic", "subworkflow", "simulation"]
 
 
 class WorkflowTaskParameters(CogniteObject, ABC):
@@ -159,6 +162,8 @@ def load_parameters(cls, data: dict) -> WorkflowTaskParameters:
             return SubworkflowTaskParameters._load(parameters)
         elif type_ == "subworkflow" and "workflowExternalId" in parameters["subworkflow"]:
             return SubworkflowReferenceParameters._load(parameters)
+        elif type_ == "simulation":
+            return SimulationTaskParameters._load(parameters)
         else:
             raise ValueError(f"Unknown task type: {type_}. Expected {ValidTaskType}")
 
@@ -238,6 +243,49 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
         return output
 
 
+class SimulationTaskParameters(WorkflowTaskParameters):
+    """
+    The simulation parameters are used to specify the simulation routine to be executed.
+    Args:
+        routine_external_id (str): The external ID of the simulation routine to be executed.
+        run_time (int | None): Reference timestamp used for data pre-processing and data sampling.
+        inputs (list[SimulationInputOverride] | None): List of input overrides.
+    """
+
+    task_type = "simulation"
+
+    def __init__(
+        self,
+        routine_external_id: str,
+        run_time: int | None = None,
+        inputs: list[SimulationInputOverride] | None = None,
+    ) -> None:
+        self.routine_external_id = routine_external_id
+        self.run_time = run_time
+        self.inputs = inputs
+
+    @classmethod
+    def _load(cls, resource: dict, cognite_client: CogniteClient | None = None) -> SimulationTaskParameters:
+        simulation: dict[str, Any] = resource["simulation"]
+
+        return cls(
+            routine_external_id=simulation["routineExternalId"],
+            run_time=simulation.get("runTime"),
+            inputs=[SimulationInputOverride._load(item) for item in simulation.get("inputs", [])]
+            if simulation.get("inputs")
+            else None,
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        simulation = {
+            "routineExternalId" if camel_case else "routine_external_id": self.routine_external_id,
+            "runTime" if camel_case else "run_time": self.run_time,
+            "inputs": [item.dump(camel_case) for item in self.inputs] if self.inputs else None,
+        }
+
+        return {"simulation": simulation}
+
+
 class TransformationTaskParameters(WorkflowTaskParameters):
     """
     The transformation parameters are used to specify the transformation to be called.
@@ -540,6 +588,8 @@ def load_output(cls, data: dict) -> WorkflowTaskOutput:
             return DynamicTaskOutput.load(data)
         elif task_type == "subworkflow":
             return SubworkflowTaskOutput.load(data)
+        elif task_type == "simulation":
+            return SimulationTaskOutput.load(data)
         else:
             raise ValueError(f"Unknown task type: {task_type}")
 
@@ -579,6 +629,44 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
         }
 
 
+class SimulationTaskOutput(WorkflowTaskOutput):
+    """
+    The class represent the output of Simulation execution.
+    Args:
+        run_id (int | None): The run ID of the simulation run.
+        log_id (int | None): The log ID of the simulation run.
+        status_message (str | None): Status message of the simulation execution.
+    """
+
+    task_type: ClassVar[str] = "simulation"
+
+    def __init__(
+        self,
+        run_id: int | None,
+        log_id: int | None,
+        status_message: str | None,
+    ) -> None:
+        self.run_id = run_id
+        self.log_id = log_id
+        self.status_message = status_message
+
+    @classmethod
+    def load(cls, data: dict[str, Any]) -> SimulationTaskOutput:
+        output = data["output"]
+        return cls(
+            run_id=output.get("runId"),
+            log_id=output.get("logId"),
+            status_message=output.get("statusMessage"),
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        return {
+            "runId" if camel_case else "run_id": self.run_id,
+            "logId" if camel_case else "log_id": self.log_id,
+            "statusMessage" if camel_case else "status_message": self.status_message,
+        }
+
+
 class TransformationTaskOutput(WorkflowTaskOutput):
     """
     The transformation output is used to specify the output of a transformation task.

pyproject.toml

@@ -1,7 +1,8 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "7.69.2"
+version = "7.70.0"
+
 description = "Cognite Python SDK"
 readme = "README.md"
 documentation = "https://cognite-sdk-python.readthedocs-hosted.com"

tests/tests_integration/test_api/test_workflows.py

@@ -13,6 +13,7 @@
 from cognite.client.data_classes.workflows import (
     CDFTaskParameters,
     FunctionTaskParameters,
+    SimulationTaskParameters,
     SubworkflowReferenceParameters,
     SubworkflowTaskParameters,
     TransformationTaskParameters,
@@ -326,6 +327,43 @@ def test_list_workflows(self, cognite_client: CogniteClient, workflow_list: Work
 
 
 class TestWorkflowVersions:
+    def test_upsert_delete_with_simulation_task(
+        self,
+        cognite_client: CogniteClient,
+    ):
+        workflow_id = "integration_test-workflow_for_simulator_integration" + random_string(5)
+
+        version = WorkflowVersionUpsert(
+            workflow_external_id=workflow_id,
+            version="1",
+            workflow_definition=WorkflowDefinitionUpsert(
+                tasks=[
+                    WorkflowTask(
+                        external_id=f"{workflow_id}-1-task1" + random_string(5),
+                        parameters=SimulationTaskParameters(
+                            routine_external_id="integration_tests_workflow",
+                        ),
+                    )
+                ],
+                description=None,
+            ),
+        )
+
+        cognite_client.workflows.versions.delete(version.as_id(), ignore_unknown_ids=True)
+        created_version: WorkflowVersion | None = None
+
+        try:
+            created_version = cognite_client.workflows.versions.upsert(version)
+            assert created_version.workflow_external_id == workflow_id
+            assert created_version.workflow_definition.tasks[0].type == "simulation"
+            assert len(created_version.workflow_definition.tasks) > 0
+        finally:
+            if created_version is not None:
+                cognite_client.workflows.versions.delete(
+                    created_version.as_id(),
+                )
+                cognite_client.workflows.delete(created_version.workflow_external_id)
+
     def test_upsert_preexisting(self, cognite_client: CogniteClient, new_workflow_version: WorkflowVersion) -> None:
         new_workflow_version.workflow_definition.description = "Updated description for testing purposes"
         updated_version = cognite_client.workflows.versions.upsert(new_workflow_version.as_write())