diff --git a/optillm.py b/optillm.py
index a0c3432..b298e52 100644
--- a/optillm.py
+++ b/optillm.py
@@ -13,7 +13,7 @@
 from concurrent.futures import ThreadPoolExecutor
 
 # Import the LiteLLM wrapper
-from litellm_wrapper import LiteLLMWrapper
+from optillm.litellm_wrapper import LiteLLMWrapper
 
 # Import approach modules
 from optillm.mcts import chat_with_mcts
diff --git a/optillm/__init__.py b/optillm/__init__.py
index e69de29..26177f9 100644
--- a/optillm/__init__.py
+++ b/optillm/__init__.py
@@ -0,0 +1,13 @@
+from importlib import util
+import os
+
+# Get the path to the root optillm.py
+spec = util.spec_from_file_location(
+    "optillm.root",
+    os.path.join(os.path.dirname(os.path.dirname(__file__)), "optillm.py")
+)
+module = util.module_from_spec(spec)
+spec.loader.exec_module(module)
+
+# Export the main function
+main = module.main
\ No newline at end of file
diff --git a/litellm_wrapper.py b/optillm/litellm_wrapper.py
similarity index 100%
rename from litellm_wrapper.py
rename to optillm/litellm_wrapper.py
diff --git a/setup.py b/setup.py
index 1c9e876..32ab43a 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="optillm",
-    version="0.0.4",
+    version="0.0.5",
     packages=find_packages(),
     py_modules=['optillm'],
     install_requires=[