diff --git a/.gitignore b/.gitignore
index 1113436..fc35145 100644
--- a/.gitignore
+++ b/.gitignore
@@ -130,4 +130,5 @@ dist
 .pnp.*
 
 storybook-static/
-stories/
+
+stats.html
diff --git a/lib/NodeGraphEditor.tsx b/lib/NodeGraphEditor.tsx
index 01740de..2c354db 100644
--- a/lib/NodeGraphEditor.tsx
+++ b/lib/NodeGraphEditor.tsx
@@ -1,8 +1,4 @@
 import {
-  Background,
-  BackgroundVariant,
-  Edge,
-  Node,
   ReactFlow,
   ReactFlowProps,
   ReactFlowProvider,
@@ -12,12 +8,9 @@ import {
   useReactFlow,
   useStoreApi,
 } from '@xyflow/react'
-import {
-  GraphConfigProvider,
-  useGraphConfig,
-} from './context/GraphConfigContext'
+import { useGraphConfig } from './context/GraphConfigContext'
 import '@xyflow/react/dist/style.css'
-import { useBuildGraphConfig, useNodeTypes } from './hooks/config'
+import { useNodeTypes } from './hooks/config'
 import {
   forwardRef,
   useImperativeHandle,
@@ -27,7 +20,6 @@ import {
   useEffect,
 } from 'react'
 import { defaultEdgeTypes } from './edge-types'
-import { IGraphConfig } from './config'
 import { useSocketConnect } from './hooks/connect'
 import { useHotkeys } from 'react-hotkeys-hook'
 import { ClipboardItem } from './clipboard'
@@ -62,26 +54,6 @@ export const NodeGraphEditor = forwardRef<
   },
 )
 
