From e97a45348897c2c9bb206ce894a2fe36d468fe6f Mon Sep 17 00:00:00 2001
From: Petros Toupas <ptoupas@gmail.com>
Date: Fri, 10 Jan 2025 11:06:09 +0200
Subject: [PATCH] Update is_hubai_available to work with hubAI API calls

---
 modelconverter/utils/hubai_utils.py | 49 ++++++++++++-----------------
 1 file changed, 20 insertions(+), 29 deletions(-)

diff --git a/modelconverter/utils/hubai_utils.py b/modelconverter/utils/hubai_utils.py
index 618c086..1a29985 100644
--- a/modelconverter/utils/hubai_utils.py
+++ b/modelconverter/utils/hubai_utils.py
@@ -1,36 +1,27 @@
-import requests
-
-
 def is_hubai_available(model_slug: str) -> bool:
-    url = "https://easyml.cloud.luxonis.com/models/api/v1/models?is_public=true&limit=1000"
-    response = requests.get(url)
-    if response.status_code != 200:
+    from modelconverter.cli import Request, slug_to_id
+
+    model_name = model_slug.split(":")[0]
+    if len(model_slug.split(":")) < 2:
         raise ValueError(
-            f"Failed to get models. Status code: {response.status_code}"
+            f"Model variant not found in {model_slug}. Please specify it."
         )
-    hub_ai_models = response.json()
-    for model in hub_ai_models:
-        slug = f"{model['team_slug']}/{model['slug']}"
+
+    model_id = slug_to_id(
+        model_slug.removeprefix("luxonis/").split(":")[0], "models"
+    )
+    model_variants = Request.get(
+        "modelVersions/", params={"model_id": model_id, "is_public": True}
+    )
+
+    for version in model_variants:
         if (
-            slug in model_slug
-            or slug.removeprefix(f"{model['team_slug']}/") in model_slug
+            f"{model_name}:{version['variant_slug']}" == model_slug
+            or f"{model_name}:{version['variant_slug']}".removeprefix(
+                "luxonis/"
+            )
+            == model_slug
         ):
-            model_id = model["id"]
+            return True
 
-            url = f"https://easyml.cloud.luxonis.com/models/api/v1/modelVersions?model_id={model_id}&is_public=true"
-            response = requests.get(url)
-            if response.status_code != 200:
-                raise ValueError(
-                    f"Failed to get model versions. Status code: {response.status_code}"
-                )
-            model_versions = response.json()
-            for version in model_versions:
-                if (
-                    f"{slug}:{version['variant_slug']}" == model_slug
-                    or f"{slug}:{version['variant_slug']}".removeprefix(
-                        f"{model['team_slug']}/"
-                    )
-                    == model_slug
-                ):
-                    return True
     return False