Added schema, team and simple file upload methods.

lib/python-beta/src/bailo/bailo.py

@@ -1,8 +1,8 @@
 """Main entry point"""
 from __future__ import annotations
 import requests
-from typing import List, Optional, Any
-from .enums import ModelVisibility
+from typing import List, Optional, Dict, Any
+from .enums import ModelVisibility, SchemaKind
 
 class Agent:
     def __init__(self):
@@ -34,7 +34,7 @@ def create_model(
 
         :param name: Name of the model
         :param description: Description of the model
-        :param visibility: Object to define model visibility (e.g public or private)
+        :param visibility: Enum to define model visibility (e.g public or private)
         :return: JSON response object
         """
         return self.agent.post(
@@ -64,7 +64,7 @@ def find_models(
         """
         return self.agent.get(
             f"{self.url}/v2/models/search",
-            json={
+            params={
                 "task": task,
                 "libraries": libraries,
                 "filters": filters,
@@ -99,7 +99,7 @@ def update_model(
         :param model_id: Unique model ID
         :param name: Name of the model, defaults to None
         :param description: Description of the model, defaults to None
-        :param visibility: Object to define model visibility (e.g. public or private), defaults to None
+        :param visibility: Enum to define model visibility (e.g. public or private), defaults to None
         :return: JSON response object
         """
         x = {}
@@ -269,7 +269,28 @@ def get_files(
             f"{self.url}/v2/model/{model_id}/files",
         ).json()
 
-    #def simple_upload(): TBC
+    def simple_upload(
+        self,
+        model_id: str,
+        name: str,
+        binary: bytes,
+        mime: Optional[str] = None,
+    ):
+        """
+        Creates a simple file upload.
+
+        :param model_id: Unique model ID
+        :param name: File name
+        :param binary: File data
+        :param mime: MIME aka media type, defaults to None
+        :return: JSON response object
+        """
+        return self.agent.post(
+            f"{self.url}/v2/model/{model_id}/files/upload/simple",
+            params={"name": name, "mime": mime},
+            data = binary,
+        ).json()  
+
 
     #def start_multi_upload(): TBC
 
@@ -291,15 +312,65 @@ def delete_file(
             f"{self.url}/v2/model/{model_id}/files/{file_id}",
         ).json()
 
-    #def get_all_schemas():
+    def get_all_schemas(
+        self,
+        kind: Optional[SchemaKind] = None,
+    ):
+        """
+        Gets all schemas.
+
+        :param kind: Enum to define schema kind (e.g. Model or AccessRequest), defaults to None
+        :return: JSON response object
+        """        
+        return self.agent.get(
+            f"{self.url}/v2/schemas",
+            params={"kind": kind},
+        ).json()
+
+    def get_schema(
+        self,
+        schema_id: str,
+    ):
+        """
+        Retrieves a specific schema using its unique ID.
+
+        :param schema_id: Unique schema ID
+        :return: JSON response object.
+        """        
+        return self.agent.get(
+            f"{self.url}/v2/schema/{schema_id}",
+        ).json()
 
-    #def get_schema():
 
-    #def create_schema():
+    def create_schema(
+        self,
+        schema_id: str,
+        name: str,
+        kind: SchemaKind,
+        json_schema: Dict[str, Any],
+    ):
+        """
+        Creates a schema.
+
+        :param schema_id: Unique schema ID
+        :param name: Name of the schema
+        :param kind: Enum to define schema kind (e.g. Model or AccessRequest)
+        :param json_schema: JSON schema
+        :return: JSON response object
+        """        
+        return self.agent.post(
+            f"{self.url}/v2/schemas",
+            json={
+                "id": schema_id,
+                "name": name,
+                "kind": kind,
+                "jsonSchema": json_schema,
+            }
+        ).json()
 
-    #def get_reviews():
+    #def get_reviews(): TBC
 
-    #def get_reviews_count():
+    #def get_reviews_count(): TBC
 
     def get_model_roles(
         self,
@@ -374,4 +445,46 @@ def get_user_teams(
         """        
         return self.agent.get(
             f"{self.url}/v2/teams/mine",
+        ).json()
+
+    def get_team(
+        self,
+        team_id: str,
+    ):
+        """
+        Retrieves a specific team given its unique ID.
+
+        :param team_id: Unique team ID
+        :return: JSON response object
+        """    
+        return self.agent.get(
+            f"{self.url}/v2/team/{team_id}",
+        ).json()
+
+    def update_team(
+        self,
+        team_id: str,
+        name: Optional[str] = None,
+        description: Optional[str] = None,      
+    ):
+        """
+        Updates a team given its unique ID.
+
+        :param team_id: Unique team ID
+        :param name: Name of team, defaults to None
+        :param description: Description of team, defaults to None
+        :return: JSON response object
+        """    
+
+        x = {}
+
+        if name is not None:
+            x.update({"name": name})
+
+        if description is not None:
+            x.update({"description": description})
+
+        return self.agent.patch(
+            f"{self.url}/v2/team/{team_id}",
+            json=x,
         ).json()

lib/python-beta/src/bailo/enums.py

@@ -3,3 +3,7 @@
 class ModelVisibility(str, Enum):
     Private = 'private'
     Public = 'public'
+
+class SchemaKind(str, Enum):
+    Model = 'model'
+    AccessRequest = 'accessRequest'

lib/python-beta/tests/test_client.py

@@ -7,7 +7,7 @@
 import json
 
 from bailo import BailoClient
-from bailo.enums import ModelVisibility
+from bailo.enums import ModelVisibility, SchemaKind
 
 mock_result = {"success": True}
 
@@ -26,11 +26,11 @@ def test_create_model(requests_mock):
 
 
 def test_find_models(requests_mock):
-    requests_mock.get("https://example.com/api/v2/models/search", json={"success": True})
+    requests_mock.get("https://example.com/api/v2/models/search?task=image_classification", json={"success": True})
 
     client = BailoClient("https://example.com")
     result = client.find_models(
-        task="Image Classification"
+        task="image_classification"
     )
 
     assert result == {"success": True}
@@ -156,7 +156,20 @@ def test_get_files(requests_mock):
 
     assert result == {"success": True}
 
-#def test_simple_upload(requests_mock):
+def test_simple_upload(requests_mock):
+    requests_mock.post("https://example.com/api/v2/model/test_id/files/upload/simple?name=test.txt", json={"success": True})
+
+    data = 'TEST'
+    data = bytes(data, 'utf-8')
+
+    client = BailoClient("https://example.com")
+    result = client.simple_upload(
+        model_id="test_id",
+        name="test.txt",
+        binary=data,
+    )
+
+    assert result == {"success": True}
 
 #def test_start_multi_upload(requests_mock):
 
@@ -173,11 +186,38 @@ def test_delete_file(requests_mock):
 
     assert result == {"success": True}
 
-#def test_get_all_schemas(requests_mock):
+def test_get_all_schemas(requests_mock):
+    requests_mock.get("https://example.com/api/v2/schemas?kind=model", json={"success": True})
+
+    client = BailoClient("https://example.com")
+    result = client.get_all_schemas(
+        kind=SchemaKind.Model
+    )
+
+    assert result == {"success": True}
+
+def test_get_schema(requests_mock):
+    requests_mock.get("https://example.com/api/v2/schema/test_id", json={"success": True})
+
+    client = BailoClient("https://example.com")
+    result = client.get_schema(
+        schema_id="test_id"
+    )
+
+    assert result == {"success": True}
+
+def test_create_schema(requests_mock):
+    requests_mock.post("https://example.com/api/v2/schemas", json={"success": True})
 
-#def test_get_schema(requests_mock):
+    client = BailoClient("https://example.com")
+    result = client.create_schema(
+        schema_id="test_id",
+        name="test",
+        kind=SchemaKind.Model,
+        json_schema={"test": "test"}
+    )
 
-#def test_create_schema(requests_mock):
+    assert result == {"success": True}
 
 #def test_get_reviews(requests_mock):
 
@@ -229,4 +269,25 @@ def test_get_user_teams(requests_mock):
     client = BailoClient("https://example.com")
     result = client.get_user_teams()
 
-    assert result == {"success": True}
+    assert result == {"success": True}
+
+def test_get_team(requests_mock):
+    requests_mock.get("https://example.com/api/v2/team/test_id", json={"success": True})
+
+    client = BailoClient("https://example.com")
+    result = client.get_team(
+        team_id="test_id",
+    )
+
+    assert result == {"success": True}
+
+def test_update_team(requests_mock):
+    requests_mock.patch("https://example.com/api/v2/team/test_id", json={"success": True})
+
+    client = BailoClient("https://example.com")
+    result = client.update_team(
+        team_id="test_id",
+        name="name",
+    )
+
+    assert result == {"success": True}