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

cognite_toolkit/_cdf_tk/client/{data_classes/api_classes.py → cdf_client/responses.py}

@@ -2,11 +2,19 @@ from typing import Generic, TypeVar
 
 from pydantic import BaseModel, Field, JsonValue
 
-from cognite_toolkit._cdf_tk.utils.http_client._data_classes2 import RequestResource
-
 T = TypeVar("T", bound=BaseModel)
 
 
+class ResponseItems(BaseModel, Generic[T]):
+    """A page of reference items from a paginated API response.
+
+    Attributes:
+        items: The list of reference items in this page.
+    """
+
+    items: list[T]
+
+
 class PagedResponse(BaseModel, Generic[T]):
     items: list[T]
     next_cursor: str | None = Field(None, alias="nextCursor")
@@ -17,14 +25,3 @@ class QueryResponse(BaseModel, Generic[T]):
     typing: dict[str, JsonValue] | None = None
     next_cursor: dict[str, str] = Field(alias="nextCursor")
     debug: dict[str, JsonValue] | None = None
-
-
-class InternalIdRequest(RequestResource):
-    id: int
-
-    def as_id(self) -> int:
-        return self.id
-
-    @classmethod
-    def from_ids(cls, ids: list[int]) -> list["InternalIdRequest"]:
-        return [cls(id=id_) for id_ in ids]

cognite_toolkit/_cdf_tk/client/data_classes/agent.py

@@ -0,0 +1,127 @@
+from typing import Annotated, Any, Literal
+
+from pydantic import BeforeValidator, Field
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestResource, ResponseResource
+from tests.test_unit.test_cdf_tk.test_tk_warnings.test_warnings_metatest import get_all_subclasses
+
+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."""
+
+    ...
+
+
+KNOWN_TOOLS = {tool.type: tool for tool in get_all_subclasses(AgentToolDefinition) if hasattr(tool, "type")}
+
+
+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 AgentRequest(RequestResource):
+    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]):
+    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:
+        return AgentRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/asset.py

@@ -0,0 +1,54 @@
+from typing import ClassVar, Literal
+
+from pydantic import JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestUpdateable, ResponseResource
+
+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"})
+    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 AssetAggregateItem(BaseModelObject):
+    child_count: int
+    depth: int
+    path: list[dict[Literal["id"], int]]
+
+
+class AssetResponse(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

@@ -1,55 +1,150 @@
 import sys
+import types
 from abc import ABC, abstractmethod
 from collections import UserList
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import Any, ClassVar, Generic, Literal, TypeVar, Union, get_args, get_origin
 
-from pydantic import ConfigDict
+from cognite.client import CogniteClient
+from pydantic import BaseModel, ConfigDict
 from pydantic.alias_generators import to_camel
 
-from cognite_toolkit._cdf_tk.utils.http_client._data_classes2 import BaseModelObject, RequestResource
-
-if TYPE_CHECKING:
-    from cognite.client import CogniteClient
-
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
     from typing_extensions import Self
 
 
-T_RequestResource = TypeVar("T_RequestResource", bound=RequestResource)
+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")
 
-class ResponseResource(BaseModelObject, Generic[T_RequestResource], ABC):
-    @abstractmethod
-    def as_request_resource(self) -> T_RequestResource:
-        """Convert the response resource to a request resource."""
-        ...
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
 
-    def as_write(self) -> T_RequestResource:
-        """Alias for as_request_resource to match protocol signature."""
-        return self.as_request_resource()
+        This is the default serialization method for request resources.
+        """
+        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)
 
 
-class Identifier(RequestResource, ABC):
-    """Base class for all identifier classes."""
+class Identifier(BaseModelObject):
+    """Base class for all identifier objects typically
+    {"externalId": "..."}, {"id": ...}, {"space": "...", "externalId: }."""
 
     model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True, frozen=True)
 
-    def dump(self, camel_case: bool = True, include_type: bool = True) -> dict[str, Any]:
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
         """Dump the resource to a dictionary.
 
         This is the default serialization method for request resources.
         """
