feat: introduce Datasets

mostlyai/sdk/_local/routes.py

@@ -37,10 +37,12 @@
     create_zip_in_memory,
     get_model_label,
     read_connector_from_json,
+    read_dataset_from_json,
     read_generator_from_json,
     read_job_progress_from_json,
     read_synthetic_dataset_from_json,
     write_connector_to_json,
+    write_dataset_to_json,
     write_generator_to_json,
     write_synthetic_dataset_to_json,
 )
@@ -56,6 +58,10 @@
     ConnectorType,
     ConnectorWriteDataConfig,
     CurrentUser,
+    Dataset,
+    DatasetConfig,
+    DatasetListItem,
+    DatasetPatchConfig,
     Generator,
     GeneratorCloneConfig,
     GeneratorCloneTrainingStatus,
@@ -289,6 +295,96 @@ async def query(id: str, sql: str = Body(..., embed=True)) -> StreamingResponse:
                 background=BackgroundTask(Path(tmp_path).unlink, missing_ok=True),
             )
 
+        ## DATASETS
+
+        @self.router.get("/datasets")
+        async def list_datasets(
+            offset: int = 0,
+            limit: int = 50,
+            searchTerm: str | None = None,
+        ) -> JSONResponse:
+            dataset_dirs = [p for p in (self.home_dir / "datasets").glob("*") if p.is_dir()]
+            dataset_list_items = []
+            for dataset_dir in dataset_dirs:
+                dataset = read_dataset_from_json(dataset_dir)
+                dataset_string = " ".join([dataset.name or "", dataset.description or ""]).lower()
+                if searchTerm and searchTerm.lower() not in dataset_string:
+                    continue
+                # use model_construct to skip validation and warnings of extra fields
+                dataset_list_items.append(DatasetListItem.model_construct(**dataset.model_dump()))
+
+            return JSONResponse(
+                status_code=200,
+                content=jsonable_encoder(
+                    {
+                        "totalCount": len(dataset_list_items),
+                        "results": dataset_list_items[int(offset) : int(offset) + int(limit)],
+                    }
+                ),
+            )
+
+        @self.router.post("/datasets", response_model=Dataset)
+        async def create_dataset(config: DatasetConfig = Body(...)) -> Dataset:
+            dataset = Dataset(**config.model_dump())
+            dataset_dir = self.home_dir / "datasets" / dataset.id
+            write_dataset_to_json(dataset_dir, dataset)
+            return dataset
+
+        @self.router.get("/datasets/{id}", response_model=Dataset)
+        async def get_dataset(id: str) -> Dataset:
+            dataset_dir = self.home_dir / "datasets" / id
+            if not dataset_dir.exists():
+                raise HTTPException(status_code=404, detail=f"Dataset `{id}` not found")
+            dataset = read_dataset_from_json(dataset_dir)
+            return dataset
+
+        @self.router.patch("/datasets/{id}", response_model=Dataset)
+        async def patch_dataset(id: str, config: DatasetPatchConfig = Body(...)) -> Dataset:
+            dataset_dir = self.home_dir / "datasets" / id
+            dataset = read_dataset_from_json(dataset_dir)
+            for key, value in config.model_dump().items():
+                if value is not None:
+                    setattr(dataset, key, value)
+            write_dataset_to_json(dataset_dir, dataset)
+            return dataset
+
+        @self.router.delete("/datasets/{id}")
+        async def delete_dataset(id: str):
+            dataset_dir = self.home_dir / "datasets" / id
+            shutil.rmtree(dataset_dir, ignore_errors=True)
+
+        @self.router.get("/datasets/{id}/file")
+        async def download_dataset_file(id: str, filepath: str) -> FileResponse:
+            dataset_dir = self.home_dir / "datasets" / id
+            filename = Path(filepath).name
+            return StreamingResponse(
+                open(dataset_dir / filepath, "rb"),
+                media_type="application/octet-stream",
+                headers={"Content-Disposition": f"attachment; filename={filename}"},
+            )
+
+        @self.router.post("/datasets/{id}/file")
+        async def upload_dataset_file(id: str, file: UploadFile = File(...)) -> None:
+            dataset_dir = self.home_dir / "datasets" / id
+            dataset = read_dataset_from_json(dataset_dir)
+            try:
+                file_content = await file.read()
+                with open(dataset_dir / file.filename, "wb") as f:
+                    f.write(file_content)
+                dataset.files.append(file.filename)
+                write_dataset_to_json(dataset_dir, dataset)
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"Error uploading file `{file.filename}`: {e}")
+
+        @self.router.delete("/datasets/{id}/file")
+        async def delete_dataset_file(id: str, filepath: str) -> None:
+            dataset_dir = self.home_dir / "datasets" / id
+            dataset = read_dataset_from_json(dataset_dir)
+            if os.path.exists(dataset_dir / filepath):
+                os.remove(dataset_dir / filepath)
+            dataset.files = [f for f in dataset.files if f != filepath]
+            write_dataset_to_json(dataset_dir, dataset)
+
         ## GENERATORS
 
         @self.router.get("/generators")