-type ExampleNodeGraphEditorProps = {
-  nodes: Node[]
-  edges: Edge[]
-  config: IGraphConfig
-}
-
-export const ExampleNodeGraphEditor = forwardRef<
-  NodeGraphHandle,
-  ExampleNodeGraphEditorProps
->(({ nodes, edges, config: _config }: ExampleNodeGraphEditorProps, ref) => {
-  const config = useBuildGraphConfig(_config)
-  return (
-    <GraphConfigProvider defaultConfig={config}>
-      <NodeGraphEditor ref={ref} defaultNodes={nodes} defaultEdges={edges}>
-        <Background color="#52525b" variant={BackgroundVariant.Dots} />
-      </NodeGraphEditor>
-    </GraphConfigProvider>
-  )
-})
-
 type FlowProps = ReactFlowProps & {
   backgroundStyles?: CSSProperties,
   /**
@@ -90,7 +62,7 @@ type FlowProps = ReactFlowProps & {
   layoutEngine?: LayoutEngine
 }
 export type NodeGraphHandle = {
-  layout: () => void
+  layout: (engine?: LayoutEngine) => void
 }
 
 const Flow = forwardRef<NodeGraphHandle, FlowProps>(
@@ -117,7 +89,8 @@ const Flow = forwardRef<NodeGraphHandle, FlowProps>(
     )
 
     // Provide methods to parent components
-    const layout = useLayoutEngine(layoutEngine ?? LayoutEngine.Dagre)
+    const layout = useLayoutEngine()
+
     useImperativeHandle(
       ref,
       () => ({
@@ -131,10 +104,10 @@ const Flow = forwardRef<NodeGraphHandle, FlowProps>(
       const shouldLayout = !!getState().nodes.find(
         (node) => node.position == undefined,
       )
-      if (initialized && shouldLayout) {
-        layout()
+      if (initialized && shouldLayout && layoutEngine) {
+        layout(layoutEngine)
       }
-    }, [initialized])
+    }, [initialized, layoutEngine])
 
     return (
       <div
diff --git a/lib/index.ts b/lib/index.ts
index 7647613..61d1a4a 100644
--- a/lib/index.ts
+++ b/lib/index.ts
@@ -17,3 +17,4 @@ export { NodeSelectField } from './components/NodeSelectField.tsx'
 export { NodeDenseLinkedField } from './components/NodeDenseLinkedField.tsx'
 export { NodeLinkedField } from './components/NodeLinkedField.tsx'
 export { NodeOutputField } from './components/NodeOutputField.tsx'
+export { registerLayoutEngine } from './layout/layout.ts'
diff --git a/lib/layout/layout.ts b/lib/layout/layout.ts
index 3bb7ec5..9f85c66 100644
--- a/lib/layout/layout.ts
+++ b/lib/layout/layout.ts
@@ -1,18 +1,16 @@
 import { useCallback } from 'react'
-import { computeDagreLayout } from './dagre'
 import { Edge, Instance, Node, useReactFlow } from '@xyflow/react'
 
-export enum LayoutEngine {
-  Dagre,
-}
+export type LayoutEngine = string
 
-type LayoutFunc = () => void
+export type LayoutFunc = (nodes: Node[], edges: Edge[]) => Node[]
 
-export function useLayoutEngine(engine: LayoutEngine): LayoutFunc {
+export function useLayoutEngine() {
   const { getNodes, getEdges, setNodes } = useReactFlow()
-  return useCallback(() => {
+  return useCallback((engine?: LayoutEngine) => {
+    if (!engine) return
     setNodes(computeLayout(engine, getNodes, getEdges))
-  }, [engine])
+  }, [])
 }
 
 function computeLayout(
@@ -25,11 +23,19 @@ function computeLayout(
   return layoutFn(getNodes(), getEdges())
 }
 
+const layoutEngines: Record<LayoutEngine, LayoutFunc> = {}
+
+export function registerLayoutEngine(
+  engine: LayoutEngine,
+  layoutFn: LayoutFunc,
+) {
+  layoutEngines[engine] = layoutFn
+}
+
 export function getLayoutFunction(
-  engine: LayoutEngine | undefined,
+  engine: LayoutEngine,
 ): ((nodes: Node[], edges: Edge[]) => Node[]) | undefined {
-  switch (engine) {
-    case LayoutEngine.Dagre:
-      return computeDagreLayout
-  }
+  const layoutFn = layoutEngines[engine]
+  if (!layoutFn) throw new Error(`Unknown layout engine ${engine}`)
+  return layoutFn
 }
diff --git a/lib/stories/ExampleNodeGraphEditor.tsx b/lib/stories/ExampleNodeGraphEditor.tsx
new file mode 100644
index 0000000..998d968
--- /dev/null
+++ b/lib/stories/ExampleNodeGraphEditor.tsx
@@ -0,0 +1,33 @@
+import {
+  Background,
+  BackgroundVariant,
+  Edge,
+  Node,
+} from '@xyflow/react'
+import { forwardRef } from "react"
+import { NodeGraphEditor, NodeGraphHandle } from "../NodeGraphEditor"
+import { GraphConfigProvider } from "../context/GraphConfigContext"
+import { useBuildGraphConfig } from '../hooks/config'
+import { IGraphConfig } from "../config"
+
+type ExampleNodeGraphEditorProps = {
+    nodes: Node[]
+    edges: Edge[]
+    config: IGraphConfig
+  }
+  
+  export const ExampleNodeGraphEditor = forwardRef<
+    NodeGraphHandle,
+    ExampleNodeGraphEditorProps
+  >(({ nodes, edges, config: _config }: ExampleNodeGraphEditorProps, ref) => {
+    const config = useBuildGraphConfig(_config)
+    return (
+      <GraphConfigProvider defaultConfig={config}>
+        <NodeGraphEditor ref={ref} defaultNodes={nodes} defaultEdges={edges}>
+          <Background color="#52525b" variant={BackgroundVariant.Dots} />
+        </NodeGraphEditor>
+      </GraphConfigProvider>
+    )
+  })
+  
+  
\ No newline at end of file
diff --git a/lib/NodeGraphEditor.stories.tsx b/lib/stories/NodeGraphEditor.stories.tsx
similarity index 99%
rename from lib/NodeGraphEditor.stories.tsx
rename to lib/stories/NodeGraphEditor.stories.tsx
index bd61a79..be5bd20 100644
--- a/lib/NodeGraphEditor.stories.tsx
+++ b/lib/stories/NodeGraphEditor.stories.tsx
@@ -1,6 +1,6 @@
-import { ExampleNodeGraphEditor } from './NodeGraphEditor'
+import { ExampleNodeGraphEditor } from './ExampleNodeGraphEditor'
 import { Meta, StoryObj } from '@storybook/react'
-import { IGraphConfig } from './config'
+import { IGraphConfig } from '../config'
 
 const simpleConfig: IGraphConfig = {
   valueTypes: {
diff --git a/lib/NodeGraphEditorCustomInput.stories.tsx b/lib/stories/NodeGraphEditorCustomInput.stories.tsx
similarity index 88%
rename from lib/NodeGraphEditorCustomInput.stories.tsx
rename to lib/stories/NodeGraphEditorCustomInput.stories.tsx
index 6fc0448..2fe3df0 100644
--- a/lib/NodeGraphEditorCustomInput.stories.tsx
+++ b/lib/stories/NodeGraphEditorCustomInput.stories.tsx
@@ -1,15 +1,15 @@
-import { NodeGraphEditor } from './NodeGraphEditor'
+import { NodeGraphEditor } from '../NodeGraphEditor.tsx'
 import { Meta, StoryObj } from '@storybook/react'
-import { GraphConfigProvider } from './context/GraphConfigContext'
+import { GraphConfigProvider } from '../context/GraphConfigContext.tsx'
 import {
   Background,
   BackgroundVariant,
   Edge,
   Node,
 } from '@xyflow/react'
-import { useBuildGraphConfig } from './hooks/config.ts'
-import { NodeInputField } from './components/NodeInputField.tsx'
-import { InputProps } from './config.ts'
+import { useBuildGraphConfig } from '../hooks/config.ts'
+import { NodeInputField } from '../components/NodeInputField.tsx'
+import { InputProps } from '../config.ts'
 
 const meta = {
   title: 'Node Graph Editor',
diff --git a/lib/NodeGraphEditorCustomNode.stories.tsx b/lib/stories/NodeGraphEditorCustomNode.stories.tsx
similarity index 88%
rename from lib/NodeGraphEditorCustomNode.stories.tsx
rename to lib/stories/NodeGraphEditorCustomNode.stories.tsx
index ff41291..50a8cbb 100644
--- a/lib/NodeGraphEditorCustomNode.stories.tsx
+++ b/lib/stories/NodeGraphEditorCustomNode.stories.tsx
@@ -1,7 +1,7 @@
-import { NodeGraphEditor } from './NodeGraphEditor'
+import { NodeGraphEditor } from '../NodeGraphEditor'
 import { Meta, StoryObj } from '@storybook/react'
-import { GraphConfigProvider } from './context/GraphConfigContext'
-import { GraphConfig } from './config'
+import { GraphConfigProvider } from '../context/GraphConfigContext'
+import { GraphConfig } from '../config'
 import { useMemo } from 'react'
 import {
   Background,
@@ -10,9 +10,9 @@ import {
   Node,
   Position,
 } from '@xyflow/react'
-import { NodeContainer } from './components/NodeContainer'
-import { useFocusBlur } from './hooks/focus'
-import { Handle } from './components/Handle'
+import { NodeContainer } from '../components/NodeContainer'
+import { useFocusBlur } from '../hooks/focus'
+import { Handle } from '../components/Handle'
 
 const meta = {
   title: 'Node Graph Editor',
diff --git a/lib/NodeGraphEditorInputGroups.stories.tsx b/lib/stories/NodeGraphEditorInputGroups.stories.tsx
similarity index 93%
rename from lib/NodeGraphEditorInputGroups.stories.tsx
rename to lib/stories/NodeGraphEditorInputGroups.stories.tsx
index 99ab04d..8f96c63 100644
--- a/lib/NodeGraphEditorInputGroups.stories.tsx
+++ b/lib/stories/NodeGraphEditorInputGroups.stories.tsx
@@ -1,16 +1,20 @@
-import { NodeGraphEditor } from './NodeGraphEditor'
+import { NodeGraphEditor } from '../NodeGraphEditor.tsx'
 import { Meta, StoryObj } from '@storybook/react'
-import { GraphConfigProvider } from './context/GraphConfigContext'
+import { GraphConfigProvider } from '../context/GraphConfigContext.tsx'
 import {
   Background,
   BackgroundVariant,
   Edge,
   Node,
 } from '@xyflow/react'
-import { useBuildGraphConfig } from './hooks/config.ts'
-import { InputProps } from './config.ts'
+import { useBuildGraphConfig } from '../hooks/config.ts'
+import { InputProps } from '../config.ts'
 import { Wheel } from '@uiw/react-color'
-import { useNodeFieldValue } from './hooks/node.ts'
+import { useNodeFieldValue } from '../hooks/node.ts'
+import { registerLayoutEngine } from '../layout/layout.ts'
+import { computeDagreLayout } from '../layout/dagre.ts'
+
+registerLayoutEngine('dagre', computeDagreLayout)
 
 const meta = {
   title: 'Node Graph Editor',
@@ -201,8 +205,8 @@ const meta = {
     )
     return (
       <GraphConfigProvider defaultConfig={config}>
-        <NodeGraphEditor defaultNodes={nodes} defaultEdges={edges}>
-          <Background color="#52525b" variant={BackgroundVariant.Dots} />
+        <NodeGraphEditor defaultNodes={nodes} defaultEdges={edges} layoutEngine="dagre">
+          <Background color="#52525b" variant={BackgroundVariant.Dots}/>
         </NodeGraphEditor>
       </GraphConfigProvider>
     )
diff --git a/lib/NodeGraphEditorLayout.stories.tsx b/lib/stories/NodeGraphEditorLayout.stories.tsx
similarity index 86%
rename from lib/NodeGraphEditorLayout.stories.tsx
rename to lib/stories/NodeGraphEditorLayout.stories.tsx
index d394a3d..d09df4f 100644
--- a/lib/NodeGraphEditorLayout.stories.tsx
+++ b/lib/stories/NodeGraphEditorLayout.stories.tsx
@@ -1,6 +1,11 @@
-import { ExampleNodeGraphEditor, NodeGraphHandle } from './NodeGraphEditor'
+import { ExampleNodeGraphEditor } from './ExampleNodeGraphEditor'
+import { NodeGraphHandle } from '../NodeGraphEditor'
 import { Meta, StoryObj } from '@storybook/react'
 import { useRef } from 'react'
+import { registerLayoutEngine } from '../layout/layout'
+import { computeDagreLayout } from '../layout/dagre'
+
+registerLayoutEngine('dagre', computeDagreLayout)
 
 const meta = {
   title: 'Node Graph Editor',
@@ -8,12 +13,14 @@ const meta = {
     const ref = useRef<NodeGraphHandle>(null)
     return (
       <div style={{ width: '100%', height: '100%' }}>
-        <button
-          style={{ position: 'absolute', top: 10, left: 10, zIndex: 1000 }}
-          onClick={() => ref.current!.layout()}
-        >
-          Layout
-        </button>
+        <div style={{position: 'absolute', top: 10, left: 10, zIndex: 1000}}>
+          <button
+            style={{  }}
+            onClick={() => ref.current!.layout('dagre')}
+          >
+            Dagre Layout
+          </button>
+        </div>
         <ExampleNodeGraphEditor
           ref={ref}
           config={config}