cognite-toolkit 0.7.42__py3-none-any.whl → 0.7.44__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/cdf_client/api.py

@@ -5,6 +5,7 @@ that handle CRUD operations for CDF Data Modeling API resources.
 """
 
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from functools import partial
@@ -97,9 +98,10 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
         method: APIMethod,
         params: dict[str, Any] | None = None,
         extra_body: dict[str, Any] | None = None,
+        endpoint: str | None = None,
     ) -> list[T_ResponseResource]:
         response_items: list[T_ResponseResource] = []
-        for response in self._chunk_requests(items, method, self._serialize_items, params, extra_body):
+        for response in self._chunk_requests(items, method, self._serialize_items, params, extra_body, endpoint):
             response_items.extend(self._page_response(response).items)
         return response_items
 
@@ -121,8 +123,9 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
         method: APIMethod,
         params: dict[str, Any] | None = None,
         extra_body: dict[str, Any] | None = None,
+        endpoint: str | None = None,
     ) -> None:
-        list(self._chunk_requests(items, method, self._serialize_items, params, extra_body))
+        list(self._chunk_requests(items, method, self._serialize_items, params, extra_body, endpoint))
         return None
 
     def _chunk_requests(
@@ -132,14 +135,28 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
         serialization: Callable[[Sequence[_T_BaseModel]], list[dict[str, JsonValue]]],
         params: dict[str, Any] | None = None,
         extra_body: dict[str, Any] | None = None,
+        endpoint_path: str | None = None,
     ) -> Iterable[SuccessResponse2]:
+        """Send requests in chunks and yield responses.
+
+        Args:
+            items: The items to process.
+            method: The API method to use. This is used ot look the up the endpoint.
+            serialization: A function to serialize the items to JSON-compatible dicts.
+            params: Optional query parameters for the request.
+            extra_body: Optional extra body content to include in the request.
+            endpoint_path: Optional override for the endpoint path.
+
+        Yields:
+            The successful responses from the API.
+        """
         # Filter out None
         request_params = self._filter_out_none_values(params)
         endpoint = self._method_endpoint_map[method]
 
         for chunk in chunker_sequence(items, endpoint.item_limit):
             request = RequestMessage2(
-                endpoint_url=f"{self._make_url(endpoint.path)}",
+                endpoint_url=f"{self._make_url(endpoint_path or endpoint.path)}",
                 method=endpoint.method,
                 body_content={"items": serialization(chunk), **(extra_body or {})},  # type: ignore[dict-item]
                 parameters=request_params,
@@ -154,12 +171,24 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
             request_params = {k: v for k, v in params.items() if v is not None}
         return request_params
 
+    @classmethod
+    def _group_items_by_text_field(
+        cls, items: Sequence[_T_BaseModel], field_name: str
+    ) -> dict[str, list[_T_BaseModel]]:
+        """Group items by a text field."""
+        grouped_items: dict[str, list[_T_BaseModel]] = defaultdict(list)
+        for item in items:
+            key = str(getattr(item, field_name))
+            grouped_items[key].append(item)
+        return grouped_items
+
     def _iterate(
         self,
         limit: int,
         cursor: str | None = None,
         params: dict[str, Any] | None = None,
         body: dict[str, Any] | None = None,
+        endpoint_path: str | None = None,
     ) -> PagedResponse[T_ResponseResource]:
         """Fetch a single page of resources.
 
@@ -171,8 +200,11 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
                 - space: Filter by space
                 - includeGlobal: Whether to include global resources
             body : Body content for the request, if applicable.
+                limit: Maximum number of items to return in the page.
+                cursor: Cursor for pagination.
             limit: Maximum number of items to return in the page.
             cursor: Cursor for pagination.
