cognite-toolkit 0.7.39__py3-none-any.whl → 0.7.51__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/resource_classes/agent.py

@@ -0,0 +1,130 @@
+from typing import Annotated, Any, Literal
+
+from pydantic import BeforeValidator, Field
+
+from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject, RequestResource, ResponseResource
+
+from .identifiers import ExternalId
+
+
+class AgentToolDefinition(BaseModelObject):
+    type: str
+    name: str
+    description: str
+
+
+class AskDocument(AgentToolDefinition):
+    type: Literal["askDocument"] = "askDocument"
+
+
+class ExamineDataSemantically(AgentToolDefinition):
+    type: Literal["examineDataSemantically"] = "examineDataSemantically"
+
+
+class AgentDataModel(BaseModelObject):
+    space: str
+    external_id: str
+    version: str
+    view_external_ids: list[str] | None = None
+
+
+class AgentInstanceSpacesDefinition(BaseModelObject):
+    type: str
+
+
+class AllInstanceSpaces(AgentInstanceSpacesDefinition):
+    type: Literal["all"] = "all"
+
+
+class ManualInstanceSpaces(AgentInstanceSpacesDefinition):
+    type: Literal["manual"] = "manual"
+    spaces: list[str]
+
+
+AgentInstanceSpaces = Annotated[
+    AllInstanceSpaces | ManualInstanceSpaces,
+    Field(discriminator="type"),
+]
+
+
+class QueryKnowledgeGraphConfig(BaseModelObject):
+    data_models: list[AgentDataModel]
+    instance_spaces: AgentInstanceSpaces | None = None
+    # This is deviating from the API documentation, but the Atlas team has confirmed that "v2" is the default
+    version: Literal["v1", "v2"] = "v2"
+
+
+class QueryKnowledgeGraph(AgentToolDefinition):
+    type: Literal["queryKnowledgeGraph"] = "queryKnowledgeGraph"
+    configuration: QueryKnowledgeGraphConfig
+
+
+class QueryTimeSeriesDatapoints(AgentToolDefinition):
+    type: Literal["queryTimeSeriesDatapoints"] = "queryTimeSeriesDatapoints"
+
+
+class SummarizeDocument(AgentToolDefinition):
+    type: Literal["summarizeDocument"] = "summarizeDocument"
+
+
+class UnknownAgentTool(AgentToolDefinition):
+    """Fallback for unknown tool types."""
+
+    ...
+
+
+# Mapping of known agent tool types to their classes
+KNOWN_TOOLS: dict[str, type[AgentToolDefinition]] = {
+    "askDocument": AskDocument,
+    "examineDataSemantically": ExamineDataSemantically,
+    "queryKnowledgeGraph": QueryKnowledgeGraph,
+    "queryTimeSeriesDatapoints": QueryTimeSeriesDatapoints,
+    "summarizeDocument": SummarizeDocument,
+}
+
+
+def _handle_unknown_tool(value: Any) -> Any:
+    if isinstance(value, dict):
+        tool_type = value.get("type")
+        if tool_type not in KNOWN_TOOLS:
+            return UnknownAgentTool(**value)
+        else:
+            return KNOWN_TOOLS[tool_type].model_validate(value)
+    return value
+
+
+AgentTool = Annotated[
+    AskDocument
+    | QueryKnowledgeGraph
+    | QueryTimeSeriesDatapoints
+    | SummarizeDocument
+    | ExamineDataSemantically
+    | UnknownAgentTool,
+    BeforeValidator(_handle_unknown_tool),
+]
+
+
+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
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class AgentRequest(Agent, RequestResource):
+    runtime_version: str | None = None
+
+
+class AgentResponse(Agent, ResponseResource[AgentRequest]):
+    created_time: int
+    last_updated_time: int
+    owner_id: str
+    runtime_version: str
+
+    def as_request_resource(self) -> AgentRequest:
+        return AgentRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/resource_classes/annotation.py