-        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=not include_type)
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
+
+    def __str__(self) -> str:
+        raise NotImplementedError()
+
+
+T_Identifier = TypeVar("T_Identifier", bound=Identifier)
+
+
+class RequestResource(BaseModelObject, ABC):
+    @abstractmethod
+    def as_id(self) -> Identifier:
+        raise NotImplementedError()
+
+    def __str__(self) -> str:
+        return str(self.as_id())
+
+
+T_RequestResource = TypeVar("T_RequestResource", bound=RequestResource)
+
+
+class RequestUpdateable(RequestResource, ABC):
+    container_fields: ClassVar[frozenset[str]] = frozenset()
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset()
+
+    def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
+        """Convert the request resource to an update item."""
+        update_item = self.as_id().dump(camel_case=True)
+        update: dict[str, Any] = {}
+        field_by_name = {info.alias or field_id: (field_id, info) for field_id, info in type(self).model_fields.items()}
+        # When mode is "patch", we only include fields that are set
+        exclude_unset = mode == "patch"
+        for key, value in self.model_dump(mode="json", by_alias=True, exclude_unset=exclude_unset).items():
+            if key in update_item:
+                # Skip identifier fields
+                continue
+            if key not in field_by_name:
+                # Skip unknown fields
+                continue
+            field_id, info = field_by_name[key]
+            if field_id in self.container_fields:
+                if mode == "patch":
+                    update[key] = {"add": value}
+                elif mode == "replace":
+                    if value is None:
+                        origin = _get_annotation_origin(info.annotation)
+                        if origin is list:
+                            update[key] = {"set": []}
+                        elif origin is dict:
+                            update[key] = {"set": {}}
+                        else:
+                            raise NotImplementedError(
+                                f'Cannot replace container field "{key}" with None when its type is unknown.'
+                            )
+                    else:
+                        update[key] = {"set": value}
+                else:
+                    raise NotImplementedError(f'Update mode "{mode}" is not supported for container fields.')
+            elif value is None:
+                if field_id not in self.non_nullable_fields:
+                    update[key] = {"setNull": True}
+            else:
+                update[key] = {"set": value}
+        update_item["update"] = update
+        return update_item
+
+
+def _get_annotation_origin(field_type: Any) -> Any:
+    origin = get_origin(field_type)
+    args = get_args(field_type)
+
+    # Check for Union type (both typing.Union and | syntax from Python 3.10+)
+    is_union = origin is Union or isinstance(field_type, getattr(types, "UnionType", ()))
+
+    if is_union:
+        # Handle Optional[T] by filtering out NoneType
+        none_types = (type(None), types.NoneType)
+        non_none_args = [arg for arg in args if arg not in none_types]
+        if len(non_none_args) == 1:
+            field_type = non_none_args[0]
+            origin = get_origin(field_type) or field_type
+    return origin
+
+
+class ResponseResource(BaseModelObject, Generic[T_RequestResource], ABC):
+    @abstractmethod
+    def as_request_resource(self) -> T_RequestResource:
+        """Convert the response resource to a request resource."""
+        raise NotImplementedError()
+
+    def as_write(self) -> T_RequestResource:
+        """Alias for as_request_resource to match protocol signature."""
+        return self.as_request_resource()
 
-    def as_id(self) -> Self:
-        return self
 
+T_ResponseResource = TypeVar("T_ResponseResource", bound=ResponseResource)
 
 T_Resource = TypeVar("T_Resource", bound=RequestResource | ResponseResource)
 
 
+# Todo: Delete this class and use list[T_Resource] directly
 class BaseResourceList(UserList[T_Resource]):
     """Base class for resource lists."""
 

cognite_toolkit/_cdf_tk/client/data_classes/event.py

