backend: make the data directory and the artifacts from the frontend customizable using environment variables

Signed-off-by: lucasew <lucas59356@gmail.com>

backend/apps/web/internal/db.py

@@ -1,4 +1,6 @@
 from peewee import *
+from config import DATA_DIR
 
-DB = SqliteDatabase("./data/ollama.db")
+
+DB = SqliteDatabase(str(DATA_DIR / "ollama.db"))
 DB.connect()

backend/apps/web/routers/utils.py

@@ -11,7 +11,7 @@ import json
 
 from utils.misc import calculate_sha256
 
-from config import OLLAMA_API_BASE_URL
+from config import OLLAMA_API_BASE_URL, DATA_DIR, UPLOAD_DIR
 from constants import ERROR_MESSAGES
 
 
@@ -96,8 +96,7 @@ async def download(
     file_name = parse_huggingface_url(url)
 
     if file_name:
-        os.makedirs("./uploads", exist_ok=True)
-        file_path = os.path.join("./uploads", f"{file_name}")
+        file_path = str(UPLOAD_DIR / file_name)
 
         return StreamingResponse(
             download_file_stream(url, file_path, file_name),
@@ -109,16 +108,15 @@ async def download(
 
 @router.post("/upload")
 def upload(file: UploadFile = File(...)):
-    os.makedirs("./data/uploads", exist_ok=True)
-    file_path = os.path.join("./data/uploads", file.filename)
+    file_path = UPLOAD_DIR / file.filename
 
     # Save file in chunks
-    with open(file_path, "wb+") as f:
+    with file_path.open("wb+") as f:
         for chunk in file.file:
             f.write(chunk)
 
     def file_process_stream():
-        total_size = os.path.getsize(file_path)
+        total_size = os.path.getsize(str(file_path))
         chunk_size = 1024 * 1024
         try:
             with open(file_path, "rb") as f:

backend/config.py

@@ -24,10 +24,12 @@ except ImportError:
 # File Upload
 ####################################
 
+DATA_DIR = Path(os.getenv("DATA_DIR", './data')).resolve()
 
-UPLOAD_DIR = "./data/uploads"
-Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
+UPLOAD_DIR = DATA_DIR / "uploads"
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 
+WEB_DIR = Path(os.getenv("WEB_DIR", "../build"))
 
 ####################################
 # ENV (dev,test,prod)
@@ -82,10 +84,10 @@ if WEBUI_AUTH and WEBUI_JWT_SECRET_KEY == "":
 # RAG
 ####################################
 
-CHROMA_DATA_PATH = "./data/vector_db"
+CHROMA_DATA_PATH = DATA_DIR / "vector_db"
 EMBED_MODEL = "all-MiniLM-L6-v2"
 CHROMA_CLIENT = chromadb.PersistentClient(
-    path=CHROMA_DATA_PATH, settings=Settings(allow_reset=True)
+    path=str(CHROMA_DATA_PATH), settings=Settings(allow_reset=True)
 )
 CHUNK_SIZE = 1500
 CHUNK_OVERLAP = 100

backend/main.py

@@ -14,7 +14,7 @@ from apps.openai.main import app as openai_app
 from apps.web.main import app as webui_app
 from apps.rag.main import app as rag_app
 
-from config import ENV
+from config import ENV, WEB_DIR
 
 
 class SPAStaticFiles(StaticFiles):
@@ -58,4 +58,4 @@ app.mount("/openai/api", openai_app)
 app.mount("/rag/api/v1", rag_app)
 
 
-app.mount("/", SPAStaticFiles(directory="../build", html=True), name="spa-static-files")
+app.mount("/", SPAStaticFiles(directory=str(WEB_DIR), html=True), name="spa-static-files")