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

cognite_toolkit/_cdf_tk/client/http_client/_data_classes2.py

@@ -0,0 +1,151 @@
+import gzip
+from abc import ABC, abstractmethod
+from collections.abc import Hashable
+from typing import Any, Literal
+
+import httpx
+from cognite.client import global_config
+from pydantic import BaseModel, JsonValue, TypeAdapter, model_validator
+
+from cognite_toolkit._cdf_tk.client.http_client._data_classes import (
+    ErrorDetails,
+    FailedRequestItems,
+    FailedRequestMessage,
+    FailedResponseItems,
+    ResponseMessage,
+    SuccessResponseItems,
+)
+from cognite_toolkit._cdf_tk.client.http_client._exception import ToolkitAPIError
+from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
+
+
+class HTTPResult2(BaseModel):
+    def get_success_or_raise(self) -> "SuccessResponse2":
+        """Raises an exception if any response in the list indicates a failure."""
+        if isinstance(self, SuccessResponse2):
+            return self
+        elif isinstance(self, FailedResponse2):
+            raise ToolkitAPIError(
+                f"Request failed with status code {self.status_code}: {self.error.message}",
+                missing=self.error.missing,  # type: ignore[arg-type]
+                duplicated=self.error.duplicated,  # type: ignore[arg-type]
+            )
+        elif isinstance(self, FailedRequest2):
+            raise ToolkitAPIError(f"Request failed with error: {self.error}")
+        else:
+            raise ToolkitAPIError("Unknown HTTPResult2 type")
+
+    # Todo: Remove when HTTPResult2 is renamed to HTTPResponse and the old HTTPResponse is deleted
+    def as_item_response(self, item_id: Hashable) -> ResponseMessage | FailedRequestMessage:
+        if isinstance(self, SuccessResponse2):
+            return SuccessResponseItems(
+                status_code=self.status_code, content=self.content, ids=[item_id], body=self.body
+            )
+        elif isinstance(self, FailedResponse2):
+            return FailedResponseItems(
+                status_code=self.status_code,
+                ids=[item_id],
+                body=self.body,
+                error=ErrorDetails(
+                    code=self.error.code,
+                    message=self.error.message,
+                    missing=self.error.missing,
+                    duplicated=self.error.duplicated,
+                ),
+            )
+        elif isinstance(self, FailedRequest2):
+            return FailedRequestItems(ids=[item_id], error=self.error)
+        else:
+            raise ToolkitAPIError("Unknown HTTPResult2 type")
+
+
+class FailedRequest2(HTTPResult2):
+    error: str
+
+
+class SuccessResponse2(HTTPResult2):
+    status_code: int
+    body: str
+    content: bytes
+
+    @property
+    def body_json(self) -> dict[str, Any]:
+        """Parse the response body as JSON."""
+        return TypeAdapter(dict[str, JsonValue]).validate_json(self.body)
+
+
+class ErrorDetails2(BaseModel):
+    """This is the expected structure of error details in the CDF API"""
+
+    code: int
+    message: str
+    missing: list[JsonValue] | None = None
+    duplicated: list[JsonValue] | None = None
+    is_auto_retryable: bool | None = None
+
+    @classmethod
+    def from_response(cls, response: httpx.Response) -> "ErrorDetails2":
+        """Populate the error details from a httpx response."""
+        try:
+            res = TypeAdapter(dict[Literal["error"], ErrorDetails2]).validate_json(response.text)
+        except ValueError:
+            return cls(code=response.status_code, message=response.text)
+        return res["error"]
+
+
+class FailedResponse2(HTTPResult2):
+    status_code: int
+    body: str
+    error: ErrorDetails2
+
+
+class BaseRequestMessage(BaseModel, ABC):
+    endpoint_url: str
+    method: Literal["GET", "POST", "PATCH", "DELETE", "PUT"]
+    connect_attempt: int = 0
+    read_attempt: int = 0
+    status_attempt: int = 0
+    api_version: str | None = None
+    disable_gzip: bool = False
+    content_length: int | None = None
+    content_type: str = "application/json"
+    accept: str = "application/json"
+
+    parameters: dict[str, PrimitiveType] | None = None
+
+    @property
+    def total_attempts(self) -> int:
+        return self.connect_attempt + self.read_attempt + self.status_attempt
+
+    @property
+    @abstractmethod
+    def content(self) -> str | bytes | None: ...
+
+
+class RequestMessage2(BaseRequestMessage):
+    data_content: bytes | None = None
+    body_content: dict[str, JsonValue] | None = None
+
+    @model_validator(mode="before")
+    def check_data_or_body(cls, values: dict[str, Any]) -> dict[str, Any]:
+        if values.get("data_content") is not None and values.get("body_content") is not None:
+            raise ValueError("Only one of data_content or body_content can be set.")
+        return values
+
+    @property
+    def content(self) -> str | bytes | None:
+        data: str | bytes | None = None
+        if self.data_content is not None:
+            data = self.data_content
+            if not global_config.disable_gzip and not self.disable_gzip:
+                data = gzip.compress(data)
+        elif self.body_content is not None:
+            # We serialize using pydantic instead of json.dumps. This is because pydantic is faster
+            # and handles more complex types such as datetime, float('nan'), etc.
+            data = _BODY_SERIALIZER.dump_json(self.body_content)
+            if not global_config.disable_gzip and not self.disable_gzip and isinstance(data, bytes):
+                data = gzip.compress(data)
+        return data
+
+
+_BODY_SERIALIZER = TypeAdapter(dict[str, JsonValue])