@@ -0,0 +1,43 @@
+from typing import ClassVar
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import RequestUpdateable, ResponseResource
+
+from .identifiers import ExternalId, InternalOrExternalId
+
+
+class EventRequest(RequestUpdateable):
+    container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "asset_ids"})
+    external_id: str | None = None
+    data_set_id: int | None = None
+    start_time: int | None = None
+    end_time: int | None = None
+    type: str | None = None
+    subtype: str | None = None
+    description: str | None = None
+    metadata: dict[str, str] | None = None
+    asset_ids: list[int] | None = None
+    source: str | None = None
+
+    def as_id(self) -> InternalOrExternalId:
+        if self.external_id is None:
+            raise ValueError("Cannot convert EventRequest to ExternalId when external_id is None")
+        return ExternalId(external_id=self.external_id)
+
+
+class EventResponse(ResponseResource[EventRequest]):
+    external_id: str | None = None
+    data_set_id: int | None = None
+    start_time: int | None = None
+    end_time: int | None = None
+    type: str | None = None
+    subtype: str | None = None
+    description: str | None = None
+    metadata: dict[str, str] | None = None
+    asset_ids: list[int] | None = None
+    source: str | None = None
+    id: int
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> EventRequest:
+        return EventRequest.model_validate(self.dump())

cognite_toolkit/_cdf_tk/client/data_classes/filemetadata.py

@@ -0,0 +1,56 @@
+from typing import ClassVar, Literal
+
+from pydantic import JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import RequestUpdateable, ResponseResource
+
+from .identifiers import ExternalId
+from .instance_api import NodeReference
+
+
+class FileMetadataRequest(RequestUpdateable):
+    container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "labels", "asset_ids", "security_categories"})
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"asset_ids", "security_categories"})
+    external_id: str | None = None
+    name: str
+    directory: str | None = None
+    source: str | None = None
+    mime_type: str | None = None
+    metadata: dict[str, str] | None = None
+    asset_ids: list[int] | None = None
+    data_set_id: int | None = None
+    labels: list[dict[Literal["externalId"], str]] | None = None
+    geo_location: JsonValue | None = None
+    source_created_time: int | None = None
+    source_modified_time: int | None = None
+    security_categories: list[int] | None = None
+
+    def as_id(self) -> ExternalId:
+        if self.external_id is None:
+            raise ValueError("Cannot convert FileMetadataRequest to ExternalId when external_id is None")
+        return ExternalId(external_id=self.external_id)
+
+
+class FileMetadataResponse(ResponseResource[FileMetadataRequest]):
+    created_time: int
+    last_updated_time: int
+    uploaded_time: int | None = None
+    uploaded: bool
+    id: int
+    external_id: str | None = None
+    name: str
+    directory: str | None = None
+    instance_id: NodeReference | None = None
+    source: str | None = None
+    mime_type: str | None = None
+    metadata: dict[str, str] | None = None
+    asset_ids: list[int] | None = None
+    data_set_id: int | None = None
+    labels: list[dict[Literal["externalId"], str]] | None = None
+    geo_location: JsonValue | None = None
+    source_created_time: int | None = None
+    source_modified_time: int | None = None
+    security_categories: list[int] | None = None
+
+    def as_request_resource(self) -> FileMetadataRequest:
+        return FileMetadataRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/identifiers.py

@@ -0,0 +1,44 @@
+from collections.abc import Sequence
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import Identifier
+
+
+class InternalOrExternalIdDefinition(Identifier):
+    type: str
+
+
+class InternalId(InternalOrExternalIdDefinition):
+    type: Literal["id"] = Field("id", exclude=True)
+    id: int
+
+    @classmethod
+    def from_ids(cls, ids: Sequence[int]) -> list["InternalId"]:
+        return [cls(id=id_) for id_ in ids]
+
+    def __str__(self) -> str:
+        return f"id={self.id}"
+
+
+class ExternalId(InternalOrExternalIdDefinition):
+    type: Literal["externalId"] = Field("externalId", exclude=True)
+    external_id: str
+
+    @classmethod
+    def from_external_ids(cls, external_ids: list[str]) -> list["ExternalId"]:
+        return [cls(external_id=ext_id) for ext_id in external_ids]
+
+    def __str__(self) -> str:
+        return f"externalId='{self.external_id}'"
+
+
+InternalOrExternalId = Annotated[InternalId | ExternalId, Field(discriminator="type")]
+
+
+class NameId(Identifier):
+    name: str
+
+    def __str__(self) -> str:
+        return f"name='{self.name}'"

