aixplain
diff --git a/‎.pre-commit-config.yaml‎
Lines changed: 8 additions & 1 deletion b/‎.pre-commit-config.yaml‎
Lines changed: 8 additions & 1 deletion
diff --git a/‎aixplain/cli_groups.py‎
Lines changed: 3 additions & 2 deletions b/‎aixplain/cli_groups.py‎
Lines changed: 3 additions & 2 deletions
diff --git a/‎aixplain/factories/cli/model_factory_cli.py‎
Lines changed: 37 additions & 15 deletions b/‎aixplain/factories/cli/model_factory_cli.py‎
Lines changed: 37 additions & 15 deletions
diff --git a/‎aixplain/factories/model_factory.py‎
Lines changed: 79 additions & 36 deletions b/‎aixplain/factories/model_factory.py‎
Lines changed: 79 additions & 36 deletions
@@ -15,4 +15,11 @@ repos:
       - id: black
         language_version: python3
         args: # arguments to configure black
-          - --line-length=128
+          - --line-length=128
+  
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.0.0 # Use the latest version
+    hooks:
+      - id: flake8
+        args: # arguments to configure black
+          - --ignore=E402,E501
@@ -21,7 +21,7 @@
     CLI Runner
 """
 import click
-from aixplain.factories.cli.model_factory_cli import list_host_machines, list_functions, create_asset_repo, asset_repo_login, onboard_model, deploy_huggingface_model, get_huggingface_model_status
+from aixplain.factories.cli.model_factory_cli import list_host_machines, list_functions, create_asset_repo, asset_repo_login, onboard_model, deploy_huggingface_model, get_huggingface_model_status, list_gpus
 
 @click.group('cli')
 def cli():
@@ -51,11 +51,12 @@ def onboard():
 create.add_command(create_asset_repo)
 list.add_command(list_host_machines)
 list.add_command(list_functions)
+list.add_command(list_gpus)
 get.add_command(asset_repo_login)
 get.add_command(get_huggingface_model_status)
 onboard.add_command(onboard_model)
 onboard.add_command(deploy_huggingface_model)
 
 
 def run_cli():
-    cli()
+    cli()
@@ -44,7 +44,7 @@ def list_host_machines(api_key: Optional[Text] = None) -> None:
     click.echo(ret_val_yaml)
 
 @click.command("functions")
-@click.option("--verbose", default=False, 
+@click.option("--verbose", is_flag=True, 
               help="List all function details, False by default.")
 @click.option("--api-key", default=None, 
               help="TEAM_API_KEY if not already set in environment.")
@@ -62,21 +62,37 @@ def list_functions(verbose: bool, api_key: Optional[Text] = None) -> None:
     ret_val_yaml = yaml.dump(ret_val)
     click.echo(ret_val_yaml)
 
+@click.command("gpus")
+@click.option("--api-key", default=None, 
+              help="TEAM_API_KEY if not already set in environment.")
+def list_gpus(api_key: Optional[Text] = None) -> None:
+    """CLI wrapper function for the LIST_GPUS function in ModelFactory.
+
+    Args:
+        api_key (Text, optional): Team API key. Defaults to None.
+    Returns:
+        None
+    """
+    ret_val = ModelFactory.list_gpus(api_key)
+    ret_val_yaml = yaml.dump(ret_val)
+    click.echo(ret_val_yaml)
+
 @click.command("image-repo")
 @click.option("--name", help="Model name.")
-@click.option("--hosting-machine", 
-              help="Hosting machine code obtained from LIST_HOSTS.")
-@click.option("--version", help="Model version.")
 @click.option("--description", help="Description of model.")
 @click.option("--function", help="Function name obtained from LIST_FUNCTIONS.")
 @click.option("--source-language", default="en", 
               help="Model source language in 2-character 639-1 code or 3-character 639-3 code.")
+@click.option("--input-modality", help="Input type (text, video, image, etc.)")
+@click.option("--output-modality", help="Output type (text, video, image, etc.)")
+@click.option("--documentation-url", default="", help="Link to model documentation.")
 @click.option("--api-key", default=None, 
               help="TEAM_API_KEY if not already set in environment.")
-def create_asset_repo(name: Text, hosting_machine: Text, version: Text, 
-                          description: Text, function: Text, 
-                          source_language: Text, 
-                          api_key: Optional[Text] = None) -> None:
+def create_asset_repo(name: Text, description: Text, function: Text,
+                      source_language: Text, input_modality: Text, 
+                      output_modality: Text, 
+                      documentation_url: Optional[Text] = "",
+                      api_key: Optional[Text] = None) -> None:
     """CLI wrapper function for the CREATE_ASSET_REPO function in ModelFactory.
 
     Args:
