add support for async methods

basalt/basaltsdk.py

@@ -5,6 +5,7 @@ class BasaltSDK(IBasaltSDK):
     """
     The BasaltSDK class implements the IBasaltSDK interface.
     It serves as the main entry point for interacting with the Basalt SDK.
+
     """
 
     def __init__(self, prompt_sdk: IPromptSDK, monitor_sdk: IMonitorSDK, dataset_sdk: IDatasetSDK):

basalt/objects/dataset.py

@@ -21,15 +21,9 @@ def to_dict(self) -> Dict[str, Any]:
         result = {
             "values": self.values,
             "metadata": self.metadata,
-						"name": self.name,
-						"idealOutput": self.ideal_output
+			"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
 

basalt/ressources/monitor/generation_types.py

@@ -34,6 +34,7 @@ class GenerationParams(BaseLogParams):
         input (Optional[str]): The input provided to the model.
         output (Optional[str]): The output generated by the model.
         variables (Optional[Dict[str, Any]]): Variables used in the prompt template.
+        options (Optional[Dict[str, Any]]): Additional options for the generation.
 
     Example:
         ```python
@@ -57,6 +58,7 @@ class GenerationParams(BaseLogParams):
     input: Optional[str] = None
     output: Optional[str] = None
     variables: Optional[Dict[str, Any]] = None
+    options: Optional[Dict[str, Any]] = None
 
 @dataclass
 class Generation(BaseLog):

basalt/sdk/datasetsdk.py

@@ -2,6 +2,7 @@
 SDK for interacting with Basalt datasets
 """
 from typing import Dict, List, Optional, Tuple, Any
+import asyncio
 
 from ..utils.dtos import (
     ListDatasetsDTO, GetDatasetDTO, CreateDatasetItemDTO,
@@ -46,6 +47,26 @@ def list(self) -> ListDatasetsResult:
             name=dataset.name,
             columns=dataset.columns
         ) for dataset in result.datasets]
+
+    async def async_list(self) -> ListDatasetsResult:
+        """
+        Asynchronously list all datasets available in the workspace.
+
+        Returns:
+            Tuple[Optional[Exception], Optional[List[DatasetDTO]]]: A tuple containing an optional 
+            exception and an optional list of DatasetDTO objects.
+        """
+        dto = ListDatasetsDTO()
+        err, result = await self._api.async_invoke(ListDatasetsEndpoint, dto)
+
+        if err is not None:
+            return err, None
+
+        return None, [DatasetDTO(
+            slug=dataset.slug,
+            name=dataset.name,
+            columns=dataset.columns
+        ) for dataset in result.datasets]
 
     def get(self, slug: str) -> GetDatasetResult:
         """
@@ -68,6 +89,28 @@ def get(self, slug: str) -> GetDatasetResult:
             return Exception(result.error), None
 
         return None, result.dataset
+
+    async def async_get(self, slug: str) -> GetDatasetResult:
+        """
+        Asynchronously get a dataset by its slug.
+
+        Args:
+            slug (str): The slug identifier for the dataset.
+
+        Returns:
+            Tuple[Optional[Exception], Optional[DatasetDTO]]: A tuple containing an optional
+            exception and an optional DatasetDTO.
+        """
+        dto = GetDatasetDTO(slug=slug)
+        err, result = await self._api.async_invoke(GetDatasetEndpoint, dto)
+
+        if err is not None:
+            return err, None
+
+        if result.error:
+            return Exception(result.error), None
+
+        return None, result.dataset
 
     def addRow(
         self,
@@ -108,3 +151,43 @@ def addRow(
             return Exception(result.error), None, None
 
         return None, result.datasetRow, result.warning
+
+    async def async_addRow(
+        self,
+        slug: str,
+        values: Dict[str, str],
+        name: Optional[str] = None,
+        ideal_output: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> CreateDatasetItemResult:
+        """
+        Asynchronously create a new item in a dataset.
+
+        Args:
+            slug (str): The slug identifier for the dataset.
+            values (Dict[str, str]): A dictionary of column values for the dataset item.
+            name (Optional[str]): An optional name for the dataset item.
+            ideal_output (Optional[str]): An optional ideal output for the dataset item.
+            metadata (Optional[Dict[str, Any]]): An optional metadata dictionary.
+
+        Returns:
+            Tuple[Optional[Exception], Optional[DatasetRowDTO], Optional[str]]: A tuple containing
+            an optional exception, an optional DatasetRowDTO, and an optional warning message.
+        """
+        dto = CreateDatasetItemDTO(
+            slug=slug,
+            values=values,
+            name=name,
+            idealOutput=ideal_output,
+            metadata=metadata
+        )
+
+        err, result = await self._api.async_invoke(CreateDatasetItemEndpoint, dto)
+
+        if err is not None:
+            return err, None, None
+
+        if result.error:
+            return Exception(result.error), None, None
+
+        return None, result.datasetRow, result.warning

basalt/sdk/monitorsdk.py

@@ -1,4 +1,5 @@
 from typing import Dict, Optional, Any, Tuple
+import asyncio
 
 from ..utils.protocols import IApi, ILogger
 from ..ressources.monitor.trace_types import TraceParams
@@ -40,6 +41,23 @@ def create_experiment(
             Experiment: A new Experiment instance.
         """
         return self._create_experiment(feature_slug, params)
+
+    async def async_create_experiment(
+        self,
+        feature_slug: str,
+        params: ExperimentParams
+    ) -> Tuple[Optional[Exception], Optional[Experiment]]:
+        """
+        Asynchronously creates a new experiment for monitoring.
+
+        Args:
+            feature_slug (str): The feature slug for the experiment.
+            params (Dict[str, Any]): Parameters for the experiment.
+
+        Returns:
+            Experiment: A new Experiment instance.
+        """
+        return await self._async_create_experiment(feature_slug, params)
 
 
     def create_trace(
@@ -63,6 +81,28 @@ def create_trace(
         trace_params = TraceParams(**params)
 
         return self._create_trace(slug, trace_params)
+
+    async def async_create_trace(
+        self,
+        slug: str,
+        params: Optional[TraceParams] = None
+    ) -> Trace:
+        """
+        Asynchronously creates a new trace for monitoring.
+
+        Args:
+            slug (str): The unique identifier for the trace.
+            params (TraceParams): Parameters for the trace.
+
+        Returns:
+            Trace: A new Trace instance.
+        """
+        if params is None:
+            params = {}
+
+        trace_params = TraceParams(**params)
+
+        return await self._async_create_trace(slug, trace_params)
 
     def create_generation(
         self,
@@ -79,6 +119,22 @@ def create_generation(
         """
         generation_params = GenerationParams(**params)
         return self._create_generation(generation_params)
+
+    async def async_create_generation(
+        self,
+        params: Dict[str, Any]
+    ) -> Generation:
+        """
+        Asynchronously creates a new generation for monitoring.
+
+        Args:
+            params (Dict[str, Any]): Parameters for the generation.
+
+        Returns:
+            Generation: A new Generation instance.
+        """
+        generation_params = GenerationParams(**params)
+        return await self._async_create_generation(generation_params)
 
     def create_log(
         self,
@@ -95,6 +151,22 @@ def create_log(
         """
         log_params = LogParams(**params)
         return self._create_log(log_params)
+
+    async def async_create_log(
+        self,
+        params: Dict[str, Any]
+    ) -> Log:
+        """
+        Asynchronously creates a new log for monitoring.
+
+        Args:
+            params (Dict[str, Any]): Parameters for the log.
+
+        Returns:
+            Log: A new Log instance.
+        """
+        log_params = LogParams(**params)
+        return await self._async_create_log(log_params)
 
     def _create_experiment(
         self,
@@ -123,6 +195,34 @@ def _create_experiment(
             return None, Experiment(result.experiment)
 
         return err, None
+
+    async def _async_create_experiment(
+        self,
+        feature_slug: str,
+        params: ExperimentParams
+    ) -> Tuple[Optional[Exception], Optional[Experiment]]:
+        """
+        Internal implementation for asynchronously creating an experiment.
+
+        Args:
+            feature_slug (str): The feature slug for the experiment.
+            params (ExperimentParams): Parameters for the experiment.
+
+        Returns:
+            Experiment: A new Experiment instance.
+        """
+        dto = CreateExperimentDTO(
+            feature_slug=feature_slug,
+            name=params.get("name"),
+        )
+
+        # Call the API endpoint
+        err, result = await self._api.async_invoke(CreateExperimentEndpoint, dto)
+
+        if err is None:
+            return None, Experiment(result.experiment)
+
+        return err, None
 
 
     def _create_trace(
@@ -157,6 +257,39 @@ def _create_trace(
         }
         trace = Trace(slug, params_dict, flusher, self._logger)
         return trace
+
+    async def _async_create_trace(
+        self,
+        slug: str,
+        params: TraceParams
+    ) -> Trace:
+        """
+        Internal implementation for asynchronously creating a trace.
+
+        Args:
+            slug (str): The unique identifier for the trace.
+            params (TraceParams): Parameters for the trace.
+
+        Returns:
+            Trace: A new Trace instance.
+        """
+        flusher = Flusher(self._api, self._logger)
+        # Convert TraceParams to a dictionary before passing to Trace
+        params_dict = {
+            "input": params.input,
+            "output": params.output,
+            "name": params.name,
+            "start_time": params.start_time,
+            "end_time": params.end_time,
+            "user": params.user,
+            "organization": params.organization,
+            "metadata": params.metadata,
+            "experiment": params.experiment,
+            "evaluators": params.evaluators,
+            "evaluationConfig": params.evaluation_config
+        }
+        trace = Trace(slug, params_dict, flusher, self._logger)
+        return trace
 
     def _create_generation(
         self,
@@ -186,6 +319,35 @@ def _create_generation(
             "options": params.options
         }
         return Generation(params_dict)
+
+    async def _async_create_generation(
+        self,
+        params: GenerationParams
+    ) -> Generation:
+        """
+        Internal implementation for asynchronously creating a generation.
+
+        Args:
+            params (GenerationParams): Parameters for the generation.
+
+        Returns:
+            Generation: A new Generation instance.
+        """
+        # Convert GenerationParams to a dictionary before passing to Generation
+        params_dict = {
+            "name": params.name,
+            "trace": params.trace,
+            "prompt": params.prompt,
+            "input": params.input,
+            "output": params.output,
+            "variables": params.variables,
+            "parent": params.parent,
+            "metadata": params.metadata,
+            "start_time": params.start_time,
+            "end_time": params.end_time,
+            "options": params.options
+        }
+        return Generation(params_dict)
 
     def _create_log(
         self,
@@ -194,6 +356,32 @@ def _create_log(
         """
         Internal implementation for creating a log.
 
+        Args:
+            params (LogParams): Parameters for the log.
+
+        Returns:
+            Log: A new Log instance.
+        """
+        # Convert LogParams to a dictionary before passing to Log
+        params_dict = {
+            "name": params.name,
+            "trace": params.trace,
+            "input": params.input,
+            "output": params.output,
+            "parent": params.parent,
+            "metadata": params.metadata,
+            "start_time": params.start_time,
+            "end_time": params.end_time
+        }
+        return Log(params_dict)
+
+    async def _async_create_log(
+        self,
+        params: LogParams
+    ) -> Log:
+        """
+        Internal implementation for asynchronously creating a log.
+
         Args:
             params (LogParams): Parameters for the log.