diff --git a/flake.nix b/flake.nix
index 7814ff8..ca93d8c 100644
--- a/flake.nix
+++ b/flake.nix
@@ -6,17 +6,7 @@
   let
     system = "x86_64-linux";
     pkgs = nixpkgs.legacyPackages.${system};
-    _nosocket = self.legacyPackages."${system}".nosocket;
-    _devContainer = (import ./api {
-      inherit pkgs _nosocket;
-      nosocket = _nosocket;
-      appEnv = "dev";
-    }).container;
-    _prodContainer = (import ./api {
-      inherit pkgs _nosocket;
-      nosocket = _nosocket;
-      appEnv = "prod";
-    }).container;
+    nosocket = self.legacyPackages."${system}".nosocket;
     args = {
       inherit pkgs;
       piston = {
@@ -81,7 +71,16 @@
       piston = (import ./api { inherit pkgs nosocket; }).package;
     };
 
-    devContainer = _devContainer;
-    prodContainer = _prodContainer;
+    containers = {
+      dev = (import ./api {
+        inherit pkgs nosocket;
+        appEnv = "dev";
+      }).container;
+
+      prod = (import ./api {
+        inherit pkgs nosocket;
+        appEnv = "prod";
+      }).container;
+    };
   };
 }
diff --git a/piston b/piston
index 048d362..54b9b01 100755
--- a/piston
+++ b/piston
@@ -13,7 +13,7 @@ SUBCOMMAND="$1"
 shift
 
 build_base() {
-    container_flake_key=$([[ $1 == "dev" ]] && echo "devContainer" || echo "prodContainer")
+    container_flake_key=$([[ $1 == "dev" ]] && echo "containers.dev" || echo "containers.prod")
     CONTAINER_PATH="$(nix build ".#$container_flake_key" --no-link --json | jq '.[0].outputs.out' -r)"
     echo "The image archive was created at: $CONTAINER_PATH"
     docker load -i $CONTAINER_PATH || exit 1