From fd46fbd961b703272e9b72f13e3bde98b0e7347c Mon Sep 17 00:00:00 2001
From: Jedd Morgan <45512892+JR-Morgan@users.noreply.github.com>
Date: Tue, 18 Feb 2025 11:47:41 +0000
Subject: [PATCH] Updated functions

---
 src/speckle_automate/automation_context.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/speckle_automate/automation_context.py b/src/speckle_automate/automation_context.py
index ad90b2f..d969ed1 100644
--- a/src/speckle_automate/automation_context.py
+++ b/src/speckle_automate/automation_context.py
@@ -5,7 +5,7 @@
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 import httpx
 from gql import gql
@@ -138,9 +138,15 @@ def create_new_model_in_project(
 
         return self.speckle_client.model.create(input)
 
+    def get_model(self, model_id: str) -> Model:
+        """Gets the model that the automation was triggered on"""
+        return self.speckle_client.model.get(
+            model_id, self.automation_run_data.project_id
+        )
+
     def create_new_version_in_project(
         self, root_object: Base, model_id: str, version_message: str = ""
-    ) -> Tuple[Model, Version]:
+    ) -> Version:
         """Save a base model to a new version on the project.
 
         Args:
@@ -149,14 +155,10 @@ def create_new_version_in_project(
             version_message (str): The message for the new version.
         """
 
-        model = self.speckle_client.model.get(
-            self.automation_run_data.project_id, model_id
-        )  # TODO: Did we want to avoid fetching the model? we needed it before in order to
-
         matching_trigger = [
             t
             for t in self.automation_run_data.triggers
-            if t.payload.model_id == model.id
+            if t.payload.model_id == model_id
         ]
         if matching_trigger:
             raise ValueError(
@@ -173,7 +175,7 @@ def create_new_version_in_project(
 
         create_version_input = CreateVersionInput(
             object_id=root_object_id,
-            model_id=model.id,
+            model_id=model_id,
             project_id=self.automation_run_data.project_id,
             message=version_message,
             source_application="SpeckleAutomate",
@@ -181,7 +183,7 @@ def create_new_version_in_project(
         version = self.speckle_client.version.create(create_version_input)
 
         self._automation_result.result_versions.append(version.id)
-        return model, version
+        return version
 
     @property
     def context_view(self) -> Optional[str]: