feat(sdk): get all reference and current datasets (#31)

radicalbit · Jun 25, 2024 · 5b1832d · 5b1832d
1 parent 6967044
commit 5b1832d
Showing 1 changed file with 45 additions and 1 deletion.
diff --git a/sdk/radicalbit_platform_sdk/apis/model.py b/sdk/radicalbit_platform_sdk/apis/model.py
@@ -5,7 +5,7 @@
 import boto3
 from botocore.exceptions import ClientError as BotoClientError
 import pandas as pd
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
 import requests
 
 from radicalbit_platform_sdk.apis import ModelCurrentDataset, ModelReferenceDataset
@@ -89,6 +89,50 @@ def delete(self) -> None:
             func=lambda _: None,
         )
 
+    def get_reference_datasets(self) -> List[ModelReferenceDataset]:
+        def __callback(response: requests.Response) -> List[ModelReferenceDataset]:
+            try:
+                adapter = TypeAdapter(List[ReferenceFileUpload])
+                references = adapter.validate_python(response.json())
+
+                return [
+                    ModelReferenceDataset(
+                        self.__base_url, self.__uuid, self.__model_type, ref
+                    )
+                    for ref in references
+                ]
+            except ValidationError as e:
+                raise ClientError(f'Unable to parse response: {response.text}') from e
+
+        return invoke(
+            method='GET',
+            url=f'{self.__base_url}/api/models/{str(self.__uuid)}/reference/all',
+            valid_response_code=200,
+            func=__callback,
+        )
+
+    def get_current_datasets(self) -> List[ModelCurrentDataset]:
+        def __callback(response: requests.Response) -> List[ModelCurrentDataset]:
+            try:
+                adapter = TypeAdapter(List[CurrentFileUpload])
+                references = adapter.validate_python(response.json())
+
+                return [
+                    ModelCurrentDataset(
+                        self.__base_url, self.__uuid, self.__model_type, ref
+                    )
+                    for ref in references
+                ]
+            except ValidationError as e:
+                raise ClientError(f'Unable to parse response: {response.text}') from e
+
+        return invoke(
+            method='GET',
+            url=f'{self.__base_url}/api/models/{str(self.__uuid)}/current/all',
+            valid_response_code=200,
+            func=__callback,
+        )
+
     def load_reference_dataset(
         self,
         file_name: str,