diff --git a/cognite/client/_api/ai/__init__.py b/cognite/client/_api/ai/__init__.py index d786696aa..5c5a936fd 100644 --- a/cognite/client/_api/ai/__init__.py +++ b/cognite/client/_api/ai/__init__.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING -from cognite.client._api.ai.tools import ToolsAPI +from cognite.client._api.ai.tools import AIToolsAPI from cognite.client._api_client import APIClient if TYPE_CHECKING: @@ -12,4 +12,4 @@ class AIAPI(APIClient): def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None: super().__init__(config, api_version, cognite_client) - self.tools = ToolsAPI(config, api_version, cognite_client) + self.tools = AIToolsAPI(config, api_version, cognite_client) diff --git a/cognite/client/_api/ai/tools/__init__.py b/cognite/client/_api/ai/tools/__init__.py index a27303a98..9d13810d1 100644 --- a/cognite/client/_api/ai/tools/__init__.py +++ b/cognite/client/_api/ai/tools/__init__.py @@ -2,14 +2,14 @@ from typing import TYPE_CHECKING -from cognite.client._api.ai.tools.documents import DocumentsAPI +from cognite.client._api.ai.tools.documents import AIDocumentsAPI from cognite.client._api_client import APIClient if TYPE_CHECKING: from cognite.client import ClientConfig, CogniteClient -class ToolsAPI(APIClient): +class AIToolsAPI(APIClient): def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None: super().__init__(config, api_version, cognite_client) - self.documents = DocumentsAPI(config, api_version, cognite_client) + self.documents = AIDocumentsAPI(config, api_version, cognite_client) diff --git a/cognite/client/_api/ai/tools/documents.py b/cognite/client/_api/ai/tools/documents.py index 4b7de2ecb..9c24a7ff9 100644 --- a/cognite/client/_api/ai/tools/documents.py +++ b/cognite/client/_api/ai/tools/documents.py @@ -8,7 +8,7 @@ from cognite.client.utils.useful_types import SequenceNotStr -class DocumentsAPI(APIClient): +class AIDocumentsAPI(APIClient): _RESOURCE_PATH = "/ai/tools/documents" def summarize( diff --git a/cognite/client/testing.py b/cognite/client/testing.py index 2a7be50ab..dff2c6aad 100644 --- a/cognite/client/testing.py +++ b/cognite/client/testing.py @@ -6,6 +6,9 @@ from unittest.mock import MagicMock from cognite.client import CogniteClient +from cognite.client._api.ai import AIAPI +from cognite.client._api.ai.tools import AIToolsAPI +from cognite.client._api.ai.tools.documents import AIDocumentsAPI from cognite.client._api.annotations import AnnotationsAPI from cognite.client._api.assets import AssetsAPI from cognite.client._api.data_modeling import DataModelingAPI @@ -95,6 +98,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # - Add spacing above and below # - Use `spec=MyAPI` only for "top level" # - Use `spec_set=MyNestedAPI` for all nested APIs + self.ai = MagicMock(spec=AIAPI) + self.ai.tools = MagicMock(spec=AIToolsAPI) + self.ai.tools.documents = MagicMock(spec_set=AIDocumentsAPI) self.annotations = MagicMock(spec_set=AnnotationsAPI) self.assets = MagicMock(spec_set=AssetsAPI)