+
         Returns:
             A Page containing the items and the cursor for the next page.
         """
@@ -194,7 +226,7 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
             raise NotImplementedError(f"Unsupported method {endpoint.method} for pagination.")
 
         request = RequestMessage2(
-            endpoint_url=self._make_url(endpoint.path),
+            endpoint_url=self._make_url(endpoint_path or endpoint.path),
             method=endpoint.method,
             parameters=request_params,
             body_content=body,
@@ -203,7 +235,9 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
         response = result.get_success_or_raise()
         return self._page_response(response)
 
-    def _list(self, limit: int | None = None, params: dict[str, Any] | None = None) -> list[T_ResponseResource]:
+    def _list(
+        self, limit: int | None = None, params: dict[str, Any] | None = None, endpoint_path: str | None = None
+    ) -> list[T_ResponseResource]:
         """List all resources, handling pagination automatically."""
         all_items: list[T_ResponseResource] = []
         next_cursor: str | None = None
@@ -211,7 +245,7 @@ class CDFResourceAPI(Generic[T_Identifier, T_RequestResource, T_ResponseResource
         endpoint = self._method_endpoint_map["list"]
         while True:
             page_limit = endpoint.item_limit if limit is None else min(limit - total, endpoint.item_limit)
-            page = self._iterate(limit=page_limit, cursor=next_cursor, params=params)
+            page = self._iterate(limit=page_limit, cursor=next_cursor, params=params, endpoint_path=endpoint_path)
             all_items.extend(page.items)
             total += len(page.items)
             if page.next_cursor is None or (limit is not None and total >= limit):

cognite_toolkit/_cdf_tk/client/data_classes/agent.py

@@ -98,29 +98,26 @@ AgentTool = Annotated[
 ]
 
 
-class AgentRequest(RequestResource):
+class Agent(BaseModelObject):
     external_id: str
     name: str
     description: str | None = None
     instructions: str | None = None
     model: str = "azure/gpt-4o-mini"
     tools: list[AgentTool] | None = None
-    runtime_version: str | None = None
 
     def as_id(self) -> ExternalId:
         return ExternalId(external_id=self.external_id)
 
 
-class AgentResponse(ResponseResource[AgentRequest]):
+class AgentRequest(Agent, RequestResource):
+    runtime_version: str | None = None
+
+
+class AgentResponse(Agent, ResponseResource[AgentRequest]):
     created_time: int
     last_updated_time: int
     owner_id: str
-    external_id: str
-    name: str
-    description: str | None = None
-    instructions: str | None = None
-    model: str = "azure/gpt-4o-mini"
-    tools: list[AgentTool] | None = None
     runtime_version: str
 
     def as_request_resource(self) -> AgentRequest:

cognite_toolkit/_cdf_tk/client/data_classes/annotation.py

@@ -0,0 +1,79 @@
+from typing import Any, Literal, TypeAlias
+
+from pydantic import Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestUpdateable, ResponseResource
+
+from .identifiers import InternalId
+
+AnnotationStatus: TypeAlias = Literal["suggested", "rejected", "approved"]
+AnnotationType: TypeAlias = Literal[
+    "diagrams.AssetLink",
+    "diagrams.FileLink",
+    "diagrams.InstanceLink",
+    "diagrams.Junction",
+    "diagrams.Line",
+    "diagrams.UnhandledSymbolObject",
+    "diagrams.UnhandledTextObject",
+    "documents.ExtractedText",
+    "forms.Detection",
+    "images.AssetLink",
+    "images.Classification",
+    "images.InstanceLink",
+    "images.KeypointCollection",
+    "images.ObjectDetection",
+    "images.TextRegion",
+    "isoplan.IsoPlanAnnotation",
+    "pointcloud.BoundingVolume",
+]
+
+
+class Annotation(BaseModelObject):
+    annotated_resource_type: str
+    annotated_resource_id: int
+    annotation_type: AnnotationType
+    creating_app: str
+    creating_app_version: str
+    creating_user: str | None
+    data: dict[str, JsonValue]
+    status: AnnotationStatus
+
+
+class AnnotationRequest(Annotation, RequestUpdateable):
+    """Request data class for annotations."""
+
+    # The 'id' field is not part of the request when creating a new resource,
+    # but is needed when updating an existing resource.
+    id: int | None = Field(default=None, exclude=True)
+
+    def as_id(self) -> InternalId:
+        if self.id is None:
+            raise ValueError("Cannot convert AnnotationRequest to InternalId when id is None")
+        return InternalId(id=self.id)
+
+    def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
+        """Converts the request to an update payload for the API."""
+        if self.id is None:
+            raise ValueError("id must be provided to create an update dictionary")
+        return {
+            "id": self.id,
+            "update": {
+                "annotationType": {"set": self.annotation_type},
+                "data": {"set": self.data},
+                "status": {"set": self.status},
+            },
+        }
+
+
+class AnnotationResponse(Annotation, ResponseResource[AnnotationRequest]):
+    """Response data class for annotations."""
+
+    id: int
+    created_time: int
+    last_updated_time: int
+
+    def as_id(self) -> InternalId:
+        return InternalId(id=self.id)
+
+    def as_request_resource(self) -> AnnotationRequest:
+        return AnnotationRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/asset.py

@@ -7,9 +7,7 @@ from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, Re
 from .identifiers import ExternalId
 
 
-class AssetRequest(RequestUpdateable):
-    container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "labels"})
-    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"parent_id", "parent_external_id"})
+class Asset(BaseModelObject):
     external_id: str | None = None
     name: str
     parent_id: int | None = None
@@ -27,28 +25,23 @@ class AssetRequest(RequestUpdateable):
         return ExternalId(external_id=self.external_id)
 
 
+class AssetRequest(Asset, RequestUpdateable):
+    container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "labels"})
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"parent_id", "parent_external_id"})
+
+
 class AssetAggregateItem(BaseModelObject):
     child_count: int
     depth: int
     path: list[dict[Literal["id"], int]]
 
 
-class AssetResponse(ResponseResource[AssetRequest]):
+class AssetResponse(Asset, ResponseResource[AssetRequest]):
     created_time: int
     last_updated_time: int
     root_id: int
     aggregates: AssetAggregateItem | None = None
     id: int
-    external_id: str | None = None
-    name: str
-    parent_id: int | None = None
-    parent_external_id: str | None = None
-    description: str | None = None
-    metadata: dict[str, str] | None = None
-    data_set_id: int | None = None
-    source: str | None = None
-    labels: list[dict[Literal["externalId"], str]] | None = None
-    geo_location: dict[str, JsonValue] | None = None
 
     def as_request_resource(self) -> AssetRequest:
         return AssetRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -18,19 +18,29 @@ class BaseModelObject(BaseModel):
     """Base class for all object. This includes resources and nested objects."""
 
     # We allow extra fields to support forward compatibility.
-    model_config = ConfigDict(alias_generator=to_camel, extra="allow")
+    model_config = ConfigDict(alias_generator=to_camel, extra="allow", populate_by_name=True)
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
         """Dump the resource to a dictionary.
 
