diff --git a/main.py b/main.py
index ff6c9c8..f2471e9 100644
--- a/main.py
+++ b/main.py
@@ -1,15 +1,10 @@
-from speckle_automate import (
-    execute_automate_function,
-)
-
-
-from src.function import automate_function
-from src.inputs import FunctionInputs
+from speckle_automate import execute_automate_function
 
+from data_shield import FunctionInputs
+from data_shield import automate_function
 
 # make sure to call the function with the executor
 if __name__ == "__main__":
     # NOTE: always pass in the automate function by its reference; do not invoke it!
 
-
     execute_automate_function(automate_function, FunctionInputs)
diff --git a/tests/test_function.py b/tests/test_function.py
index e577732..dc7fb0d 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -7,8 +7,7 @@ from speckle_automate import (
 )
 from speckle_automate.fixtures import *  # noqa: F401, F403
 
-from src.function import automate_function
-from src.inputs import FunctionInputs, SanitizationMode
+from data_shield import automate_function, FunctionInputs, SanitizationMode
 
 
 class TestFunction: