diff --git a/main.go b/main.go
index cea755b..d847cbe 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"log"
 	"os"
+	"path/filepath"
 
 	"github.com/mudler/local-agent-framework/core/agent"
 	"github.com/mudler/local-agent-framework/core/state"
@@ -17,8 +18,7 @@ var apiKey = os.Getenv("API_KEY")
 var vectorStore = os.Getenv("VECTOR_STORE")
 var timeout = os.Getenv("TIMEOUT")
 var embeddingModel = os.Getenv("EMBEDDING_MODEL")
-
-const defaultChunkSize = 4098
+var stateDir = os.Getenv("STATE_DIR")
 
 func init() {
 	if testModel == "" {
@@ -30,19 +30,17 @@ func init() {
 	if timeout == "" {
 		timeout = "5m"
 	}
+	if stateDir == "" {
+		cwd, err := os.Getwd()
+		if err != nil {
+			panic(err)
+		}
+
+		stateDir = filepath.Join(cwd, "pool")
+	}
 }
 
-func main() {
-	// current dir
-	cwd, err := os.Getwd()
-	if err != nil {
-		panic(err)
-	}
-
-	stateDir := cwd + "/pool"
-	os.MkdirAll(stateDir, 0755)
-
-	var ragDB agent.RAGDB
+func ragDB() (ragDB agent.RAGDB) {
 	lai := llm.NewClient(apiKey, apiURL+"/v1", timeout)
 
 	switch vectorStore {
@@ -57,16 +55,30 @@ func main() {
 		}
 	}
 
+	return
+}
+
+func main() {
+	// make sure state dir exists
+	os.MkdirAll(stateDir, 0755)
+
+	// Initialize rag DB connection
+	ragDB := ragDB()
+
+	// Create the agent pool
 	pool, err := state.NewAgentPool(testModel, apiURL, stateDir, ragDB, webui.Actions, webui.Connectors, timeout)
 	if err != nil {
 		panic(err)
 	}
 
+	// Create the application
 	app := webui.NewApp(webui.WithPool(pool))
 
+	// Start the agents
 	if err := pool.StartAll(); err != nil {
 		panic(err)
 	}
 
+	// Start the web server
 	log.Fatal(app.Listen(":3000"))
 }