-        This is the default serialization method for request resources.
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+            exclude_extra (bool): Whether to exclude extra fields not defined in the model. Default is False.
+
         """
+        if exclude_extra:
+            return self.model_dump(
+                mode="json",
+                by_alias=camel_case,
+                exclude_unset=True,
+                exclude=set(self.__pydantic_extra__) if self.__pydantic_extra__ else None,
+            )
         return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
 
     @classmethod
     def _load(cls, resource: dict[str, Any]) -> Self:
         """Load method to match CogniteResource signature."""
-        return cls.model_validate(resource)
+        return cls.model_validate(resource, by_alias=True)
 
 
 class Identifier(BaseModelObject):
@@ -39,7 +49,7 @@ class Identifier(BaseModelObject):
 
     model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True, frozen=True)
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
         """Dump the resource to a dictionary.
 
         This is the default serialization method for request resources.

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/__init__.py

@@ -0,0 +1,164 @@
+from ._constraints import (
+    Constraint,
+    ConstraintAdapter,
+    ConstraintDefinition,
+    RequiresConstraintDefinition,
+    UniquenessConstraintDefinition,
+)
+from ._container import (
+    Container,
+    ContainerPropertyDefinition,
+    ContainerRequest,
+    ContainerResponse,
+)
+from ._data_model import DataModelRequest, DataModelResponse
+from ._data_types import (
+    BooleanProperty,
+    DataType,
+    DataTypeAdapter,
+    DateProperty,
+    DirectNodeRelation,
+    EnumProperty,
+    EnumValue,
+    FileCDFExternalIdReference,
+    Float32Property,
+    Float64Property,
+    FloatProperty,
+    Int32Property,
+    Int64Property,
+    JSONProperty,
+    ListablePropertyTypeDefinition,
+    PropertyTypeDefinition,
+    SequenceCDFExternalIdReference,
+    TextProperty,
+    TimeseriesCDFExternalIdReference,
+    TimestampProperty,
+    Unit,
+)
+from ._indexes import (
+    BtreeIndex,
+    Index,
+    IndexAdapter,
+    IndexDefinition,
+    InvertedIndex,
+)
+from ._instance import (
+    EdgeRequest,
+    EdgeResponse,
+    InstanceDefinition,
+    InstanceResponseDefinition,
+    InstanceSource,
+    NodeRequest,
+    NodeResponse,
+)
+from ._references import (
+    ContainerConstraintReference,
+    ContainerDirectReference,
+    ContainerIndexReference,
+    ContainerReference,
+    DataModelReference,
+    NodeReference,
+    SpaceReference,
+    ViewDirectReference,
+    ViewReference,
+)
+from ._space import Space, SpaceRequest, SpaceResponse
+from ._view import View, ViewRequest, ViewResponse
+from ._view_property import (
+    ConnectionPropertyDefinition,
+    ConstraintOrIndexState,
+    EdgeProperty,
+    MultiEdgeProperty,
+    MultiReverseDirectRelationPropertyRequest,
+    MultiReverseDirectRelationPropertyResponse,
+    ReverseDirectRelationProperty,
+    SingleEdgeProperty,
+    SingleReverseDirectRelationPropertyRequest,
+    SingleReverseDirectRelationPropertyResponse,
+    ViewCoreProperty,
+    ViewCorePropertyRequest,
+    ViewCorePropertyResponse,
+    ViewPropertyDefinition,
+    ViewRequestProperty,
+    ViewRequestPropertyAdapter,
+    ViewResponseProperty,
+)
+
+__all__ = [
+    "BooleanProperty",
+    "BtreeIndex",
+    "ConnectionPropertyDefinition",
+    "Constraint",
+    "ConstraintAdapter",
+    "ConstraintDefinition",
+    "ConstraintOrIndexState",
+    "Container",
+    "ContainerConstraintReference",
+    "ContainerDirectReference",
+    "ContainerIndexReference",
+    "ContainerPropertyDefinition",
+    "ContainerReference",
+    "ContainerRequest",
+    "ContainerResponse",
+    "DataModelReference",
+    "DataModelRequest",
+    "DataModelResponse",
+    "DataType",
+    "DataTypeAdapter",
+    "DateProperty",
+    "DirectNodeRelation",
+    "EdgeProperty",
+    "EdgeRequest",
+    "EdgeResponse",
+    "EnumProperty",
+    "EnumValue",
+    "FileCDFExternalIdReference",
+    "Float32Property",
+    "Float64Property",
+    "FloatProperty",
+    "Index",
+    "IndexAdapter",
+    "IndexDefinition",
+    "InstanceDefinition",
+    "InstanceResponseDefinition",
+    "InstanceSource",
+    "Int32Property",
+    "Int64Property",
+    "InvertedIndex",
+    "JSONProperty",
+    "ListablePropertyTypeDefinition",
+    "MultiEdgeProperty",
+    "MultiReverseDirectRelationPropertyRequest",
+    "MultiReverseDirectRelationPropertyResponse",
+    "NodeReference",
+    "NodeRequest",
+    "NodeResponse",
+    "PropertyTypeDefinition",
+    "RequiresConstraintDefinition",
+    "ReverseDirectRelationProperty",
+    "SequenceCDFExternalIdReference",
+    "SingleEdgeProperty",
+    "SingleReverseDirectRelationPropertyRequest",
+    "SingleReverseDirectRelationPropertyResponse",
+    "Space",
+    "SpaceReference",
+    "SpaceRequest",
+    "SpaceResponse",
+    "TextProperty",
+    "TimeseriesCDFExternalIdReference",
+    "TimestampProperty",
+    "UniquenessConstraintDefinition",
+    "Unit",
+    "View",
+    "ViewCoreProperty",
+    "ViewCorePropertyRequest",
+    "ViewCorePropertyResponse",
+    "ViewDirectReference",
+    "ViewPropertyDefinition",
+    "ViewReference",
+    "ViewRequest",
+    "ViewRequestProperty",
+    "ViewRequestPropertyAdapter",
+    "ViewResponse",
+    "ViewResponseProperty",
+]

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/_constraints.py

