restrict component imports

PrefectHQ · Nov 15, 2023 · a4bc53a · a4bc53a
1 parent a5f080a
commit a4bc53a
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/src/marvin/components/__init__.py b/src/marvin/components/__init__.py
@@ -1 +1,14 @@
-from .ai_function import ai_fn
+from .ai_function import ai_fn, AIFunction
+from .ai_classifier import ai_classifier, AIClassifier
+from .ai_model import ai_model
+from .prompt_function import prompt_fn, PromptFn
+
+__all__ = [
+    "ai_fn",
+    "ai_classifier",
+    "ai_model",
+    "prompt_fn",
+    "AIFunction",
+    "AIClassifier",
+    "PromptFn",
+]