mostlyai/sdk/_local/storage.py

@@ -23,6 +23,7 @@
 
 from mostlyai.sdk.domain import (
     Connector,
+    Dataset,
     Generator,
     JobProgress,
     ModelType,
@@ -82,6 +83,17 @@ def write_connector_to_json(connector_dir: Path, connector: Connector) -> None:
     write_to_json(json_file, connector)
 
 
+def read_dataset_from_json(dataset_dir: Path) -> Dataset:
+    json_file = dataset_dir / "dataset.json"
+    return Dataset(**json.loads(json_file.read_text()))
+
+
+def write_dataset_to_json(dataset_dir: Path, dataset: Dataset) -> None:
+    json_file = dataset_dir / "dataset.json"
+    dataset_dir.mkdir(parents=True, exist_ok=True)
+    write_to_json(json_file, dataset)
+
+
 def read_job_progress_from_json(resource_dir: Path) -> JobProgress:
     progress_file = resource_dir / "job_progress.json"
     lock_file = progress_file.with_suffix(".lock")

mostlyai/sdk/client/api.py

@@ -30,6 +30,7 @@
 )
 from mostlyai.sdk.client.base import DEFAULT_BASE_URL, GET, _MostlyBaseClient
 from mostlyai.sdk.client.connectors import _MostlyConnectorsClient
+from mostlyai.sdk.client.datasets import _MostlyDatasetsClient
 from mostlyai.sdk.client.exceptions import APIError
 from mostlyai.sdk.client.generators import _MostlyGeneratorsClient
 from mostlyai.sdk.client.synthetic_datasets import (
@@ -218,6 +219,7 @@ def __init__(
         super().__init__(**client_kwargs)
         self.connectors = _MostlyConnectorsClient(**client_kwargs)
         self.generators = _MostlyGeneratorsClient(**client_kwargs)
+        self.datasets = _MostlyDatasetsClient(**client_kwargs)
         self.synthetic_datasets = _MostlySyntheticDatasetsClient(**client_kwargs)
         self.synthetic_probes = _MostlySyntheticProbesClient(**client_kwargs)
         if mode == "LOCAL":

mostlyai/sdk/client/datasets.py

@@ -0,0 +1,236 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+import rich
+
+from mostlyai.sdk.client.base import (
+    DELETE,
+    GET,
+    PATCH,
+    POST,
+    Paginator,
+    _MostlyBaseClient,
+)
+from mostlyai.sdk.domain import (
+    Dataset,
+    DatasetConfig,
+    DatasetListItem,
+    DatasetPatchConfig,
+)
+
+
+class _MostlyDatasetsClient(_MostlyBaseClient):
+    SECTION = ["datasets"]
+
+    # PUBLIC METHODS #
+
+    def list(
+        self,
+        offset: int = 0,
+        limit: int | None = None,
+        search_term: str | None = None,
+        owner_id: str | list[str] | None = None,
+        visibility: str | list[str] | None = None,
+        created_from: str | None = None,
+        created_to: str | None = None,
+        sort_by: str | list[str] | None = None,
+    ) -> Iterator[DatasetListItem]:
+        """
+        List datasets.
+
+        Args:
+            offset: Offset for the entities in the response.
+            limit: Limit for the number of entities in the response.
+            status: Filter by generation status.
+            search_term: Filter by name or description.
+            owner_id: Filter by owner ID.
+            visibility: Filter by visibility (e.g., PUBLIC, PRIVATE or UNLISTED).
+            created_from: Filter by creation date, not older than this date. Format: YYYY-MM-DD.
+            created_to: Filter by creation date, not younger than this date. Format: YYYY-MM-DD.
+            sort_by: Sort by field. Either RECENCY, NO_OF_LIKES.
+
+        Returns:
+            An iterator over datasets.
+
+        Example for listing all datasets:
+            ```python
+            from mostlyai.sdk import MostlyAI
+            mostly = MostlyAI()
+            for ds in mostly.datasets.list():
+                print(f"Dataset `{ds.name}` ({ds.id})")
+            ```
+
+        Example for searching datasets via key word:
+            ```python
+            from mostlyai.sdk import MostlyAI
+            mostly = MostlyAI()
+            datasets = list(mostly.datasets.list(search_term="census"))
+            print(f"Found {len(datasets)} datasets")
+        """
+        with Paginator(
+            self,
+            DatasetListItem,
+            offset=offset,
+            limit=limit,
+            search_term=search_term,
+            owner_id=owner_id,
+            visibility=visibility,
+            created_from=created_from,
+            created_to=created_to,
+            sort_by=sort_by,
+        ) as paginator:
+            yield from paginator
+
+    def get(self, dataset_id: str) -> Dataset:
+        """
+        Retrieve a dataset by its ID.
+
+        Args:
+            dataset_id: The unique identifier of the dataset.
+
+        Returns:
+            Dataset: The retrieved dataset object.
+
+        Example for retrieving a dataset:
+            ```python
+            from mostlyai.sdk import MostlyAI
+            mostly = MostlyAI()
+            ds = mostly.datasets.get('INSERT_YOUR_DATASET_ID')
+            ds
+            ```
+        """
+        if not isinstance(dataset_id, str) or len(dataset_id) != 36:
+            raise ValueError("The provided dataset_id must be a UUID string")
+        response = self.request(verb=GET, path=[dataset_id], response_type=Dataset)
+        return response
+
+    def create(self, config: DatasetConfig | dict[str, Any]) -> Dataset:
+        """
+        Create a dataset.
+
+        Args:
+            config: Configuration for the dataset.
+
+        Returns:
+            The created dataset object.
+
+        Example for creating a dataset:
+            ```python
+            from mostlyai.sdk import MostlyAI
+            mostly = MostlyAI()
+            ds = mostly.datasets.create(
+                config=DatasetConfig(
+                    name="INSERT_YOUR_DATASET_NAME",
+                    description="INSERT_YOUR_DATASET_DESCRIPTION",
+                    connectors=[
+                        DatasetConnector(
+                            connector_id="INSERT_YOUR_CONNECTOR_ID",
+                            location="INSERT_YOUR_LOCATION",
+                        )
+                    ],
+                )
+            )
+            ```
+        """
+        dataset = self.request(
+            verb=POST,
+            path=[],
+            json=config,
+            response_type=Dataset,
+        )
+        dsid = dataset.id
+        if self.local:
+            rich.print(f"Created dataset [dodger_blue2]{dsid}[/]")
+        else:
+            rich.print(f"Created dataset [link={self.base_url}/d/datasets/{dsid} dodger_blue2 underline]{dsid}[/]")
+        return dataset
+
+    def _update(
+        self,
+        dataset_id: str,
+        config: DatasetPatchConfig | dict[str, Any],
+    ) -> Dataset:
+        response = self.request(
+            verb=PATCH,
+            path=[dataset_id],
+            json=config,
+            response_type=Dataset,
+        )
+        return response
+
+    def _delete(self, dataset_id: str) -> None:
+        response = self.request(verb=DELETE, path=[dataset_id])
+        return response
+
+    def _config(self, dataset_id: str) -> DatasetConfig:
+        response = self.request(verb=GET, path=[dataset_id, "config"], response_type=DatasetConfig)
+        return response
+
+    def _download_file(
+        self,
+        dataset_id: str,
+        file_path: str,
+    ) -> tuple[bytes, str | None]:
+        response = self.request(
+            verb=GET,
+            path=[dataset_id, "file"],
+            params={"filepath": file_path},
+            headers={
+                "Accept": "application/json, text/plain, */*",
+            },
+            raw_response=True,
+        )
+        content_bytes = response.content
+        # Check if 'Content-Disposition' header is present
+        if "Content-Disposition" in response.headers:
+            content_disposition = response.headers["Content-Disposition"]
+            filename = re.findall(r"filename(?:=|\*=UTF-8'')(.+)", content_disposition)[0]
+        else:
+            filename = Path(file_path).name
+        return content_bytes, filename
+
+    def _upload_file(
+        self,
+        dataset_id: str,
+        file_path: str | Path,
+    ) -> None:
+        with open(file_path, "rb") as f:
+            _ = self.request(
+                verb=POST,
+                path=[dataset_id, "file"],
+                headers={
+                    "Accept": "application/json, text/plain, */*",
+                },
+                files={"file": f},
+            )
+            rich.print(f"Uploaded file `{file_path}` to Dataset [dodger_blue2]{dataset_id}[/]")
+
+    def _delete_file(
+        self,
+        dataset_id: str,
+        file_path: str | Path,
+    ) -> None:
+        _ = self.request(
+            verb=DELETE,
+            path=[dataset_id, "file"],
+            params={
+                "filepath": file_path,
+            },
+        )
+        rich.print(f"Deleted file `{file_path}` from Dataset [dodger_blue2]{dataset_id}[/]")