implement datasets sdk

basalt/_version.py

@@ -1 +1 @@
-__version__ = "0.2.1"
+__version__ = "0.2.2"

basalt/basalt_facade.py

@@ -1,7 +1,8 @@
 from .utils.api import Api
-from .utils.protocols import IPromptSDK, IBasaltSDK, LogLevel
+from .utils.protocols import IPromptSDK, IBasaltSDK, LogLevel, IDatasetSDK
 from .sdk.promptsdk import PromptSDK
 from .sdk.monitorsdk import MonitorSDK
+from .sdk.datasetsdk import DatasetSDK
 from .basaltsdk import BasaltSDK
 from .utils.memcache import MemoryCache
 from .utils.networker import Networker
@@ -40,8 +41,9 @@ def __init__(self, api_key: str, log_level: LogLevel = 'all'):
 
         prompt = PromptSDK(api, cache, global_fallback_cache, logger)
         monitor = MonitorSDK(api, logger)
+        datasets = DatasetSDK(api, logger)
 
-        self._basalt = BasaltSDK(prompt, monitor)
+        self._basalt = BasaltSDK(prompt, monitor, datasets)
 
     @property
     def prompt(self) -> IPromptSDK:
@@ -56,3 +58,10 @@ def monitor(self) -> IMonitorSDK:
         Read-only access to the MonitorSDK instance.
         """
         return self._basalt.monitor
+
+    @property
+    def datasets(self) -> IDatasetSDK:
+        """
+        Read-only access to the DatasetSDK instance.
+        """
+        return self._basalt.datasets

basalt/basaltsdk.py

@@ -1,4 +1,4 @@
-from .utils.protocols import IPromptSDK, IBasaltSDK
+from .utils.protocols import IPromptSDK, IBasaltSDK, IDatasetSDK
 from .ressources.monitor.monitorsdk_types import IMonitorSDK
 
 class BasaltSDK(IBasaltSDK):
@@ -7,9 +7,10 @@ class BasaltSDK(IBasaltSDK):
     It serves as the main entry point for interacting with the Basalt SDK.
     """
 
-    def __init__(self, prompt_sdk: IPromptSDK, monitor_sdk: IMonitorSDK):
+    def __init__(self, prompt_sdk: IPromptSDK, monitor_sdk: IMonitorSDK, dataset_sdk: IDatasetSDK):
         self._prompt = prompt_sdk
         self._monitor = monitor_sdk
+        self._datasets = dataset_sdk
 
     @property
     def prompt(self) -> IPromptSDK:
@@ -20,3 +21,8 @@ def prompt(self) -> IPromptSDK:
     def monitor(self) -> IMonitorSDK:
         """Read-only access to the MonitorSDK instance"""
         return self._monitor
+
+    @property
+    def datasets(self) -> IDatasetSDK:
+        """Read-only access to the DatasetSDK instance"""
+        return self._datasets

basalt/endpoints/create_dataset_item.py

@@ -0,0 +1,87 @@
+"""
+Endpoint for creating a new dataset item
+"""
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+from ..utils.dtos import DatasetRowDTO, CreateDatasetItemDTO
+
+@dataclass
+class CreateDatasetItemEndpointResponse:
+    """
+    Response from the create dataset item endpoint
+    """
+    datasetRow: DatasetRowDTO
+    warning: Optional[str] = None
+    error: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CreateDatasetItemEndpointResponse":
+        """
+        Create an instance of CreateDatasetItemEndpointResponse from a dictionary.
+
+        Args:
+            data (Dict[str, Any]): The dictionary containing the response data.
+
+        Returns:
+            CreateDatasetItemEndpointResponse
+        """
+        if "error" in data:
+            return cls(datasetRow=None, error=data["error"])
+
+        return cls(
+            datasetRow=DatasetRowDTO.from_dict(data["datasetRow"]),
+            warning=data.get("warning"),
+            error=None
+        )
+
+
+class CreateDatasetItemEndpoint:
+    """
+    Endpoint class for creating a dataset item.
+    """
+    @staticmethod
+    def prepare_request(dto: CreateDatasetItemDTO) -> Dict[str, Any]:
+        """
+        Prepare the request dictionary for the CreateDatasetItem endpoint.
+
+        Args:
+            dto (CreateDatasetItemDTO): The DTO containing dataset item data.
+
+        Returns:
+            The path, method, and body for creating a dataset item on the API.
+        """
+        body = {
+            "values": dto.values
+        }
+
+        if dto.name:
+            body["name"] = dto.name
+
+        if dto.idealOutput:
+            body["idealOutput"] = dto.idealOutput
+
+        if dto.metadata:
+            body["metadata"] = dto.metadata
+
+        return {
+            "path": f"/datasets/{dto.slug}/items",
+            "method": "POST",
+            "body": body
+        }
+
+    @staticmethod
+    def decode_response(response: dict) -> Tuple[Optional[Exception], Optional[CreateDatasetItemEndpointResponse]]:
+        """
+        Decode the response returned from the API
+
+        Args:
+            response (dict): The JSON response to encode into a CreateDatasetItemEndpointResponse
+
+        Returns:
+            A tuple containing an optional exception and an optional CreateDatasetItemEndpointResponse.
+        """
+        try:
+            return None, CreateDatasetItemEndpointResponse.from_dict(response)
+        except Exception as e:
+            return e, None

basalt/endpoints/get_dataset.py