cognite_toolkit/_cdf_tk/client/http_client/_exception.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+
+class ToolkitAPIError(Exception):
+    """Base class for all exceptions raised by the Cognite Toolkit API client."""
+
+    def __init__(
+        self, message: str, missing: list[dict[str, Any]] | None = None, duplicated: list[dict[str, Any]] | None = None
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.missing = missing
+        self.duplicated = duplicated

cognite_toolkit/_cdf_tk/client/http_client/_item_classes.py

@@ -0,0 +1,118 @@
+"""Data classes for handling item-based requests and responses in the Cognite Toolkit HTTP client."""
+
+import gzip
+from collections import UserList
+from collections.abc import Sequence
+from typing import Any
+
+from cognite.client import global_config
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client._resource_base import RequestItem
+from cognite_toolkit._cdf_tk.client.http_client._data_classes2 import (
+    _BODY_SERIALIZER,
+    BaseRequestMessage,
+    ErrorDetails2,
+)
+from cognite_toolkit._cdf_tk.client.http_client._exception import ToolkitAPIError
+from cognite_toolkit._cdf_tk.client.http_client._tracker import ItemsRequestTracker
+
+
+class ItemsResultMessage2(BaseModel):
+    ids: list[str]
+
+
+class ItemsFailedRequest2(ItemsResultMessage2):
+    error_message: str
+
+
+class ItemsSuccessResponse2(ItemsResultMessage2):
+    status_code: int
+    body: str
+    content: bytes
+
+
+class ItemsFailedResponse2(ItemsResultMessage2):
+    status_code: int
+    error: ErrorDetails2
+    body: str
+
+
+def _set_default_tracker(data: dict[str, Any]) -> ItemsRequestTracker:
+    if "tracker" not in data or data["tracker"] is None:
+        return ItemsRequestTracker(data.get("max_failures_before_abort", 50))
+    return data["tracker"]
+
+
+class ItemsRequest2(BaseRequestMessage):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    items: Sequence[RequestItem]
+    extra_body_fields: dict[str, JsonValue] | None = None
+    max_failures_before_abort: int = 50
+    tracker: ItemsRequestTracker = Field(init=False, default_factory=_set_default_tracker, exclude=True)
+
+    @property
+    def content(self) -> str | bytes | None:
+        body: dict[str, JsonValue] = {"items": [item.dump(camel_case=True) for item in self.items]}
+        if self.extra_body_fields:
+            body.update(self.extra_body_fields)
+        res = _BODY_SERIALIZER.dump_json(body)
+        if not global_config.disable_gzip and not self.disable_gzip and isinstance(res, bytes):
+            return gzip.compress(res)
+        return res
+
+    def split(self, status_attempts: int) -> list["ItemsRequest2"]:
+        """Split the request into multiple requests with a single item each."""
+        mid = len(self.items) // 2
+        if mid == 0:
+            return [self]
+        self.tracker.register_failure()
+        messages: list[ItemsRequest2] = []
+        for part in (self.items[:mid], self.items[mid:]):
+            new_request = self.model_copy(update={"items": part, "status_attempt": status_attempts})
+            new_request.tracker = self.tracker
+            messages.append(new_request)
+        return messages
+
+
+class ItemResponse(BaseModel):
+    items: list[dict[str, JsonValue]]
+
+
+class ItemsResultList(UserList[ItemsResultMessage2]):
+    def __init__(self, collection: Sequence[ItemsResultMessage2] | None = None) -> None:
+        super().__init__(collection or [])
+
+    def raise_for_status(self) -> None:
+        """Raises an exception if any response in the list indicates a failure."""
+        failed_responses = [resp for resp in self.data if isinstance(resp, ItemsFailedResponse2)]
+        failed_requests = [resp for resp in self.data if isinstance(resp, ItemsFailedRequest2)]
+        if not failed_responses and not failed_requests:
+            return
+        error_messages = "; ".join(f"Status {err.status_code}: {err.error.message}" for err in failed_responses)
+        if failed_requests:
+            if error_messages:
+                error_messages += "; "
+            error_messages += "; ".join(f"Request error: {err.error_message}" for err in failed_requests)
+        raise ToolkitAPIError(f"One or more requests failed: {error_messages}")
+
+    @property
+    def has_failed(self) -> bool:
+        """Indicates whether any response in the list indicates a failure.
+
+        Returns:
+            bool: True if there are any failed responses or requests, False otherwise.
+        """
+        for resp in self.data:
+            if isinstance(resp, ItemsFailedResponse2 | ItemsFailedRequest2):
+                return True
+        return False
+
+    def get_items(self) -> list[dict[str, JsonValue]]:
+        """Get the items from all successful responses."""
+        items: list[dict[str, JsonValue]] = []
+        for resp in self.data:
+            if isinstance(resp, ItemsSuccessResponse2):
+                body_json = ItemResponse.model_validate_json(resp.body)
+                items.extend(body_json.items)
+        return items

cognite_toolkit/_cdf_tk/{utils → client}/http_client/_tracker.py

@@ -1,5 +1,6 @@
 import threading
 from dataclasses import dataclass, field
+from typing import Any
 
 
 @dataclass
@@ -9,13 +10,15 @@ class ItemsRequestTracker:
     Attributes:
         max_failures_before_abort (int): Maximum number of allowed failed split requests before aborting
             the entire operation. A value of -1 indicates no early abort.
-        lock (threading.Lock): A lock to ensure thread-safe updates to the failure count.
+        lock (Any): A lock to ensure thread-safe updates to the failure count.
         failed_split_count (int): The current count of failed split requests.
 
     """
 
     max_failures_before_abort: int = -1  # -1 means no early abort
-    lock: threading.Lock = field(default_factory=threading.Lock, init=False)
+    # NOTE: `threading.Lock` is a factory function (backed by `_thread.allocate_lock`), not a type.
+    # Annotate as `Any` so Pydantic won't attempt to generate a schema for the lock field.
+    lock: Any = field(default_factory=threading.Lock, init=False)
     failed_split_count: int = field(default=0, init=False)
 
     def register_failure(self) -> None:

cognite_toolkit/_cdf_tk/client/request_classes/base.py

@@ -0,0 +1,19 @@
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class BaseModelRequest(BaseModel):
+    """Base class for all object. This includes resources and nested objects."""
+
+    model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True)
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the object to a dictionary.
+
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)

cognite_toolkit/_cdf_tk/client/request_classes/filters.py

@@ -0,0 +1,113 @@
+import sys
+from typing import Any, Literal
+
+from pydantic import Field
+
+from cognite_toolkit._cdf_tk.client.resource_classes.annotation import AnnotationStatus, AnnotationType
+from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import NodeReference, ViewReference
+from cognite_toolkit._cdf_tk.client.resource_classes.identifiers import ExternalId, InternalId
+
+from .base import BaseModelRequest
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+class Filter(BaseModelRequest): ...
+
+
+class ClassicFilter(Filter):
+    asset_subtree_ids: list[ExternalId | InternalId] | None = None
+    data_set_ids: list[ExternalId | InternalId] | None = None
+
+    @classmethod
+    def from_asset_subtree_and_data_sets(
+        cls,
+        asset_subtree_id: str | int | list[str | int] | None = None,
+        data_set_id: str | int | list[str | int] | None = None,
+    ) -> Self:
+        return cls(
+            asset_subtree_ids=cls._as_internal_or_external_id_list(asset_subtree_id),
+            data_set_ids=cls._as_internal_or_external_id_list(data_set_id),
+        )
+
+    @classmethod
+    def _as_internal_or_external_id_list(
+        cls, id: str | int | list[str | int] | None
+    ) -> list[ExternalId | InternalId] | None:
+        if id is None:
+            return None
+        ids = id if isinstance(id, list) else [id]
+        return [ExternalId(external_id=item) if isinstance(item, str) else InternalId(id=item) for item in ids]
+
+
+class DataModelingFilter(Filter):
+    space: str | None = None
+    include_global: bool | None = None
+
+
+class ContainerFilter(DataModelingFilter):
+    used_for: Literal["node", "edge", "record", "all"] | None
+
+
+class ViewFilter(DataModelingFilter):
+    include_inherited_properties: bool | None = None
+    all_versions: bool | None = None
+
+
+class DataModelFilter(DataModelingFilter):
+    inline_views: bool | None = None
+    all_versions: bool | None = None
+
+
+class InstanceFilter(Filter):
+    instance_type: Literal["node", "edge"] | None = None
+    source: ViewReference | None = None
+    space: list[str] | None = None
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        body: dict[str, Any] = {}
+        if self.instance_type is not None:
+            body["instanceType"] = self.instance_type
+        if self.source is not None:
+            body["sources"] = [{"source": self.source.dump()}]
+        if self.space is not None:
+            instance_type = self.instance_type or "node"
+            body["filter"] = {
+                "in": {
+                    "property": [instance_type, "space"],
+                    "values": self.space,
+                }
+            }
+        return body
+
+
+class AnnotationFilter(Filter):
+    annotated_resource_type: Literal["file", "threedmodel"]
+    annotated_resource_ids: list[ExternalId | InternalId]
+    annotation_type: AnnotationType | None = None
+    created_app: str | None = None
+    creating_app_version: str | None = None
+    creating_user: str | None = None
+    status: AnnotationStatus | None = None
+
+
+class ThreeDAssetMappingFilter(Filter): ...
+
+
+class ThreeDAssetMappingAssetIdFilter(ThreeDAssetMappingFilter):
+    asset_ids: list[int] = Field(max_length=100)
+
+
+class ThreeDAssetMappingAssetInstanceIdFilter(ThreeDAssetMappingFilter):
+    asset_instance_ids: list[NodeReference] = Field(max_length=100)
+
+
+class ThreeDAssetMapping3DNodeFilter(ThreeDAssetMappingFilter):
+    node_ids: list[int] = Field(max_length=100)
+
+
+class ThreeDAssetMappingTreeIndexFilter(ThreeDAssetMappingFilter):
+    tree_indexes: list[int] = Field(max_length=100)

cognite_toolkit/_cdf_tk/client/request_classes/graphql.py

@@ -0,0 +1,28 @@
+UPSERT_BODY = """
+mutation UpsertGraphQlDmlVersion($dmCreate: GraphQlDmlVersionUpsert!) {
+    upsertGraphQlDmlVersion(graphQlDmlVersion: $dmCreate) {
+        errors {
+            kind
+            message
+            hint
+            location {
+                start {
+                    line
+                    column
+                }
+            }
+        }
+        result {
+            space
+            externalId
+            version
+            name
+            description
+            graphQlDml
+            isGlobal
+            createdTime
+            lastUpdatedTime
+        }
+    }
+}
+"""

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")