seekrai 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

seekrai/resources/models.py

@@ -1,56 +1,14 @@
 from __future__ import annotations
 
-import os
-from pathlib import Path
-from typing import Any, List
-
-from tqdm import tqdm
-from tqdm.utils import CallbackIOWrapper
-
 from seekrai.abstract import api_requestor
-from seekrai.constants import DISABLE_TQDM
+from seekrai.resources.resource_base import ResourceBase
 from seekrai.seekrflow_response import SeekrFlowResponse
-from seekrai.types import ModelList, ModelResponse, SeekrFlowClient, SeekrFlowRequest
-from seekrai.types.models import ModelType
-
-
-class Models:
-    def __init__(self, client: SeekrFlowClient) -> None:
-        self._client = client
-
-    def upload(
-        self,
-        file: Path | str,
-        *,
-        model_type: ModelType | str = ModelType.OBJECT_DETECTION,
-    ) -> ModelResponse:
-        if isinstance(file, str):
-            file = Path(file)
+from seekrai.types import ModelList, ModelResponse, SeekrFlowRequest
+from seekrai.types.common import ObjectType
+from seekrai.utils import parse_timestamp
 
-        requestor = api_requestor.APIRequestor(
-            client=self._client,
-        )
-        file_size = os.stat(file.as_posix()).st_size
-
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            desc=f"Uploading model file {file.name}",
-            disable=bool(DISABLE_TQDM),
-        ) as t:
-            with file.open("rb") as f:
-                reader_wrapper = CallbackIOWrapper(t.update, f, "read")
-                response, _, _ = requestor.request(
-                    options=SeekrFlowRequest(
-                        method="PUT",
-                        url="flow/pt-models",
-                        files={"files": reader_wrapper, "filename": file.name},
-                        params={"purpose": model_type},
-                    ),
-                )
-        return ModelResponse(**response.data)
 
+class Models(ResourceBase):
     def list(
         self,
     ) -> ModelList:
@@ -58,7 +16,7 @@ class Models:
         Method to return list of models on the API
 
         Returns:
-            List[ModelResponse]: List of model objects
+            ModelList: List of model objects
         """
 
         requestor = api_requestor.APIRequestor(
@@ -68,95 +26,36 @@ class Models:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url="flow/pt-models",
-            ),
-            stream=False,
-        )
-
-        assert isinstance(response, SeekrFlowResponse)
-        return ModelList(**response.data)
-
-    def promote(self, id: str) -> ModelResponse:
-        requestor = api_requestor.APIRequestor(
-            client=self._client,
-        )
-
-        response, _, _ = requestor.request(
-            options=SeekrFlowRequest(
-                method="GET",
-                url=f"flow/pt-models/{id}/promote-model",
-                params={"model_id": id},
-            ),
-            stream=False,
-        )
-
-        assert isinstance(response, SeekrFlowResponse)
-
-        return ModelResponse(**response.data)
-
-    def demote(self, id: str) -> ModelResponse:
-        requestor = api_requestor.APIRequestor(
-            client=self._client,
-        )
-
-        response, _, _ = requestor.request(
-            options=SeekrFlowRequest(
-                method="GET",
-                url=f"flow/pt-models/{id}/demote-model",
-                params={"model_id": id},
+                url="flow/models",
             ),
             stream=False,
         )
-
         assert isinstance(response, SeekrFlowResponse)
 
-        return ModelResponse(**response.data)
+        model_responses = [
+            ModelResponse(
+                id=str(model["id"]),
+                object=ObjectType.Model,
+                name=model["name"],
+                bytes=model["size"],
+                created_at=parse_timestamp(model["created_at"]),
+                model_type=model["model_type"],
+            )
+            for model in response.data["data"]
+        ]
 
-    def predict(self, id: str, file: Path | str) -> Any:
-        requestor = api_requestor.APIRequestor(
-            client=self._client,
-        )
+        return ModelList(data=model_responses)
 
-        if isinstance(file, str):
-            file = Path(file)
-
-        file_size = os.stat(file.as_posix()).st_size
-
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            desc=f"Uploading file {file.name}",
-            disable=bool(DISABLE_TQDM),
-        ):
-            with file.open("rb") as f:
-                response, _, _ = requestor.request(
-                    options=SeekrFlowRequest(
-                        method="POST",
-                        url="flow/pt-models/predict",
-                        files={"files": f, "filename": file.name},
-                        params={"model_id": id},
-                    ),
-                    stream=False,
-                )
-
-        assert isinstance(response, SeekrFlowResponse)
-
-        return response.data
-
-
-class AsyncModels:
-    def __init__(self, client: SeekrFlowClient) -> None:
-        self._client = client
 
+class AsyncModels(ResourceBase):
     async def list(
         self,
-    ) -> List[ModelResponse]:
+    ) -> ModelList:
         """
         Async method to return list of models on API
 
         Returns:
-            List[ModelResponse]: List of model objects
+            ModelList: List of model objects
         """
 
         requestor = api_requestor.APIRequestor(
@@ -174,4 +73,16 @@ class AsyncModels:
         assert isinstance(response, SeekrFlowResponse)
         assert isinstance(response.data, list)
 
-        return [ModelResponse(**model) for model in response.data]
+        model_responses = [
+            ModelResponse(
+                id=str(model["id"]),
+                object=ObjectType.Model,
+                name=model["name"],
+                bytes=model["size"],
+                created_at=parse_timestamp(model["created_at"]),
+                model_type=model["model_type"],
+            )
+            for model in response.data["data"]
+        ]
+
+        return ModelList(data=model_responses)

seekrai/resources/projects.py

@@ -1,6 +1,7 @@
 from seekrai.abstract import api_requestor
+from seekrai.resources.resource_base import ResourceBase
 from seekrai.seekrflow_response import SeekrFlowResponse
-from seekrai.types import SeekrFlowClient, SeekrFlowRequest
+from seekrai.types import SeekrFlowRequest
 from seekrai.types.projects import (
     GetProjectsResponse,
     PostProjectRequest,
@@ -10,10 +11,7 @@ from seekrai.types.projects import (
 )
 
 
-class Projects:
-    def __init__(self, client: SeekrFlowClient) -> None:
-        self._client = client
-
+class Projects(ResourceBase):
     def list(self, skip: int = 0, limit: int = 100) -> GetProjectsResponse:
         requestor = api_requestor.APIRequestor(
             client=self._client,
@@ -70,10 +68,7 @@ class Projects:
         return ProjectSchema(**response.data)
 
 
-class AsyncProjects:
-    def __init__(self, client: SeekrFlowClient) -> None:
-        self._client = client
-
+class AsyncProjects(ResourceBase):
     async def list(self, skip: int = 0, limit: int = 100) -> GetProjectsResponse:
         requestor = api_requestor.APIRequestor(
             client=self._client,

seekrai/resources/resource_base.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from seekrai.types import SeekrFlowClient
+
+
+class ResourceBase:
+    """Base class for resources."""
+
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client