From fd4ef3fb320c735303cd37756ff632e6cd139d6b Mon Sep 17 00:00:00 2001
From: "mauro.cortellazzi" <mauro.cortellazzi@radicalbit.io>
Date: Tue, 25 Jun 2024 13:25:43 +0200
Subject: [PATCH] feat(sdk): get all reference and current datasets

---
 sdk/radicalbit_platform_sdk/apis/model.py | 46 ++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/sdk/radicalbit_platform_sdk/apis/model.py b/sdk/radicalbit_platform_sdk/apis/model.py
index 5e3d4211..739aaf3f 100644
--- 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,