From 1176d58a943ca323b206de337181fda0759bea1d Mon Sep 17 00:00:00 2001
From: Ibrahim Muhammad <imuhammad@labelbox.com>
Date: Wed, 6 Dec 2023 21:24:08 -0800
Subject: [PATCH] [AL-0] Add ontology_id for model app

---
 labelbox/schema/foundry/app.py            |  1 +
 labelbox/schema/foundry/foundry_client.py | 27 ++-------------
 tests/integration/test_foundry.py         | 41 ++++++++++++++++-------
 3 files changed, 32 insertions(+), 37 deletions(-)

diff --git a/labelbox/schema/foundry/app.py b/labelbox/schema/foundry/app.py
index dc72769c1..67536d0b0 100644
--- a/labelbox/schema/foundry/app.py
+++ b/labelbox/schema/foundry/app.py
@@ -12,6 +12,7 @@ class App(_CamelCaseMixin, BaseModel):
     description: str
     inference_params: Dict[str, Any]
     class_to_schema_id: Dict[str, str]
+    ontology_id: str
     created_by: Optional[str] = None
 
     @classmethod
diff --git a/labelbox/schema/foundry/foundry_client.py b/labelbox/schema/foundry/foundry_client.py
index b1d530439..50df486d4 100644
--- a/labelbox/schema/foundry/foundry_client.py
+++ b/labelbox/schema/foundry/foundry_client.py
@@ -15,11 +15,12 @@ def _create_app(self, app: App) -> App:
         field_names_str = "\n".join(APP_FIELD_NAMES)
         query_str = f"""
             mutation CreateDataRowAttachmentPyApi(
-                $name: String!, $modelId: ID!, $description: String, $inferenceParams: Json!, $classToSchemaId: Json!
+                $name: String!, $modelId: ID!, $ontologyId: ID!, $description: String, $inferenceParams: Json!, $classToSchemaId: Json!
             ){{
                 createModelFoundryApp(input: {{
                     name: $name 
                     modelId: $modelId 
+                    ontologyId: $ontologyId
                     description: $description
                     inferenceParams: $inferenceParams
                     classToSchemaId: $classToSchemaId
@@ -73,31 +74,9 @@ def _delete_app(self, id: str) -> None:
             raise exceptions.LabelboxError(f'Unable to delete app with id {id}',
                                            e)
 
-    def _get_model(self, id: str) -> Model:
-        field_names_str = "\n".join(MODEL_FIELD_NAMES)
-
-        query_str = f"""
-            query GetModelByIdPyApi($id: ID!) {{
-                modelFoundryModel(where: {{id: $id}}) {{
-                    model {{
-                        {field_names_str}
-                    }}
-                }}
-            }}
-        """
-        params = {"id": id}
-
-        try:
-            response = self.client.execute(query_str, params)
-        except Exception as e:
-            raise exceptions.LabelboxError(f'Unable to get model with id {id}',
-                                           e)
-        return Model(**response["modelFoundryModel"]["model"])
-
     def run_app(self, model_run_name: str,
                 data_rows: Union[DataRowIds, GlobalKeys], app_id: str) -> Task:
         app = self._get_app(app_id)
-        model = self._get_model(app.model_id)
 
         data_rows_query = self.client.build_catalog_query(data_rows)
 
@@ -110,7 +89,7 @@ def run_app(self, model_run_name: str,
                 "query": [data_rows_query],
                 "scope": None
             },
-            "ontologyId": model.ontology_id
+            "ontologyId": app.ontology_id
         }
 
         query = """
diff --git a/tests/integration/test_foundry.py b/tests/integration/test_foundry.py
index 246673f58..67c633e18 100644
--- a/tests/integration/test_foundry.py
+++ b/tests/integration/test_foundry.py
@@ -19,14 +19,34 @@ def foundry_client(client):
 
 
 @pytest.fixture()
-def unsaved_app(random_str):
-    return App(
-        model_id=TEST_MODEL_ID,
-        name=f"Test App {random_str}",
-        description="Test App Description",
-        inference_params={"confidence": 0.2},
-        class_to_schema_id={},
-    )
+def ontology(client, random_str):
+    object_features = [
+        lb.Tool(tool=lb.Tool.Type.BBOX,
+                name="text",
+                color="#ff0000",
+                classifications=[
+                    lb.Classification(class_type=lb.Classification.Type.TEXT,
+                                      name="value")
+                ])
+    ]
+
+    ontology_builder = lb.OntologyBuilder(tools=object_features,)
+
+    ontology = client.create_ontology(
+        f"Test ontology for tesseract model {random_str}",
+        ontology_builder.asdict(),
+        media_type=lb.MediaType.Image)
+    return ontology
+
+
+@pytest.fixture()
+def unsaved_app(random_str, ontology):
+    return App(model_id=TEST_MODEL_ID,
+               name=f"Test App {random_str}",
+               description="Test App Description",
+               inference_params={"confidence": 0.2},
+               class_to_schema_id={},
+               ontology_id=ontology.uid)
 
 
 @pytest.fixture()
@@ -55,11 +75,6 @@ def test_get_app_with_invalid_id(foundry_client):
         foundry_client._get_app("invalid-id")
 
 
-def test_get_model_by_id(foundry_client):
-    model = foundry_client._get_model(TEST_MODEL_ID)
-    assert str(model.id) == TEST_MODEL_ID
-
-
 def test_run_foundry_app_with_data_row_id(foundry_client, data_row, app,
                                           random_str):
     data_rows = lb.DataRowIds([data_row.uid])