From 1ce8af15d6b447714162121953ab022b1899a400 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 12:10:12 +0200 Subject: [PATCH 1/6] sync API --- mostlyai/sdk/domain.py | 70 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/mostlyai/sdk/domain.py b/mostlyai/sdk/domain.py index 20249198..7e050f86 100644 --- a/mostlyai/sdk/domain.py +++ b/mostlyai/sdk/domain.py @@ -13,7 +13,7 @@ # limitations under the License. # generated by datamodel-codegen: -# timestamp: 2025-06-26T08:04:20+00:00 +# timestamp: 2025-07-10T10:09:28+00:00 from __future__ import annotations @@ -175,6 +175,19 @@ class PaginatedTotalCount(RootModel[int]): root: int = Field(..., description="The total number of entities within the list") +class DatasetUsage(CustomBaseModel): + """ + Usage statistics of a dataset. + """ + + no_of_likes: int | None = Field(None, alias="noOfLikes", description="Number of likes of this dataset.") + + +class DatasetSortField(str, Enum): + recency = "RECENCY" + no_of_likes = "NO_OF_LIKES" + + class ConnectorAccessType(str, Enum): """ The access permissions of a connector. @@ -889,6 +902,7 @@ class AssistantMessageContentType(str, Enum): artifact_generator = "artifact/generator" artifact_synthetic_dataset = "artifact/synthetic-dataset" artifact_connector = "artifact/connector" + artifact_dataset = "artifact/dataset" class AssistantMessageDeltaContentType(str, Enum): @@ -1794,6 +1808,60 @@ class Metadata(CustomBaseModel): ) +class DatasetConnector(CustomBaseModel): + """ + Configuration for a dataset connector. + """ + + connector_id: str | None = Field(None, alias="connectorId", description="The unique identifier of a connector.") + locations: list[str] | None = None + + +class Dataset(CustomBaseModel): + """ + A dataset to be consumed via the assistant. + """ + + id: str = Field(..., description="The unique identifier of a dataset.") + name: str | None = Field(None, description="The name of a dataset.") + description: str | None = Field(None, description="The description of / instructions for a dataset.") + connectors: list[DatasetConnector] | None = None + files: list[str] | None = None + usage: DatasetUsage | None = None + metadata: Metadata | None = None + + +class DatasetListItem(CustomBaseModel): + """ + Essential dataset details for listings. + """ + + id: str = Field(..., description="The unique identifier of a dataset.") + name: str | None = Field(None, description="The name of a dataset.") + usage: DatasetUsage | None = None + metadata: Metadata | None = None + + +class DatasetConfig(CustomBaseModel): + """ + The configuration for creating a dataset. + """ + + name: str | None = Field(None, description="The name of a dataset.") + description: str | None = Field(None, description="The description of / instructions for a dataset.") + connectors: list[DatasetConnector] | None = None + + +class DatasetPatchConfig(CustomBaseModel): + """ + The configuration for updating a dataset. + """ + + name: str | None = Field(None, description="The name of a dataset.") + description: str | None = Field(None, description="The description of / instructions for a dataset.") + connectors: list[DatasetConnector] | None = None + + class ConnectorListItem(CustomBaseModel): """ Essential connector details for listings. From 561b8e34150aa2d98a8abfa960ab9d0c4a8a1a82 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 12:33:39 +0200 Subject: [PATCH 2/6] draft --- mostlyai/sdk/_local/datasets.py | 26 +++ mostlyai/sdk/_local/routes.py | 81 +++++++- mostlyai/sdk/_local/storage.py | 12 ++ mostlyai/sdk/client/api.py | 2 + mostlyai/sdk/client/datasets.py | 181 ++++++++++++++++++ mostlyai/sdk/domain.py | 42 +++- .../pydantic_v2/BaseModel.jinja2 | 42 ++++ tools/model.py | 45 +++++ 8 files changed, 429 insertions(+), 2 deletions(-) create mode 100644 mostlyai/sdk/_local/datasets.py create mode 100644 mostlyai/sdk/client/datasets.py diff --git a/mostlyai/sdk/_local/datasets.py b/mostlyai/sdk/_local/datasets.py new file mode 100644 index 00000000..511a94c0 --- /dev/null +++ b/mostlyai/sdk/_local/datasets.py @@ -0,0 +1,26 @@ +# 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. + +from pathlib import Path + +from mostlyai.sdk._local.storage import write_dataset_to_json +from mostlyai.sdk.domain import Dataset, DatasetConfig + + +def create_dataset(home_dir: Path, config: DatasetConfig) -> Dataset: + dataset = Dataset(**config.model_dump()) + dataset_dir = home_dir / "datasets" / dataset.id + write_dataset_to_json(dataset_dir, dataset) + + return dataset diff --git a/mostlyai/sdk/_local/routes.py b/mostlyai/sdk/_local/routes.py index cd5255dd..11bd8943 100644 --- a/mostlyai/sdk/_local/routes.py +++ b/mostlyai/sdk/_local/routes.py @@ -29,7 +29,7 @@ from mostlyai import sdk from mostlyai.sdk._data.conversions import create_container_from_connector from mostlyai.sdk._data.file.utils import read_data_table_from_path -from mostlyai.sdk._local import connectors, generators, synthetic_datasets +from mostlyai.sdk._local import connectors, datasets, generators, synthetic_datasets from mostlyai.sdk._local.execution.jobs import execute_probing_job from mostlyai.sdk._local.generators import create_generator as create_generator_model from mostlyai.sdk._local.storage import ( @@ -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,79 @@ 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, + ownerId: str | list[str] | None = None, + visibility: str | list[str] | None = None, + createdFrom: str | None = None, + createdTo: str | None = None, + sortBy: str | list[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 + if ( + ownerId + and dataset.metadata + and dataset.metadata.owner_id not in ([ownerId] if isinstance(ownerId, str) else ownerId) + ): + continue + if ( + visibility + and dataset.metadata + and dataset.metadata.visibility not in ([visibility] if isinstance(visibility, str) else visibility) + ): + 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 = datasets.create_dataset(self.home_dir, config) + 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) + ## GENERATORS @self.router.get("/generators") diff --git a/mostlyai/sdk/_local/storage.py b/mostlyai/sdk/_local/storage.py index 7d1ff024..2d0c1861 100644 --- a/mostlyai/sdk/_local/storage.py +++ b/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") diff --git a/mostlyai/sdk/client/api.py b/mostlyai/sdk/client/api.py index adec657c..5a7df921 100644 --- a/mostlyai/sdk/client/api.py +++ b/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": diff --git a/mostlyai/sdk/client/datasets.py b/mostlyai/sdk/client/datasets.py new file mode 100644 index 00000000..0cbf0b29 --- /dev/null +++ b/mostlyai/sdk/client/datasets.py @@ -0,0 +1,181 @@ +# 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. + +from collections.abc import Iterator +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 diff --git a/mostlyai/sdk/domain.py b/mostlyai/sdk/domain.py index 7e050f86..28823061 100644 --- a/mostlyai/sdk/domain.py +++ b/mostlyai/sdk/domain.py @@ -13,7 +13,7 @@ # limitations under the License. # generated by datamodel-codegen: -# timestamp: 2025-07-10T10:09:28+00:00 +# timestamp: 2025-07-10T11:59:57+00:00 from __future__ import annotations @@ -1829,6 +1829,46 @@ class Dataset(CustomBaseModel): files: list[str] | None = None usage: DatasetUsage | None = None metadata: Metadata | None = None + OPEN_URL_PARTS: ClassVar[list] = ["d", "datasets"] + + @model_validator(mode="before") + @classmethod + def add_required_fields(cls, values): + if isinstance(values, dict): + if "id" not in values: + values["id"] = str(uuid.uuid4()) + return values + + def update( + self, + name: str | None = None, + description: str | None = None, + connectors: list[DatasetConnector] | None = None, + ) -> None: + """ + Update a dataset with specific parameters. + + Args: + name (str | None): The name of the connector. + description (str | None): The description of the connector. + connectors (list[DatasetConnector] | None): The connectors of the dataset. + """ + patch_config = DatasetPatchConfig( + name=name, + description=description, + connectors=connectors, + ) + self.client._update( + dataset_id=self.id, + config=patch_config, + ) + self.reload() + + def delete(self) -> None: + """ + Delete the dataset. + """ + return self.client._delete(dataset_id=self.id) class DatasetListItem(CustomBaseModel): diff --git a/tools/custom_template/pydantic_v2/BaseModel.jinja2 b/tools/custom_template/pydantic_v2/BaseModel.jinja2 index 65310c32..b3ad1601 100644 --- a/tools/custom_template/pydantic_v2/BaseModel.jinja2 +++ b/tools/custom_template/pydantic_v2/BaseModel.jinja2 @@ -1293,3 +1293,45 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme values.differential_privacy.value_protection_epsilon = 1.0 return values {%- endif %} +{%- if class_name == "Dataset" %} + OPEN_URL_PARTS: ClassVar[list] = ["d", "datasets"] + + @model_validator(mode="before") + @classmethod + def add_required_fields(cls, values): + if isinstance(values, dict): + if "id" not in values: + values["id"] = str(uuid.uuid4()) + return values + + def update( + self, + name: str | None = None, + description: str | None = None, + connectors: list[DatasetConnector] | None = None, + ) -> None: + """ + Update a dataset with specific parameters. + + Args: + name (str | None): The name of the connector. + description (str | None): The description of the connector. + connectors (list[DatasetConnector] | None): The connectors of the dataset. + """ + patch_config = DatasetPatchConfig( + name=name, + description=description, + connectors=connectors, + ) + self.client._update( + dataset_id=self.id, + config=patch_config, + ) + self.reload() + + def delete(self) -> None: + """ + Delete the dataset. + """ + return self.client._delete(dataset_id=self.id) +{%- endif %} diff --git a/tools/model.py b/tools/model.py index 72757f55..d71e485e 100644 --- a/tools/model.py +++ b/tools/model.py @@ -25,6 +25,8 @@ from mostlyai.sdk.domain import ( ConnectorAccessType, ConnectorPatchConfig, + DatasetConnector, + DatasetPatchConfig, Generator, GeneratorConfig, GeneratorPatchConfig, @@ -1215,3 +1217,46 @@ def delegated_method(*args, **kwargs): return delegated_method return object.__getattribute__(self, item) + + +class Dataset: + OPEN_URL_PARTS: ClassVar[list] = ["d", "datasets"] + + @model_validator(mode="before") + @classmethod + def add_required_fields(cls, values): + if isinstance(values, dict): + if "id" not in values: + values["id"] = str(uuid.uuid4()) + return values + + def update( + self, + name: str | None = None, + description: str | None = None, + connectors: list[DatasetConnector] | None = None, + ) -> None: + """ + Update a dataset with specific parameters. + + Args: + name (str | None): The name of the connector. + description (str | None): The description of the connector. + connectors (list[DatasetConnector] | None): The connectors of the dataset. + """ + patch_config = DatasetPatchConfig( + name=name, + description=description, + connectors=connectors, + ) + self.client._update( + dataset_id=self.id, + config=patch_config, + ) + self.reload() + + def delete(self) -> None: + """ + Delete the dataset. + """ + return self.client._delete(dataset_id=self.id) From 16a8bbe0f390540884e1ae124ba906377253e893 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 16:24:56 +0200 Subject: [PATCH 3/6] upload/download/delete file --- mostlyai/sdk/client/datasets.py | 55 +++++++++++++++++++ mostlyai/sdk/domain.py | 43 ++++++++++++++- .../pydantic_v2/BaseModel.jinja2 | 41 ++++++++++++++ tools/model.py | 41 ++++++++++++++ 4 files changed, 179 insertions(+), 1 deletion(-) diff --git a/mostlyai/sdk/client/datasets.py b/mostlyai/sdk/client/datasets.py index 0cbf0b29..76964fab 100644 --- a/mostlyai/sdk/client/datasets.py +++ b/mostlyai/sdk/client/datasets.py @@ -12,7 +12,9 @@ # 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 @@ -179,3 +181,56 @@ def _delete(self, dataset_id: str) -> None: 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}[/]") diff --git a/mostlyai/sdk/domain.py b/mostlyai/sdk/domain.py index 28823061..a0baa20c 100644 --- a/mostlyai/sdk/domain.py +++ b/mostlyai/sdk/domain.py @@ -13,7 +13,7 @@ # limitations under the License. # generated by datamodel-codegen: -# timestamp: 2025-07-10T11:59:57+00:00 +# timestamp: 2025-07-10T14:23:10+00:00 from __future__ import annotations @@ -1870,6 +1870,47 @@ def delete(self) -> None: """ return self.client._delete(dataset_id=self.id) + def download_file( + self, + dataset_file_path: str | Path, + output_file_path: str | Path | None = None, + ) -> Path: + """ + Download the dataset file. + + Args: + file_path (str | Path | None): The file path to save the dataset file. + + Returns: + Path: The path to the saved file. + """ + bytes, filename = self.client._download_file(dataset_id=self.id, file_path=str(dataset_file_path)) + output_file_path = Path(output_file_path or ".") + if output_file_path.is_dir(): + output_file_path = output_file_path / filename + output_file_path.write_bytes(bytes) + return output_file_path + + def upload_file( + self, + file_path: str | Path, + ) -> None: + """ + Upload the dataset file. + """ + self.client._upload_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() + + def delete_file( + self, + file_path: str | Path, + ) -> None: + """ + Delete the dataset file. + """ + self.client._delete_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() + class DatasetListItem(CustomBaseModel): """ diff --git a/tools/custom_template/pydantic_v2/BaseModel.jinja2 b/tools/custom_template/pydantic_v2/BaseModel.jinja2 index b3ad1601..bfa563df 100644 --- a/tools/custom_template/pydantic_v2/BaseModel.jinja2 +++ b/tools/custom_template/pydantic_v2/BaseModel.jinja2 @@ -1334,4 +1334,45 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme Delete the dataset. """ return self.client._delete(dataset_id=self.id) + + def download_file( + self, + dataset_file_path: str | Path, + output_file_path: str | Path | None = None, + ) -> Path: + """ + Download the dataset file. + + Args: + file_path (str | Path | None): The file path to save the dataset file. + + Returns: + Path: The path to the saved file. + """ + bytes, filename = self.client._download_file(dataset_id=self.id, file_path=str(dataset_file_path)) + output_file_path = Path(output_file_path or ".") + if output_file_path.is_dir(): + output_file_path = output_file_path / filename + output_file_path.write_bytes(bytes) + return output_file_path + + def upload_file( + self, + file_path: str | Path, + ) -> None: + """ + Upload the dataset file. + """ + self.client._upload_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() + + def delete_file( + self, + file_path: str | Path, + ) -> None: + """ + Delete the dataset file. + """ + self.client._delete_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() {%- endif %} diff --git a/tools/model.py b/tools/model.py index d71e485e..42693795 100644 --- a/tools/model.py +++ b/tools/model.py @@ -1260,3 +1260,44 @@ def delete(self) -> None: Delete the dataset. """ return self.client._delete(dataset_id=self.id) + + def download_file( + self, + dataset_file_path: str | Path, + output_file_path: str | Path | None = None, + ) -> Path: + """ + Download the dataset file. + + Args: + file_path (str | Path | None): The file path to save the dataset file. + + Returns: + Path: The path to the saved file. + """ + bytes, filename = self.client._download_file(dataset_id=self.id, file_path=str(dataset_file_path)) + output_file_path = Path(output_file_path or ".") + if output_file_path.is_dir(): + output_file_path = output_file_path / filename + output_file_path.write_bytes(bytes) + return output_file_path + + def upload_file( + self, + file_path: str | Path, + ) -> None: + """ + Upload the dataset file. + """ + self.client._upload_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() + + def delete_file( + self, + file_path: str | Path, + ) -> None: + """ + Delete the dataset file. + """ + self.client._delete_file(dataset_id=self.id, file_path=str(file_path)) + self.reload() From f55262f1b6f39c24e7ad8ae878ec12363605a41c Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 17:05:29 +0200 Subject: [PATCH 4/6] upload/download/delete file for local mode --- mostlyai/sdk/_local/routes.py | 32 +++++++++++++++++++ mostlyai/sdk/domain.py | 9 +++++- .../pydantic_v2/BaseModel.jinja2 | 7 ++++ tools/model.py | 7 ++++ 4 files changed, 54 insertions(+), 1 deletion(-) diff --git a/mostlyai/sdk/_local/routes.py b/mostlyai/sdk/_local/routes.py index 11bd8943..f5a5b4bd 100644 --- a/mostlyai/sdk/_local/routes.py +++ b/mostlyai/sdk/_local/routes.py @@ -368,6 +368,38 @@ 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") diff --git a/mostlyai/sdk/domain.py b/mostlyai/sdk/domain.py index a0baa20c..2718e8aa 100644 --- a/mostlyai/sdk/domain.py +++ b/mostlyai/sdk/domain.py @@ -13,7 +13,7 @@ # limitations under the License. # generated by datamodel-codegen: -# timestamp: 2025-07-10T14:23:10+00:00 +# timestamp: 2025-07-10T14:44:02+00:00 from __future__ import annotations @@ -1839,6 +1839,13 @@ def add_required_fields(cls, values): values["id"] = str(uuid.uuid4()) return values + @field_validator("files", mode="after") + @classmethod + def initialize_file_list(cls, values): + if values is None: + values = [] + return values + def update( self, name: str | None = None, diff --git a/tools/custom_template/pydantic_v2/BaseModel.jinja2 b/tools/custom_template/pydantic_v2/BaseModel.jinja2 index bfa563df..64958290 100644 --- a/tools/custom_template/pydantic_v2/BaseModel.jinja2 +++ b/tools/custom_template/pydantic_v2/BaseModel.jinja2 @@ -1304,6 +1304,13 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comme values["id"] = str(uuid.uuid4()) return values + @field_validator("files", mode="after") + @classmethod + def initialize_file_list(cls, values): + if values is None: + values = [] + return values + def update( self, name: str | None = None, diff --git a/tools/model.py b/tools/model.py index 42693795..f050247f 100644 --- a/tools/model.py +++ b/tools/model.py @@ -1230,6 +1230,13 @@ def add_required_fields(cls, values): values["id"] = str(uuid.uuid4()) return values + @field_validator("files", mode="after") + @classmethod + def initialize_file_list(cls, values): + if values is None: + values = [] + return values + def update( self, name: str | None = None, From 62c1edb5f3342b6da0dfaa5994c1204917824008 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 17:08:34 +0200 Subject: [PATCH 5/6] refine --- mostlyai/sdk/_local/datasets.py | 26 -------------------------- mostlyai/sdk/_local/routes.py | 6 ++++-- 2 files changed, 4 insertions(+), 28 deletions(-) delete mode 100644 mostlyai/sdk/_local/datasets.py diff --git a/mostlyai/sdk/_local/datasets.py b/mostlyai/sdk/_local/datasets.py deleted file mode 100644 index 511a94c0..00000000 --- a/mostlyai/sdk/_local/datasets.py +++ /dev/null @@ -1,26 +0,0 @@ -# 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. - -from pathlib import Path - -from mostlyai.sdk._local.storage import write_dataset_to_json -from mostlyai.sdk.domain import Dataset, DatasetConfig - - -def create_dataset(home_dir: Path, config: DatasetConfig) -> Dataset: - dataset = Dataset(**config.model_dump()) - dataset_dir = home_dir / "datasets" / dataset.id - write_dataset_to_json(dataset_dir, dataset) - - return dataset diff --git a/mostlyai/sdk/_local/routes.py b/mostlyai/sdk/_local/routes.py index f5a5b4bd..38c23611 100644 --- a/mostlyai/sdk/_local/routes.py +++ b/mostlyai/sdk/_local/routes.py @@ -29,7 +29,7 @@ from mostlyai import sdk from mostlyai.sdk._data.conversions import create_container_from_connector from mostlyai.sdk._data.file.utils import read_data_table_from_path -from mostlyai.sdk._local import connectors, datasets, generators, synthetic_datasets +from mostlyai.sdk._local import connectors, generators, synthetic_datasets from mostlyai.sdk._local.execution.jobs import execute_probing_job from mostlyai.sdk._local.generators import create_generator as create_generator_model from mostlyai.sdk._local.storage import ( @@ -342,7 +342,9 @@ async def list_datasets( @self.router.post("/datasets", response_model=Dataset) async def create_dataset(config: DatasetConfig = Body(...)) -> Dataset: - dataset = datasets.create_dataset(self.home_dir, config) + 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) From 3b6e6918723fde963ca4ab212f55917fd681bea4 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Thu, 10 Jul 2025 17:18:58 +0200 Subject: [PATCH 6/6] refine --- mostlyai/sdk/_local/routes.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/mostlyai/sdk/_local/routes.py b/mostlyai/sdk/_local/routes.py index 38c23611..7c9d4c2d 100644 --- a/mostlyai/sdk/_local/routes.py +++ b/mostlyai/sdk/_local/routes.py @@ -302,11 +302,6 @@ async def list_datasets( offset: int = 0, limit: int = 50, searchTerm: str | None = None, - ownerId: str | list[str] | None = None, - visibility: str | list[str] | None = None, - createdFrom: str | None = None, - createdTo: str | None = None, - sortBy: str | list[str] | None = None, ) -> JSONResponse: dataset_dirs = [p for p in (self.home_dir / "datasets").glob("*") if p.is_dir()] dataset_list_items = [] @@ -315,18 +310,6 @@ async def list_datasets( dataset_string = " ".join([dataset.name or "", dataset.description or ""]).lower() if searchTerm and searchTerm.lower() not in dataset_string: continue - if ( - ownerId - and dataset.metadata - and dataset.metadata.owner_id not in ([ownerId] if isinstance(ownerId, str) else ownerId) - ): - continue - if ( - visibility - and dataset.metadata - and dataset.metadata.visibility not in ([visibility] if isinstance(visibility, str) else visibility) - ): - continue # use model_construct to skip validation and warnings of extra fields dataset_list_items.append(DatasetListItem.model_construct(**dataset.model_dump()))