@@ -0,0 +1,72 @@
+"""
+Endpoint for fetching a specific dataset by slug
+"""
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+from ..utils.dtos import DatasetDTO, GetDatasetDTO
+
+@dataclass
+class GetDatasetEndpointResponse:
+    """
+    Response from the get dataset endpoint
+    """
+    dataset: DatasetDTO
+    error: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GetDatasetEndpointResponse":
+        """
+        Create an instance of GetDatasetEndpointResponse from a dictionary.
+
+        Args:
+            data (Dict[str, Any]): The dictionary containing the response data.
+
+        Returns:
+            GetDatasetEndpointResponse
+        """
+        if "error" in data:
+            return cls(dataset=None, error=data["error"])
+
+        return cls(
+            dataset=DatasetDTO.from_dict(data["dataset"]),
+            error=None
+        )
+
+
+class GetDatasetEndpoint:
+    """
+    Endpoint class for fetching a specific dataset.
+    """
+    @staticmethod
+    def prepare_request(dto: GetDatasetDTO) -> Dict[str, Any]:
+        """
+        Prepare the request dictionary for the GetDataset endpoint.
+
+        Args:
+            dto (GetDatasetDTO): The DTO containing dataset slug.
+
+        Returns:
+            The path, method, and query parameters for getting a dataset on the API.
+        """
+        return {
+            "path": f"/datasets/{dto.slug}",
+            "method": "GET",
+            "query": {}
+        }
+
+    @staticmethod
+    def decode_response(response: dict) -> Tuple[Optional[Exception], Optional[GetDatasetEndpointResponse]]:
+        """
+        Decode the response returned from the API
+
+        Args:
+            response (dict): The JSON response to encode into a GetDatasetEndpointResponse
+
+        Returns:
+            A tuple containing an optional exception and an optional GetDatasetEndpointResponse.
+        """
+        try:
+            return None, GetDatasetEndpointResponse.from_dict(response)
+        except Exception as e:
+            return e, None

basalt/endpoints/list_datasets.py

@@ -0,0 +1,64 @@
+"""
+Endpoint for listing all datasets
+"""
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from ..utils.dtos import DatasetDTO, ListDatasetsDTO
+
+@dataclass
+class ListDatasetsEndpointResponse:
+    """
+    Response from the list datasets endpoint
+    """
+    datasets: List[DatasetDTO]
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ListDatasetsEndpointResponse":
+        """
+        Create an instance of ListDatasetsEndpointResponse from a dictionary.
+
+        Args:
+            data (Dict[str, Any]): The dictionary containing the response data.
+
+        Returns:
+            ListDatasetsEndpointResponse
+        """
+        return cls(
+            datasets=[DatasetDTO.from_dict(dataset) for dataset in data["datasets"]],
+        )
+
+
+class ListDatasetsEndpoint:
+    """
+    Endpoint class for fetching all datasets.
+    """
+    @staticmethod
+    def prepare_request(dto: ListDatasetsDTO) -> Dict[str, Any]:
+        """
+        Prepare the request dictionary for the ListDatasets endpoint.
+
+        Returns:
+            The path, method, and query parameters for getting datasets on the API.
+        """
+        return {
+            "path": "/datasets",
+            "method": "GET",
+            "query": {}
+        }
+
+    @staticmethod
+    def decode_response(response: dict) -> Tuple[Optional[Exception], Optional[ListDatasetsEndpointResponse]]:
+        """
+        Decode the response returned from the API
+
+        Args:
+            response (dict): The JSON response to encode into a ListDatasetsEndpointResponse
+
+        Returns:
+            A tuple containing an optional exception and an optional ListDatasetsEndpointResponse.
+        """
+        try:
+            return None, ListDatasetsEndpointResponse.from_dict(response)
+        except Exception as e:
+            return e, None

basalt/objects/dataset.py

@@ -0,0 +1,92 @@
+"""
+Dataset object for Basalt SDK
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional
+
+
+@dataclass
+class DatasetRow:
+    """
+    A row in a dataset with values and metadata
+    """
+    values: Dict[str, str] 
+    name: Optional[str] = None
+    ideal_output: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the DatasetRow to a dictionary for API requests"""
+        result = {
+            "values": self.values,
+            "metadata": self.metadata,
+						"name": self.name,
+						"idealOutput": self.ideal_output
+        }
+
+        # if self.name:
+        #     result["name"] = self.name
+
+        # if self.ideal_output:
+        #     result["idealOutput"] = self.ideal_output
+
+        return result
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "DatasetRow":
+        """
+        Create a DatasetRow instance from a dictionary
+
+        Args:
+            data: Dictionary containing dataset row data
+
+        Returns:
+            DatasetRow: A new DatasetRow instance
+        """
+        return cls(
+            values=data.get("values", {}),
+            name=data.get("name", None),
+            ideal_output=data.get("idealOutput", None),
+            metadata=data.get("metadata", {})
+        )
+
+
+@dataclass
+class Dataset:
+    """
+    A dataset with rows and metadata
+    """
+    slug: str
+    name: str
+    columns: List[str]
+    rows: List[DatasetRow] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the Dataset to a dictionary for API responses"""
+        return {
+            "slug": self.slug,
+            "name": self.name,
+            "columns": self.columns,
+            "rows": [row.to_dict() for row in self.rows]
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Dataset":
+        """
+        Create a Dataset instance from a dictionary
+
+        Args:
+            data: Dictionary containing dataset data
+
+        Returns:
+            Dataset: A new Dataset instance
+        """
+        rows = [DatasetRow.from_dict(row) for row in data.get("rows", [])]
+
+        return cls(
+            slug=data["slug"],
+            name=data["name"],
+            columns=data["columns"],
+            rows=rows
+        )

basalt/ressources/datasets/__init__.py

@@ -0,0 +1,3 @@
+"""
+Datasets resource module
+"""