@@ -0,0 +1,79 @@
+from typing import Any, Literal, TypeAlias
+
+from pydantic import Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject, ResponseResource, UpdatableRequestResource
+
+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, UpdatableRequestResource):
+    """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/resource_classes/apm_config.py

@@ -0,0 +1,128 @@
+from typing import ClassVar, Literal
+
+from cognite_toolkit._cdf_tk.client._resource_base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+)
+
+from .data_modeling import ViewReference
+from .identifiers import ExternalId
+
+
+class ThreeDModelIdentifier(BaseModelObject):
+    revision_id: int | None = None
+    model_id: int | None = None
+    name: str | None = None
+
+
+class ThreeDConfiguration(BaseModelObject):
+    full_weight_models: list[ThreeDModelIdentifier] | None = None
+    light_weight_models: list[ThreeDModelIdentifier] | None = None
+
+
+class ResourceFilters(BaseModelObject):
+    data_set_ids: list[int] | None = None
+    asset_subtree_external_ids: list[str] | None = None
+    root_asset_external_ids: list[str] | None = None
+    external_id_prefix: str | None = None
+    spaces: list[str] | None = None
+
+
+class RootLocationDataFilters(BaseModelObject):
+    general: ResourceFilters | None = None
+    assets: ResourceFilters | None = None
+    files: ResourceFilters | None = None
+    timeseries: ResourceFilters | None = None
+
+
+class ObservationFeatureToggles(BaseModelObject):
+    is_enabled: bool | None = None
+    is_write_back_enabled: bool | None = None
+    notifications_endpoint_external_id: str | None = None
+    attachments_endpoint_external_id: str | None = None
+
+
+class RootLocationFeatureToggles(BaseModelObject):
+    three_d: bool | None = None
+    trends: bool | None = None
+    documents: bool | None = None
+    workorders: bool | None = None
+    notifications: bool | None = None
+    media: bool | None = None
+    template_checklist_flow: bool | None = None
+    workorder_checklist_flow: bool | None = None
+    observations: ObservationFeatureToggles | None = None
+
+
+class ObservationConfigFieldProperty(BaseModelObject):
+    display_title: str | None = None
+    display_description: str | None = None
+    is_required: bool | None = None
+
+
+class ObservationConfigDropdownPropertyOption(BaseModelObject):
+    id: str | None = None
+    value: str | None = None
+    label: str | None = None
+
+
+class ObservationConfigDropdownProperty(ObservationConfigFieldProperty):
+    options: list[ObservationConfigDropdownPropertyOption] | None = None
+
+
+class ObservationsConfig(BaseModelObject):
+    files: ObservationConfigFieldProperty | None = None
+    description: ObservationConfigFieldProperty | None = None
+    asset: ObservationConfigFieldProperty | None = None
+    troubleshooting: ObservationConfigFieldProperty | None = None
+    type: ObservationConfigDropdownProperty | None = None
+    priority: ObservationConfigDropdownProperty | None = None
+
+
+class RootLocationConfiguration(BaseModelObject):
+    asset_external_id: str | None = None
+    external_id: str | None = None
+    display_name: str | None = None
+    three_d_configuration: ThreeDConfiguration | None = None
+    data_set_id: int | None = None
+    template_admins: list[str] | None = None  # list of Group Names
+    checklist_admins: list[str] | None = None  # list of Group Names
+    app_data_instance_space: str | None = None
+    source_data_instance_space: str | None = None
+    data_filters: RootLocationDataFilters | None = None
+    feature_toggles: RootLocationFeatureToggles | None = None
+    observations: ObservationsConfig | None = None
+
+
+class FeatureConfiguration(BaseModelObject):
+    root_location_configurations: list[RootLocationConfiguration] | None = None
+
+
+class APMConfig(BaseModelObject):
+    space: ClassVar[str] = "APM_Config"
+    view_ref: ClassVar[ViewReference] = ViewReference(space="APM_Config", external_id="APM_Config", version="1")
+    instance_type: Literal["node"] = "node"
+    external_id: str
+    name: str | None = None
+    app_data_space_id: str | None = None
+    app_data_space_version: str | None = None
+    customer_data_space_id: str | None = None
+    customer_data_space_version: str | None = None
+    feature_configuration: FeatureConfiguration | None = None
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class APMConfigRequest(APMConfig, RequestResource):
+    existing_version: int | None = None
+
+
+class APMConfigResponse(APMConfig, ResponseResource[APMConfigRequest]):
+    version: int
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> APMConfigRequest:
+        return APMConfigRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/resource_classes/asset.py

@@ -0,0 +1,47 @@
+from typing import ClassVar, Literal
+
+from pydantic import JsonValue
+
+from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject, ResponseResource, UpdatableRequestResource
+
+from .identifiers import ExternalId
+
+
+class Asset(BaseModelObject):
+    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_id(self) -> ExternalId:
+        if self.external_id is None:
+            raise ValueError("Cannot convert AssetRequest to ExternalId when external_id is None")
+        return ExternalId(external_id=self.external_id)
+
+
+class AssetRequest(Asset, UpdatableRequestResource):
+    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(Asset, ResponseResource[AssetRequest]):
+    created_time: int
+    last_updated_time: int
+    root_id: int
+    aggregates: AssetAggregateItem | None = None
+    id: int
+
+    def as_request_resource(self) -> AssetRequest:
+        return AssetRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/{data_classes → resource_classes}/charts_data.py

@@ -3,7 +3,7 @@ from typing import Any
 from cognite.client.data_classes.data_modeling import NodeId, ViewId
 from pydantic import JsonValue, field_serializer, field_validator
 
-from .base import BaseModelObject
+from cognite_toolkit._cdf_tk.client._resource_base import BaseModelObject
 
 
 class UserInfo(BaseModelObject):

cognite_toolkit/_cdf_tk/client/resource_classes/cognite_file.py

@@ -0,0 +1,53 @@
+from datetime import datetime
+from typing import ClassVar, Literal
+
+from cognite_toolkit._cdf_tk.client._resource_base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+)
+
+from .data_modeling import ViewReference
+from .instance_api import NodeReference, TypedNodeIdentifier
+
+
+class CogniteFile(BaseModelObject):
+    view_ref: ClassVar[ViewReference] = ViewReference(space="cdf_cdm", external_id="CogniteFile", version="v1")
+    instance_type: Literal["node"] = "node"
+    space: str
+    external_id: str
+    name: str | None = None
+    description: str | None = None
+    tags: list[str] | None = None
+    aliases: list[str] | None = None
+    source_id: str | None = None
+    source_context: str | None = None
+    source: NodeReference | None = None
+    source_created_time: datetime | None = None
+    source_updated_time: datetime | None = None
+    source_created_user: str | None = None
+    source_updated_user: str | None = None
+    assets: list[NodeReference] | None = None
+    mime_type: str | None = None
+    directory: str | None = None
+    category: NodeReference | None = None
+    type: NodeReference | None = None
+
+    def as_id(self) -> TypedNodeIdentifier:
+        return TypedNodeIdentifier(space=self.space, external_id=self.external_id)
+
+
+class CogniteFileRequest(CogniteFile, RequestResource):
+    existing_version: int | None = None
+
+
+class CogniteFileResponse(CogniteFile, ResponseResource[CogniteFileRequest]):
+    version: int
+    created_time: int
+    last_updated_time: int
+    is_uploaded: bool | None = None
+    uploaded_time: datetime | None = None
+    deleted_time: int | None = None
+
+    def as_request_resource(self) -> CogniteFileRequest:
+        return CogniteFileRequest.model_validate(self.dump(), extra="ignore")

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

@@ -0,0 +1,168 @@
+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,
+    InstanceRequest,
+    InstanceResponse,
+    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",
+    "InstanceRequest",
+    "InstanceResponse",
+    "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/resource_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._resource_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)