cognite_toolkit/_cdf_tk/client/data_classes/instance_api.py

@@ -14,19 +14,38 @@ class TypedInstanceIdentifier(Identifier):
     space: str
     external_id: str
 
+    def __str__(self) -> str:
+        return f"Instance({self.instance_type}, {self.space}, {self.external_id})"
+
+    def dump(self, camel_case: bool = True, include_type: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
+
+        This is the default serialization method for request resources.
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=not include_type)
+
 
 class TypedNodeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["node"] = "node"
 
+    def __str__(self) -> str:
+        return f"Node({self.space}, {self.external_id})"
+
 
 class TypedEdgeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["edge"] = "edge"
 
+    def __str__(self) -> str:
+        return f"Edge({self.space}, {self.external_id})"
+
 
 class InstanceIdentifier(Identifier):
     space: str
     external_id: str
 
+    def __str__(self) -> str:
+        return f"Instance({self.space}, {self.external_id})"
+
 
 class InstanceResult(BaseModelObject):
     instance_type: InstanceType
@@ -51,6 +70,16 @@ class ViewReference(Identifier):
     external_id: str
     version: str
 
+    def __str__(self) -> str:
+        return f"View({self.space}, {self.external_id}, v{self.version})"
+
+    def dump(self, camel_case: bool = True, include_type: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
+
+        This is the default serialization method for request resources.
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=not include_type)
+
 
 ######################################################
 # The classes below are helper classes for making instances request/responses.
@@ -160,3 +189,8 @@ class InstanceResponseItem(BaseModelObject):
             space=self.space,
             external_id=self.external_id,
         )
+
+
+class NodeReference(BaseModelObject):
+    space: str
+    external_id: str

cognite_toolkit/_cdf_tk/client/data_classes/raw.py

@@ -0,0 +1,43 @@
+import sys
+
+from pydantic import Field
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    Identifier,
+    RequestResource,
+    ResponseResource,
+)
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+class RAWDatabase(RequestResource, Identifier, ResponseResource["RAWDatabase"]):
+    name: str
+
+    def as_id(self) -> Self:
+        return self
+
+    def __str__(self) -> str:
+        return f"name='{self.name}'"
+
+    def as_request_resource(self) -> "RAWDatabase":
+        return type(self).model_validate(self.dump(), extra="ignore")
+
+
+class RAWTable(RequestResource, Identifier, ResponseResource["RAWTable"]):
+    # This is a query parameter, so we exclude it from serialization
+    db_name: str = Field(exclude=True)
+    name: str
+
+    def as_id(self) -> Self:
+        return self
+
+    def __str__(self) -> str:
+        return f"dbName='{self.db_name}', tableName='{self.name}'"
+
+    def as_request_resource(self) -> "RAWTable":
+        dumped = {**self.dump(), "dbName": self.db_name}
+        return type(self).model_validate(dumped, extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/streams.py

@@ -7,6 +7,7 @@ from cognite_toolkit._cdf_tk.protocols import (
 )
 
 from .base import BaseModelObject, BaseResourceList, RequestResource, ResponseResource
+from .identifiers import ExternalId
 
 
 class StreamRequest(RequestResource):
@@ -15,8 +16,8 @@ class StreamRequest(RequestResource):
     external_id: str
     settings: dict[Literal["template"], dict[Literal["name"], StreamTemplateName]]
 
-    def as_id(self) -> str:
-        return self.external_id
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
 
 
 class StreamRequestList(BaseResourceList[StreamRequest], ResourceRequestListProtocol):