@@ -0,0 +1,37 @@
+from abc import ABC
+from typing import Annotated, Any, Literal
+
+from pydantic import Field, TypeAdapter, field_serializer
+from pydantic_core.core_schema import FieldSerializationInfo
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject
+
+from ._references import ContainerReference
+
+
+class ConstraintDefinition(BaseModelObject, ABC):
+    constraint_type: str
+
+
+class UniquenessConstraintDefinition(ConstraintDefinition):
+    constraint_type: Literal["uniqueness"] = "uniqueness"
+    properties: list[str]
+    by_space: bool | None = None
+
+
+class RequiresConstraintDefinition(ConstraintDefinition):
+    constraint_type: Literal["requires"] = "requires"
+    require: ContainerReference
+
+    @field_serializer("require", mode="plain")
+    @classmethod
+    def serialize_require(cls, require: ContainerReference, info: FieldSerializationInfo) -> dict[str, Any]:
+        return {**require.model_dump(**vars(info)), "type": "container"}
+
+
+Constraint = Annotated[
+    UniquenessConstraintDefinition | RequiresConstraintDefinition,
+    Field(discriminator="constraint_type"),
+]
+
+ConstraintAdapter: TypeAdapter[Constraint] = TypeAdapter(Constraint)

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/_container.py

