diff --git a/main.py b/main.py
index f2471e9..0d9926c 100644
--- a/main.py
+++ b/main.py
@@ -1,7 +1,7 @@
 from speckle_automate import execute_automate_function
 
-from data_shield import FunctionInputs
-from data_shield import automate_function
+from data_shield.inputs import FunctionInputs
+from data_shield.function import automate_function
 
 # make sure to call the function with the executor
 if __name__ == "__main__":
diff --git a/src/data_shield/__init__.py b/src/data_shield/__init__.py
index 37b2a90..e69de29 100644
--- a/src/data_shield/__init__.py
+++ b/src/data_shield/__init__.py
@@ -1,9 +0,0 @@
-# You can expose frequently used imports for convenience
-from .actions import PrefixRemovalAction
-from .function import automate_function
-from .inputs import FunctionInputs, SanitizationMode
-from .rules import ParameterRules
-from .traversal import get_data_traversal_rules
-
-__all__ = ["PrefixRemovalAction", "automate_function", "FunctionInputs", "SanitizationMode", "ParameterRules",
-           "get_data_traversal_rules"]
diff --git a/src/data_shield/actions.py b/src/data_shield/actions.py
index 26ea2a1..8a9dab5 100644
--- a/src/data_shield/actions.py
+++ b/src/data_shield/actions.py
@@ -5,7 +5,7 @@ from typing import Dict, List, Optional
 from speckle_automate import AutomationContext
 from specklepy.objects import Base
 
-from data_shield import ParameterRules
+from data_shield.rules import ParameterRules
 
 
 # Our main goal is to define actions that can be taken on parameters.
diff --git a/src/data_shield/function.py b/src/data_shield/function.py
index b0ade91..3610042 100644
--- a/src/data_shield/function.py
+++ b/src/data_shield/function.py
@@ -1,7 +1,9 @@
 from speckle_automate import AutomationContext
 from specklepy.objects import Base
 
-from data_shield import FunctionInputs, SanitizationMode, PrefixRemovalAction, get_data_traversal_rules
+from data_shield.inputs import FunctionInputs, SanitizationMode
+from data_shield.actions import PrefixRemovalAction
+from data_shield.traversal import  get_data_traversal_rules
 
 
 class ParameterProcessor:
diff --git a/tests/test_function.py b/tests/test_function.py
index dc7fb0d..8a74af5 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -7,7 +7,7 @@ from speckle_automate import (
 )
 from speckle_automate.fixtures import *  # noqa: F401, F403
 
-from data_shield import automate_function, FunctionInputs, SanitizationMode
+from data_shield.function import automate_function, FunctionInputs, SanitizationMode
 
 
 class TestFunction: