cognite-toolkit 0.7.31__py3-none-any.whl → 0.7.33__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/data_classes/instance_api.py

@@ -7,7 +7,7 @@ from .base import BaseModelObject, Identifier, RequestResource
 InstanceType: TypeAlias = Literal["node", "edge"]
 
 
-class InstanceIdentifier(Identifier):
+class TypedInstanceIdentifier(Identifier):
     """Identifier for an Instance instance."""
 
     instance_type: InstanceType
@@ -15,11 +15,16 @@ class InstanceIdentifier(Identifier):
     external_id: str
 
 
-class NodeIdentifier(InstanceIdentifier):
+class InstanceIdentifier(Identifier):
+    space: str
+    external_id: str
+
+
+class NodeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["node"] = "node"
 
 
-class EdgeIdentifier(InstanceIdentifier):
+class EdgeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["edge"] = "edge"
 
 
@@ -32,8 +37,8 @@ class InstanceResult(BaseModelObject):
     created_time: int
     last_updated_time: int
 
-    def as_id(self) -> InstanceIdentifier:
-        return InstanceIdentifier(
+    def as_id(self) -> TypedInstanceIdentifier:
+        return TypedInstanceIdentifier(
             instance_type=self.instance_type,
             space=self.space,
             external_id=self.external_id,
@@ -62,8 +67,8 @@ class InstanceRequestResource(RequestResource):
     space: str
     external_id: str
 
-    def as_id(self) -> InstanceIdentifier:
-        return InstanceIdentifier(
+    def as_id(self) -> TypedInstanceIdentifier:
+        return TypedInstanceIdentifier(
             instance_type=self.instance_type,
             space=self.space,
             external_id=self.external_id,
@@ -107,7 +112,7 @@ class InstanceSource(BaseModelObject):
         return value
 
 
-class InstanceRequestItem(BaseModelObject):
+class InstanceRequestItem(RequestResource):
     model_config = ConfigDict(populate_by_name=True)
     instance_type: InstanceType
     space: str
@@ -115,8 +120,8 @@ class InstanceRequestItem(BaseModelObject):
     existing_version: int | None = None
     sources: list[InstanceSource] | None = None
 
-    def as_id(self) -> InstanceIdentifier:
-        return InstanceIdentifier(
+    def as_id(self) -> TypedInstanceIdentifier:
+        return TypedInstanceIdentifier(
             instance_type=self.instance_type,
             space=self.space,
             external_id=self.external_id,
@@ -128,7 +133,7 @@ class InstanceResponseItem(BaseModelObject):
     space: str
     external_id: str
     version: int
-    type: InstanceIdentifier | None = None
+    type: TypedInstanceIdentifier | None = None
     created_time: int
     last_updated_time: int
     deleted_time: int | None = None
@@ -149,8 +154,8 @@ class InstanceResponseItem(BaseModelObject):
         output.update(space_properties.get(view_version, {}))
         return output
 
-    def as_id(self) -> InstanceIdentifier:
-        return InstanceIdentifier(
+    def as_id(self) -> TypedInstanceIdentifier:
+        return TypedInstanceIdentifier(
             instance_type=self.instance_type,
             space=self.space,
             external_id=self.external_id,

cognite_toolkit/_cdf_tk/client/data_classes/three_d.py

@@ -19,11 +19,17 @@ class RevisionStatus(BaseModelObject):
 class ThreeDModelRequest(RequestResource):
     name: str
 
+    def as_id(self) -> str:
+        return self.name
+
 
 class ThreeDModelClassicRequest(ThreeDModelRequest):
     data_set_id: int | None = None
     metadata: dict[str, str] | None = None
 
+    def as_id(self) -> str:
+        return self.name
+
 
 class ThreeDModelDMSRequest(ThreeDModelRequest):
     space: str

cognite_toolkit/_cdf_tk/client/testing.py

@@ -13,6 +13,7 @@ from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client._toolkit_client import ToolkitClient
 
+from ._toolkit_client import ToolAPI
 from .api.canvas import CanvasAPI, IndustrialCanvasAPI
 from .api.charts import ChartsAPI
 from .api.dml import DMLAPI
@@ -53,6 +54,7 @@ from .api.robotics.maps import MapsAPI
 from .api.search import SearchAPI
 from .api.search_config import SearchConfigurationsAPI
 from .api.streams import StreamsAPI
+from .api.three_d import ThreeDAPI, ThreeDModelAPI
 from .api.token import TokenAPI
 from .api.verify import VerifyAPI
 
@@ -131,6 +133,10 @@ class ToolkitClientMock(CogniteClientMock):
         self.time_series.data.synthetic = MagicMock(spec_set=SyntheticDatapointsAPI)
         self.time_series.subscriptions = MagicMock(spec_set=DatapointsSubscriptionAPI)
 
+        self.tool = MagicMock(spec=ToolAPI)
+        self.tool.three_d = MagicMock(spec=ThreeDAPI)
+        self.tool.three_d.models = MagicMock(spec_set=ThreeDModelAPI)
+
         self.streams = MagicMock(spec=StreamsAPI)
 
         # This is a helper API, not a real API.

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -9,8 +9,10 @@ from cognite.client.data_classes._base import (
 from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, NodeId, ViewId
 from cognite.client.utils._identifier import InstanceId
 from cognite.client.utils._text import to_camel_case
-from pydantic import BaseModel, BeforeValidator, field_validator, model_validator
+from pydantic import BaseModel, BeforeValidator, Field, field_validator, model_validator
 
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestResource
+from cognite_toolkit._cdf_tk.client.data_classes.instance_api import InstanceIdentifier
 from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceApplyList
 from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId
 from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
@@ -261,3 +263,32 @@ class AssetCentricMappingList(
 
     def as_write(self) -> InstanceApplyList:
         return InstanceApplyList([item.as_write() for item in self])
+
+
+class Model(BaseModelObject):
+    instance_id: InstanceIdentifier
+
+
+class Thumbnail(BaseModelObject):
+    instance_id: InstanceIdentifier
+
+
+class ThreeDRevisionMigrationRequest(RequestResource):
+    space: str
+    type: Literal["CAD", "PointCloud", "Image360"]
+    revision_id: int
+    model: Model
+
+    def as_id(self) -> int:
+        return self.revision_id
+
+
+class ThreeDMigrationRequest(RequestResource):
+    model_id: int
+    type: Literal["CAD", "PointCloud", "Image360"]
+    space: str
+    thumbnail: Thumbnail | None = None
+    revision: ThreeDRevisionMigrationRequest = Field(exclude=True)
+
+    def as_id(self) -> int:
+        return self.model_id

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Callable, Sequence
-from typing import Generic, cast
+from typing import Generic, Literal, cast
 from uuid import uuid4
 
 from cognite.client.data_classes.data_modeling import (
@@ -27,26 +27,33 @@ from cognite_toolkit._cdf_tk.client.data_classes.charts_data import (
     ChartSource,
     ChartTimeseries,
 )
+from cognite_toolkit._cdf_tk.client.data_classes.instance_api import InstanceIdentifier
 from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMappingApply
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import RevisionStatus, ThreeDModelResponse
 from cognite_toolkit._cdf_tk.commands._migrate.conversion import DirectRelationCache, asset_centric_to_dm
-from cognite_toolkit._cdf_tk.commands._migrate.data_classes import AssetCentricMapping
+from cognite_toolkit._cdf_tk.commands._migrate.data_classes import (
+    Model,
+    ThreeDMigrationRequest,
+    ThreeDRevisionMigrationRequest,
+)
 from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_default_mappings
 from cognite_toolkit._cdf_tk.commands._migrate.issues import (
     CanvasMigrationIssue,
     ChartMigrationIssue,
     ConversionIssue,
     MigrationIssue,
+    ThreeDModelMigrationIssue,
 )
-from cognite_toolkit._cdf_tk.commands._migrate.selectors import AssetCentricMigrationSelector
 from cognite_toolkit._cdf_tk.constants import MISSING_INSTANCE_SPACE
 from cognite_toolkit._cdf_tk.exceptions import ToolkitMigrationError, ToolkitValueError
 from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
 from cognite_toolkit._cdf_tk.storageio._base import T_Selector
-from cognite_toolkit._cdf_tk.storageio.selectors import CanvasSelector, ChartSelector
+from cognite_toolkit._cdf_tk.storageio.selectors import CanvasSelector, ChartSelector, ThreeDSelector
 from cognite_toolkit._cdf_tk.utils import humanize_collection
-from cognite_toolkit._cdf_tk.utils.useful_types import (
-    T_AssetCentricResourceExtended,
-)
+from cognite_toolkit._cdf_tk.utils.useful_types import T_AssetCentricResourceExtended
+
+from .data_classes import AssetCentricMapping
+from .selectors import AssetCentricMigrationSelector
 
 
 class DataMapper(Generic[T_Selector, T_ResourceResponse, T_ResourceRequest], ABC):
@@ -383,3 +390,82 @@ class CanvasMapper(DataMapper[CanvasSelector, IndustrialCanvas, IndustrialCanvas
             max_width=reference.max_width,
             max_height=reference.max_height,
         )
+
+
+class ThreeDMapper(DataMapper[ThreeDSelector, ThreeDModelResponse, ThreeDMigrationRequest]):
+    def __init__(self, client: ToolkitClient) -> None:
+        self.client = client
+
+    def map(
+        self, source: Sequence[ThreeDModelResponse]
+    ) -> Sequence[tuple[ThreeDMigrationRequest | None, MigrationIssue]]:
+        self._populate_cache(source)
+        output: list[tuple[ThreeDMigrationRequest | None, MigrationIssue]] = []
+        for item in source:
+            mapped_item, issue = self._map_single_item(item)
+            output.append((mapped_item, issue))
+        return output
+
+    def _populate_cache(self, source: Sequence[ThreeDModelResponse]) -> None:
+        dataset_ids: set[int] = set()
+        for model in source:
+            if model.data_set_id is not None:
+                dataset_ids.add(model.data_set_id)
+        self.client.migration.space_source.retrieve(list(dataset_ids))
+
+    def _map_single_item(
+        self, item: ThreeDModelResponse
+    ) -> tuple[ThreeDMigrationRequest | None, ThreeDModelMigrationIssue]:
+        issue = ThreeDModelMigrationIssue(model_name=item.name, model_id=item.id)
+        instance_space: str | None = None
+        last_revision_id: int | None = None
+        model_type: Literal["CAD", "PointCloud", "Image360"] | None = None
+        if item.data_set_id is None:
+            issue.error_message.append("3D model is not associated with any dataset.")
+        else:
+            space_source = self.client.migration.space_source.retrieve(item.data_set_id)
+            if space_source is not None:
+                instance_space = space_source.instance_space
+        if instance_space is None and item.data_set_id is not None:
+            issue.error_message.append(f"Missing instance space for dataset ID {item.data_set_id!r}")
+        if item.last_revision_info is None:
+            issue.error_message.append("3D model has no revisions.")
+        else:
+            model_type = self._get_type(item.last_revision_info)
+            last_revision_id = item.last_revision_info.revision_id
+            if last_revision_id is None:
+                issue.error_message.append("3D model's last revision has no revision ID.")
+
+        if model_type is None:
+            issue.error_message.append("3D model's last revision has no recognized type.")
+
+        if instance_space is None or last_revision_id is None or model_type is None or issue.has_issues:
+            return None, issue
+
+        mapped_request = ThreeDMigrationRequest(
+            model_id=item.id,
+            type=model_type,
+            space=instance_space,
+            revision=ThreeDRevisionMigrationRequest(
+                space=instance_space,
+                type=model_type,
+                revision_id=last_revision_id,
+                model=Model(
+                    instance_id=InstanceIdentifier(
+                        space=instance_space,
+                        external_id=f"cog_3d_model_{item.id!s}",
+                    )
+                ),
+            ),
+        )
+        return mapped_request, issue
+
+    @staticmethod
+    def _get_type(revision: RevisionStatus) -> Literal["CAD", "PointCloud", "Image360"] | None:
+        types = revision.types or []
+        if any("gltf-directory" in t for t in types):
+            return "CAD"
+        elif any("ept-pointcloud" in t for t in types):
+            return "PointCloud"
+        else:
+            return None

cognite_toolkit/_cdf_tk/commands/_migrate/issues.py

@@ -30,6 +30,24 @@ class MigrationIssue(MigrationObject):
         return True
 
 
+class ThreeDModelMigrationIssue(MigrationIssue):
+    """Represents a 3D model migration issue encountered during migration.
+
+    Attributes:
+        model_external_id (str): The external ID of the 3D model that could not be migrated.
+    """
+
+    type: ClassVar[str] = "threeDModelMigration"
+    model_name: str
+    model_id: int
+    error_message: list[str] = Field(default_factory=list)
+
+    @property
+    def has_issues(self) -> bool:
+        """Check if there are any issues recorded in this ThreeDModelMigrationIssue."""
+        return bool(self.error_message)
+
+
 class ChartMigrationIssue(MigrationIssue):
     """Represents a chart migration issue encountered during migration.
 

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -6,6 +6,8 @@ from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, Nod
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import ThreeDModelResponse
+from cognite_toolkit._cdf_tk.commands._migrate.data_classes import ThreeDMigrationRequest
 from cognite_toolkit._cdf_tk.constants import MISSING_EXTERNAL_ID, MISSING_INSTANCE_SPACE
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError, ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio import (
@@ -15,9 +17,22 @@ from cognite_toolkit._cdf_tk.storageio import (
     UploadableStorageIO,
 )
 from cognite_toolkit._cdf_tk.storageio._base import Page, UploadItem
+from cognite_toolkit._cdf_tk.storageio.selectors import (
+    ThreeDModelFilteredSelector,
+    ThreeDModelIdSelector,
+    ThreeDSelector,
+)
 from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest, SuccessResponseItems
+from cognite_toolkit._cdf_tk.utils.http_client import (
+    FailedResponse,
+    HTTPClient,
+    HTTPMessage,
+    ItemsRequest,
+    SimpleBodyRequest,
+    SuccessResponseItems,
+    ToolkitAPIError,
+)
 from cognite_toolkit._cdf_tk.utils.useful_types import (
     AssetCentricKindExtended,
     AssetCentricType,
@@ -348,3 +363,94 @@ class AnnotationMigrationIO(
         selector: AssetCentricMigrationSelector | None = None,
     ) -> list[dict[str, JsonVal]]:
         raise NotImplementedError("Serializing Annotation Migrations to JSON is not supported.")
+
+
+class ThreeDMigrationIO(UploadableStorageIO[ThreeDSelector, ThreeDModelResponse, ThreeDMigrationRequest]):
+    KIND = "3DMigration"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".ndjson"})
+    DOWNLOAD_LIMIT = 1000
+    CHUNK_SIZE = 1
+    UPLOAD_ENDPOINT = "/3d/migrate/models"
+    REVISION_ENDPOINT = "/3d/migrate/revisions"
+
+    def as_id(self, item: ThreeDModelResponse) -> str:
+        return f"{item.name}_{item.id!s}"
+
+    def stream_data(self, selector: ThreeDSelector, limit: int | None = None) -> Iterable[Page[ThreeDModelResponse]]:
+        published: bool | None = None
+        if isinstance(selector, ThreeDModelFilteredSelector):
+            published = selector.published
+        included_models: set[int] | None = None
+        if isinstance(selector, ThreeDModelIdSelector):
+            included_models = set(selector.ids)
+        cursor: str | None = None
+        total = 0
+        while True:
+            request_limit = min(self.DOWNLOAD_LIMIT, limit - total) if limit is not None else self.DOWNLOAD_LIMIT
+            response = self.client.tool.three_d.models.iterate(
+                published=published, include_revision_info=True, limit=request_limit, cursor=cursor
+            )
+            # Only include asset-centric 3D models
+            items = [
+                item
+                for item in response.items
+                if item.space is None and (included_models is None or item.id in included_models)
+            ]
+            total += len(items)
+            if items:
+                yield Page(worker_id="main", items=items, next_cursor=response.next_cursor)
+            if response.next_cursor is None:
+                break
+            cursor = response.next_cursor
+
+    def count(self, selector: ThreeDSelector) -> int | None:
+        # There is no efficient way to count 3D models in CDF.
+        return None
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[ThreeDModelResponse], selector: ThreeDSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        raise NotImplementedError("Deserializing Annotation Migrations from JSON is not supported.")
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> ThreeDMigrationRequest:
+        raise NotImplementedError("Deserializing ThreeD Migrations from JSON is not supported.")
+
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[ThreeDMigrationRequest]],
+        http_client: HTTPClient,
+        selector: ThreeDSelector | None = None,
+    ) -> Sequence[HTTPMessage]:
+        """Migrate 3D models by uploading them to the migrate/models endpoint."""
+        if len(data_chunk) > self.CHUNK_SIZE:
+            raise RuntimeError(f"Uploading more than {self.CHUNK_SIZE} 3D models at a time is not supported.")
+
+        results: list[HTTPMessage] = []
+        responses = http_client.request_with_retries(
+            message=ItemsRequest(
+                endpoint_url=self.client.config.create_api_url(self.UPLOAD_ENDPOINT),
+                method="POST",
+                items=list(data_chunk),
+            )
+        )
+        if (
+            failed_response := next((res for res in responses if isinstance(res, FailedResponse)), None)
+        ) and failed_response.status_code == 400:
+            raise ToolkitAPIError("3D model migration failed. You need to enable the 3D migration alpha feature flag.")
+
+        results.extend(responses)
+        success_ids = {id for res in responses if isinstance(res, SuccessResponseItems) for id in res.ids}
+        for data in data_chunk:
+            if data.source_id not in success_ids:
+                continue
+            revision = http_client.request_with_retries(
+                message=SimpleBodyRequest(
+                    endpoint_url=self.client.config.create_api_url(self.REVISION_ENDPOINT),
+                    method="POST",
+                    body_content={"items": [data.item.revision.dump(camel_case=True)]},
+                )
+            )
+            results.extend(revision.as_item_responses(data.source_id))
+        return results

cognite_toolkit/_cdf_tk/resource_classes/workflow_version.py

@@ -1,4 +1,5 @@
-from typing import Literal
+from abc import ABC
+from typing import Annotated, Literal
 
 from cognite.client.data_classes import WorkflowVersionUpsert
 from pydantic import Field, JsonValue
@@ -6,18 +7,141 @@ from pydantic import Field, JsonValue
 from .base import BaseModelResource, ToolkitResource
 
 
+class WorkflowVersionId(BaseModelResource):
+    workflow_external_id: str = Field(
+        max_length=255,
+        description="Identifier for a workflow. Must be unique for the project. No trailing or leading whitespace and no null characters allowed.",
+    )
+    version: str = Field(
+        max_length=255,
+        description="Identifier for a version. Must be unique for the workflow. No trailing or leading whitespace and no null characters allowed.",
+    )
+
+
+class CogniteFunctionRef(BaseModelResource):
+    external_id: str = Field(
+        description="The external id of the Cognite Function in the project. This can be either a function external ID or a reference like ${myTaskExternalId.output.someKey}"
+    )
+    data: str | JsonValue | None = Field(
+        None, description="Input data that will be passed to the Cognite Function. Limited to 100KB in size."
+    )
+
+
+class FunctionTaskParameters(BaseModelResource):
+    function: CogniteFunctionRef
+    is_async_complete: bool = Field(
+        False, description="Defines if the execution of the task should be completed asynchronously."
+    )
+
+
+class TransformationRef(BaseModelResource):
+    external_id: str = Field(
+        description="The external id of the Transformation in the project. This can be either a transformation external ID or a reference like ${myTaskExternalId.output.someKey}"
+    )
+    concurrency_policy: Literal["fail", "waitForCurrent", "restartAfterCurrent"] = Field(
+        "fail",
+        description="""Determines the behavior of the task if the Transformation is already running.
+
+fail: The task fails if another instance of the Transformation is currently running.
+waitForCurrent: The task will pause and wait for the already running Transformation to complete. Once completed, the task is completed. This mode is useful for preventing redundant Transformation runs.
+restartAfterCurrent: The task waits for the ongoing Transformation to finish. After completion, the task restarts the Transformation. This mode ensures that the most recent data can be used by following tasks.""",
+    )
+    use_transformation_credentials: bool = Field(
+        False,
+        description="If set to true, the transformation will run using the client credentials configured on the transformation. If set to false, the transformation will run using the client credentials used to trigger the workflow.",
+    )
+
+
+class TransformationTaskParameters(BaseModelResource):
+    transformation: TransformationRef
+
+
+class CDFRequest(BaseModelResource):
+    resource_path: str = Field(
+        description="The path of the request. The path should be prefixed by {cluster}.cognitedata.com/api/v1/project/{project} based on the relevant cluster and project. It can also contain references like ${myTaskExternalId.output.someKey}"
+    )
+    method: Literal["POST", "GET", "PUT"] | str = Field(
+        description="The HTTP method of the request. It can also be a reference like ${myTaskExternalId.output.someKey}"
+    )
+    query_parameters: dict[str, JsonValue] | str | None = Field(
+        None,
+        description="The query parameters of the request. It can also be a reference like ${myTaskExternalId.output.someKey}",
+    )
+    body: JsonValue | str | None = Field(
+        None, description="The body of the request. It can also be a reference like ${myTaskExternalId.output.someKey}"
+    )
+    request_timeout_in_millis: float | str | None = Field(
+        None,
+        description="The timeout for the request in milliseconds. It can also be a reference like ${myTaskExternalId.output.someKey}",
+    )
+    cdf_version_header: Literal["alpha", "beta"] | str | None = Field(
+        None, description="The Cognite Data Fusion version header to use for the request."
+    )
+
+
+class CDFTaskParameters(BaseModelResource):
+    cdf_request: CDFRequest
+
+
+class DynamicRef(BaseModelResource):
+    tasks: str = Field(
+        description="A Reference is an expression that allows dynamically injecting input to a task during execution. References can be used to reference the input of the Workflow, the output of a previous task in the Workflow, or the input of a previous task in the Workflow. Note that the injected value must be valid in the context of the property it is injected into. Example Task reference: ${myTaskExternalId.output.someKey} Example Workflow input reference: ${workflow.input.myKey}"
+    )
+
+
+class DynamicTaskParameters(BaseModelResource):
+    dynamic: DynamicRef = Field(description="Reference to another task to use as the definition for this task.")
+
+
+class SubworkflowRef(BaseModelResource):
+    tasks: "WorkflowVersionId | list[Task]" = Field(
+        description="Reference to the subworkflow to execute. This can be either a reference to an existing workflow version or an inline definition of tasks."
+    )
+
+
+class SubworkflowTaskParameters(BaseModelResource):
+    subworkflow: SubworkflowRef = Field(description="Reference to the subworkflow to execute.")
+
+
+class SimulatorInputUnit(BaseModelResource):
+    name: str = Field(description="Name of the unit.")
+
+
+class SimulatorInput(BaseModelResource):
+    reference_id: str = Field(description="Reference id of the value to override.")
+    value: str | int | float | list[str] | list[int] | list[float] = Field(
+        description="Override the value used for a simulation run."
+    )
+    unit: SimulatorInputUnit | None = Field(None, description="Override the unit of the value")
+
+
+class SimulationRef(BaseModelResource):
+    routine_external_id: str = Field(description="The external id of the routine to be executed.")
+    run_time: int | None = Field(
+        None,
+        description="Run time in milliseconds. Reference timestamp used for data pre-processing and data sampling.",
+    )
+    inputs: list[SimulatorInput] | None = Field(
+        None, description="List of inputs to be provided to the simulation.", max_length=200
+    )
+
+
+class SimulationTaskParameters(BaseModelResource):
+    simulation: SimulationRef = Field(description="Reference to the simulation to execute.")
+
+
 class TaskId(BaseModelResource):
     external_id: str = Field(
         max_length=255, description="The external ID provided by the client. Must be unique for the resource type."
     )
 
 
-class TaskDefinition(BaseModelResource):
+class TaskDefinition(BaseModelResource, ABC):
     external_id: str = Field(
         max_length=255,
         description="Identifier for the task. Must be unique within the version. No trailing or leading whitespace and no null characters allowed.",
     )
-    type: Literal["function", "transformation", "cdf", "dynamic", "subworkflow", "simulation"]
+    type: str
     name: str | None = Field(
         default=None,
         max_length=255,
@@ -28,7 +152,6 @@ class TaskDefinition(BaseModelResource):
         max_length=500,
         description="Description of the intention of the task",
     )
-    parameters: JsonValue = Field()
     retries: int = Field(
         3,
         ge=0,
@@ -52,13 +175,49 @@ class TaskDefinition(BaseModelResource):
     )
 
 
+class FunctionTask(TaskDefinition):
+    type: Literal["function"] = "function"
+    parameters: FunctionTaskParameters
+
+
+class TransformationTask(TaskDefinition):
+    type: Literal["transformation"] = "transformation"
+    parameters: TransformationTaskParameters
+
+
+class CDFTask(TaskDefinition):
+    type: Literal["cdfRequest"] = "cdfRequest"
+    parameters: CDFTaskParameters
+
+
+class DynamicTask(TaskDefinition):
+    type: Literal["dynamic"] = "dynamic"
+    parameters: DynamicTaskParameters
+
+
+class SubworkflowTask(TaskDefinition):
+    type: Literal["subworkflow"] = "subworkflow"
+    parameters: SubworkflowTaskParameters
+
+
+class SimulationTask(TaskDefinition):
+    type: Literal["simulation"] = "simulation"
+    parameters: SimulationTaskParameters
+
+
+Task = Annotated[
+    FunctionTask | TransformationTask | CDFTask | DynamicTask | SubworkflowTask | SimulationTask,
+    Field(discriminator="type"),
+]
+
+
 class WorkflowDefinition(BaseModelResource):
     description: str | None = Field(
         default=None,
         max_length=500,
         description="The description of the workflow version.",
     )
-    tasks: list[TaskDefinition]
+    tasks: list[Task]
 
 
 class WorkflowVersionYAML(ToolkitResource):