@@ -0,0 +1,50 @@
+from abc import ABC
+from typing import Literal
+
+from pydantic import JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestResource, ResponseResource
+
+from ._constraints import Constraint
+from ._data_types import DataType
+from ._indexes import Index
+from ._references import ContainerReference
+
+
+class ContainerPropertyDefinition(BaseModelObject):
+    immutable: bool | None = None
+    nullable: bool | None = None
+    auto_increment: bool | None = None
+    default_value: str | int | float | bool | dict[str, JsonValue] | None = None
+    description: str | None = None
+    name: str | None = None
+    type: DataType
+
+
+class Container(BaseModelObject, ABC):
+    space: str
+    external_id: str
+    name: str | None = None
+    description: str | None = None
+    used_for: Literal["node", "edge", "all"] | None = None
+    properties: dict[str, ContainerPropertyDefinition]
+    constraints: dict[str, Constraint] | None = None
+    indexes: dict[str, Index] | None = None
+
+    def as_id(self) -> ContainerReference:
+        return ContainerReference(
+            space=self.space,
+            external_id=self.external_id,
+        )
+
+
+class ContainerRequest(Container, RequestResource): ...
+
+
+class ContainerResponse(Container, ResponseResource[ContainerRequest]):
+    created_time: int
+    last_updated_time: int
+    is_global: bool
+
+    def as_request_resource(self) -> "ContainerRequest":
+        return ContainerRequest.model_validate(self.model_dump(by_alias=True))

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/_data_model.py

@@ -0,0 +1,73 @@
+from abc import ABC
+from typing import Any
+
+from pydantic import field_serializer
+from pydantic_core.core_schema import FieldSerializationInfo
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestResource, ResponseResource
+
+from ._references import DataModelReference, ViewReference
+from ._view import ViewResponse
+
+
+class DataModel(BaseModelObject, ABC):
+    """Cognite Data Model resource.
+
+    Data models group and structure views into reusable collections.
+    A data model contains a set of views where the node types can
+    refer to each other with direct relations and edges.
+    """
+
+    space: str
+    external_id: str
+    version: str
+    description: str | None = None
+
+    def as_id(self) -> DataModelReference:
+        return DataModelReference(
+            space=self.space,
+            external_id=self.external_id,
+            version=self.version,
+        )
+
+
+class DataModelRequest(DataModel, RequestResource):
+    views: list[ViewReference] | None = None
+
+    @field_serializer("views", mode="plain")
+    @classmethod
+    def serialize_views(
+        cls, views: list[ViewReference] | None, info: FieldSerializationInfo
+    ) -> list[dict[str, Any]] | None:
+        if views is None:
+            return None
+        return [{**view.model_dump(**vars(info)), "type": "view"} for view in views]
+
+
+class DataModelResponse(DataModel, ResponseResource[DataModelRequest]):
+    views: list[ViewReference] | None = None
+    created_time: int
+    last_updated_time: int
+    is_global: bool
+
+    def as_request_resource(self) -> DataModelRequest:
+        return DataModelRequest.model_validate(self.model_dump(by_alias=True), extra="ignore")
+
+    @field_serializer("views", mode="plain")
+    @classmethod
+    def serialize_views(
+        cls, views: list[ViewReference] | None, info: FieldSerializationInfo
+    ) -> list[dict[str, Any]] | None:
+        if views is None:
+            return None
+        return [{**view.model_dump(**vars(info)), "type": "view"} for view in views]
+
+
+class DataModelResponseWithViews(DataModel, ResponseResource[DataModelRequest]):
+    views: list[ViewResponse] | None = None
+    created_time: int
+    last_updated_time: int
+    is_global: bool
+
+    def as_request_resource(self) -> DataModelRequest:
+        return DataModelRequest.model_validate(self.model_dump(by_alias=True), extra="ignore")