diff --git a/dagger/src/main/__init__.py b/dagger/src/main/__init__.py
index 57226c4..d74918a 100644
--- a/dagger/src/main/__init__.py
+++ b/dagger/src/main/__init__.py
@@ -19,10 +19,15 @@ async def publish(self, source: dagger.Directory) -> str:
     @function
     def build(self, source: dagger.Directory) -> dagger.Container:
         """Build the application container"""
-        return (
+        build = (
             self.build_env(source)
             .with_exec(["hatch", "build"])
         )
+        return (
+            dag.container()
+            .from_("python:3.10-slim")
+            .with_directory("/dist", build.directory("/src/dist"))
+        )
 
     @function
     async def test(self, source: dagger.Directory) -> str:
@@ -53,18 +58,3 @@ def build_env(self, source: dagger.Directory) -> dagger.Container:
             .with_exec(["pip", "install", "inquirer"])
             .with_exec(["pip", "install", "pytest"])
         )
-
-
-# Example usage:
-if __name__ == "__main__":
-    import anyio
-
-    async def main():
-        async with dagger.Connection() as client:
-            source = await client.host().directory(".", exclude=[".git", "__pycache__"])
-            killercoda_cli = KillercodaCli()
-
-            build_result = await killercoda_cli.publish(source)
-            print(build_result)
-
-    anyio.run(main)