@@ -93,9 +109,10 @@ def create_asset_repo(name: Text, hosting_machine: Text, version: Text,
     Returns:
         None
     """
-    ret_val = ModelFactory.create_asset_repo(name, hosting_machine, version, 
-                                             description, function, 
-                                             source_language, api_key)
+    ret_val = ModelFactory.create_asset_repo(name, description, function, 
+                                             source_language, input_modality, 
+                                             output_modality, documentation_url, 
+                                             api_key)
     ret_val_yaml = yaml.dump(ret_val)
     click.echo(ret_val_yaml)
 
@@ -119,8 +136,10 @@ def asset_repo_login(api_key: Optional[Text] = None) -> None:
 @click.option("--model-id", help="Model ID from CREATE_IMAGE_REPO.")
 @click.option("--image-tag", help="The tag of the image that you would like hosted.")
 @click.option("--image-hash", help="The hash of the image you would like onboarded.")
+@click.option("--host-machine", default="", help="The machine on which to host the model.")
 @click.option("--api-key", default=None, help="TEAM_API_KEY if not already set in environment.")
 def onboard_model(model_id: Text, image_tag: Text, image_hash: Text, 
+                  host_machine: Optional[Text] = "", 
                   api_key: Optional[Text] = None) -> None:
     """CLI wrapper function for the ONBOARD_MODEL function in ModelFactory.
 
@@ -132,17 +151,20 @@ def onboard_model(model_id: Text, image_tag: Text, image_hash: Text,
     Returns:
         None
     """   
-    ret_val = ModelFactory.onboard_model(model_id, image_tag, image_hash, api_key)
+    ret_val = ModelFactory.onboard_model(model_id, image_tag, image_hash, 
+                                         host_machine, api_key)
     ret_val_yaml = yaml.dump(ret_val)
     click.echo(ret_val_yaml)
 
 @click.command("hf-model")
 @click.option("--name", help="User-defined name for Hugging Face model.")
 @click.option("--hf-repo-id", help="Repository ID from Hugging Face in {supplier}/{model name} form.")
-@click.option("--hf-token", help="Hugging Face token used to authenticate to this model.")
+@click.option("--revision", default="", help="Commit hash of repository.")
+@click.option("--hf-token", default=None, help="Hugging Face token used to authenticate to this model.")
 @click.option("--api-key", default=None, help="TEAM_API_KEY if not already set in environment.")
 def deploy_huggingface_model(name: Text, hf_repo_id: Text, 
                              hf_token: Optional[Text] = None, 
+                             revision: Optional[Text] = None,
                              api_key: Optional[Text] = None) -> None:
     """CLI wrapper function for the DEPLOY_HUGGINGFACE_MODEL function in ModelFactory.
 
@@ -153,7 +175,7 @@ def deploy_huggingface_model(name: Text, hf_repo_id: Text,
     Returns:
         None
     """
-    ret_val = ModelFactory.deploy_huggingface_model(name, hf_repo_id, hf_token, api_key)
+    ret_val = ModelFactory.deploy_huggingface_model(name, hf_repo_id, revision, hf_token, api_key)
     ret_val_yaml = yaml.dump(ret_val)
     click.echo(ret_val_yaml)
 
@@ -172,4 +194,4 @@ def get_huggingface_model_status(model_id: Text, api_key: Optional[Text] = None)
     """
     ret_val = ModelFactory.get_huggingface_model_status(model_id, api_key)
     ret_val_yaml = yaml.dump(ret_val)
-    click.echo(ret_val_yaml)
+    click.echo(ret_val_yaml)
@@ -24,6 +24,7 @@
 import json
 import logging
 from aixplain.modules.model import Model
+from aixplain.modules.model.llm_model import LLM
 from aixplain.enums import Function, Language, OwnershipType, Supplier, SortBy, SortOrder
 from aixplain.utils import config
 from aixplain.utils.file_utils import _request_with_retry
@@ -60,13 +61,18 @@ def _create_model_from_response(cls, response: Dict) -> Model:
                 if "language" in param["name"]:
                     parameters[param["name"]] = [w["value"] for w in param["values"]]
 
-        return Model(
+        function = Function(response["function"]["id"])
+        ModelClass = Model
+        if function == Function.TEXT_GENERATION:
+            ModelClass = LLM
+
+        return ModelClass(
             response["id"],
             response["name"],
             supplier=response["supplier"],
             api_key=response["api_key"],
             cost=response["pricing"],
-            function=Function(response["function"]["id"]),
+            function=function,
             parameters=parameters,
             is_subscribed=True if "subscription" in response else False,
             version=response["version"]["id"],
@@ -100,7 +106,7 @@ def get(cls, model_id: Text, api_key: Optional[Text] = None) -> Model:
             model = cls._create_model_from_response(resp)
             logging.info(f"Model Creation: Model {model_id} instantiated.")
             return model
-        except Exception as e:
+        except Exception:
             if resp is not None and "statusCode" in resp:
                 status_code = resp["statusCode"]
                 message = resp["message"]
@@ -135,7 +141,7 @@ def _get_assets_from_page(
         sort_order: SortOrder = SortOrder.ASCENDING,
     ) -> List[Model]:
         try:
-            url = urljoin(cls.backend_url, f"sdk/models/paginate")
+            url = urljoin(cls.backend_url, "sdk/models/paginate")
             filter_params = {"q": query, "pageNumber": page_number, "pageSize": page_size}
             if is_finetunable is not None:
                 filter_params["isFineTunable"] = is_finetunable
@@ -253,7 +259,7 @@ def list_host_machines(cls, api_key: Optional[Text] = None) -> List[Dict]:
             List[Dict]: List of dictionaries containing information about
             each hosting machine.
         """
-        machines_url = urljoin(config.BACKEND_URL, f"sdk/hosting-machines")
+        machines_url = urljoin(config.BACKEND_URL, "sdk/hosting-machines")
         logging.debug(f"URL: {machines_url}")
         if api_key:
             headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
@@ -264,6 +270,25 @@ def list_host_machines(cls, api_key: Optional[Text] = None) -> List[Dict]:
         for dictionary in response_dicts:
             del dictionary["id"]
         return response_dicts
+    
+    @classmethod
+    def list_gpus(cls, api_key: Optional[Text] = None) -> List[List[Text]]:
+        """List GPU names on which you can host your language model.
+
+        Args:
+            api_key (Text, optional): Team API key. Defaults to None.
+
+        Returns:
+            List[List[Text]]: List of all available GPUs and their prices.
+        """
+        gpu_url = urljoin(config.BACKEND_URL, "sdk/model-onboarding/gpus")
+        if api_key:
+            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
+        else:
+            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+        response = _request_with_retry("get", gpu_url, headers=headers)
+        response_list = json.loads(response.text)
+        return response_list
 
     @classmethod
     def list_functions(cls, verbose: Optional[bool] = False, api_key: Optional[Text] = None) -> List[Dict]:
@@ -278,7 +303,7 @@ def list_functions(cls, verbose: Optional[bool] = False, api_key: Optional[Text]
             List[Dict]: List of dictionaries containing information about
             each supported function.
         """
-        functions_url = urljoin(config.BACKEND_URL, f"sdk/functions")
+        functions_url = urljoin(config.BACKEND_URL, "sdk/functions")
         logging.debug(f"URL: {functions_url}")
         if api_key:
             headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
@@ -304,12 +329,13 @@ def list_functions(cls, verbose: Optional[bool] = False, api_key: Optional[Text]
     def create_asset_repo(
         cls,
         name: Text,
-        hosting_machine: Text,
-        version: Text,
         description: Text,
         function: Text,
         source_language: Text,
-        api_key: Optional[Text] = None,
+        input_modality: Text,
+        output_modality: Text,
+        documentation_url: Optional[Text] = "",
+        api_key: Optional[Text] = None
     ) -> Dict:
         """Creates an image repository for this model and registers it in the
         platform backend.
@@ -336,27 +362,36 @@ def create_asset_repo(
                 function_id = function_dict["id"]
         if function_id is None:
             raise Exception("Invalid function name")
-        create_url = urljoin(config.BACKEND_URL, f"sdk/models/register")
+        create_url = urljoin(config.BACKEND_URL, f"sdk/models/onboard")
         logging.debug(f"URL: {create_url}")
         if api_key:
             headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
         else:
             headers = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
-        always_on = False
-        is_async = False  # Hard-coded to False for first release
+
         payload = {
-            "name": name,
-            "hostingMachine": hosting_machine,
-            "alwaysOn": always_on,
-            "version": version,
-            "description": description,
-            "function": function_id,
-            "isAsync": is_async,
-            "sourceLanguage": source_language,
+            "model": {
+                "name": name,
+                "description": description,
+                "connectionType": [
+                    "synchronous"
+                ],
+                "function": function_id,
+                "modalities": [
+                    f"{input_modality}-{output_modality}"
+                ],
+                "documentationUrl": documentation_url,
+                "sourceLanguage": source_language
+            },
+            "source": "aixplain-ecr",
+            "onboardingParams": {
+            }
         }
-        payload = json.dumps(payload)
         logging.debug(f"Body: {str(payload)}")
-        response = _request_with_retry("post", create_url, headers=headers, data=payload)
+        response = _request_with_retry("post", create_url, headers=headers, json=payload)
+
+        assert response.status_code == 201
+
         return response.json()
 
     @classmethod
@@ -370,23 +405,26 @@ def asset_repo_login(cls, api_key: Optional[Text] = None) -> Dict:
         Returns:
             Dict: Backend response
         """
-        login_url = urljoin(config.BACKEND_URL, f"sdk/ecr/login")
+        login_url = urljoin(config.BACKEND_URL, "sdk/ecr/login")
         logging.debug(f"URL: {login_url}")
         if api_key:
-            headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
+            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
         else:
-            headers = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
         response = _request_with_retry("post", login_url, headers=headers)
+        print(f"Response: {response}")
         response_dict = json.loads(response.text)
         return response_dict
 
     @classmethod
-    def onboard_model(cls, model_id: Text, image_tag: Text, image_hash: Text, api_key: Optional[Text] = None) -> Dict:
+    def onboard_model(cls, model_id: Text, image_tag: Text, image_hash: Text, host_machine: Optional[Text] = "", api_key: Optional[Text] = None) -> Dict:
         """Onboard a model after its image has been pushed to ECR.
 
         Args:
             model_id (Text): Model ID obtained from CREATE_ASSET_REPO.
             image_tag (Text): Image tag to be onboarded.
+            image_hash (Text): Image digest.
+            host_machine (Text, optional): Machine on which to host model.
             api_key (Text, optional): Team API key. Defaults to None.
         Returns:
             Dict: Backend response
@@ -397,18 +435,18 @@ def onboard_model(cls, model_id: Text, image_tag: Text, image_hash: Text, api_ke
             headers = {"x-api-key": f"{api_key}", "Content-Type": "application/json"}
         else:
             headers = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
-        payload = {"image": image_tag, "sha": image_hash}
-        payload = json.dumps(payload)
+        payload = {"image": image_tag, "sha": image_hash, "hostMachine": host_machine}
         logging.debug(f"Body: {str(payload)}")
-        response = _request_with_retry("post", onboard_url, headers=headers, data=payload)
-        message = "Your onboarding request has been submitted to an aiXplain specialist for finalization. We will notify you when the process is completed."
-        logging.info(message)
+        response = _request_with_retry("post", onboard_url, headers=headers, json=payload)
+        if response.status_code == 201:
+            message = "Your onboarding request has been submitted to an aiXplain specialist for finalization. We will notify you when the process is completed."
+            logging.info(message)
+        else:
+            message = "An error has occurred. Please make sure your model_id is valid and your host_machine, if set, is a valid option from the LIST_GPUS function."
         return response
 
     @classmethod
-    def deploy_huggingface_model(
-        cls, name: Text, hf_repo_id: Text, hf_token: Optional[Text] = "", api_key: Optional[Text] = None
-    ) -> Dict:
+    def deploy_huggingface_model(cls, name: Text, hf_repo_id: Text, revision: Optional[Text] = "", hf_token: Optional[Text] = "", api_key: Optional[Text] = None) -> Dict:
         """Onboards and deploys a Hugging Face large language model.
 
         Args:
@@ -420,7 +458,7 @@ def deploy_huggingface_model(
             Dict: Backend response
         """
         supplier, model_name = hf_repo_id.split("/")
-        deploy_url = urljoin(config.BACKEND_URL, f"sdk/model-onboarding/onboard")
+        deploy_url = urljoin(config.BACKEND_URL, "sdk/model-onboarding/onboard")
         if api_key:
             headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
         else:
@@ -435,7 +473,12 @@ def deploy_huggingface_model(
                 "sourceLanguage": "en",
             },
             "source": "huggingface",
-            "onboardingParams": {"hf_model_name": model_name, "hf_supplier": supplier, "hf_token": hf_token},
+            "onboardingParams": {
+                "hf_supplier": supplier,
+                "hf_model_name": model_name,
+                "hf_token": hf_token,
+                "revision": revision
+            }
         }
         response = _request_with_retry("post", deploy_url, headers=headers, json=body)
         logging.debug(response.text)