浏览代码

Merge pull request #536 from lucasew/backend-data-dir

Unhardcode some stuff, flexibilize dotenv requirement
Timothy Jaeryang Baek 1 年之前
父节点
当前提交
a243e53bfe
共有 4 个文件被更改,包括 30 次插入23 次删除
  1. 3 1
      backend/apps/web/internal/db.py
  2. 3 5
      backend/apps/web/routers/utils.py
  3. 18 15
      backend/config.py
  4. 6 2
      backend/main.py

+ 3 - 1
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(f"{DATA_DIR}/ollama.db")
 DB.connect()

+ 3 - 5
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 = f"{UPLOAD_DIR}/{file_name}"
 
         return StreamingResponse(
             download_file_stream(url, file_path, file_name),
@@ -109,8 +108,7 @@ 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 = f"{UPLOAD_DIR}/{file.filename}"
 
     # Save file in chunks
     with open(file_path, "wb+") as f:

+ 18 - 15
backend/config.py

@@ -1,36 +1,39 @@
-from dotenv import load_dotenv, find_dotenv
 import os
-
-
 import chromadb
 from chromadb import Settings
-
-
 from secrets import token_bytes
 from base64 import b64encode
-
 from constants import ERROR_MESSAGES
-
-
 from pathlib import Path
 
-load_dotenv(find_dotenv("../.env"))
+try:
+    from dotenv import load_dotenv, find_dotenv
+
+    load_dotenv(find_dotenv("../.env"))
+except ImportError:
+    print("dotenv not installed, skipping...")
 
 
 ####################################
-# File Upload
+# ENV (dev,test,prod)
 ####################################
 
+ENV = os.environ.get("ENV", "dev")
 
-UPLOAD_DIR = "./data/uploads"
-Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
 
+####################################
+# DATA/FRONTEND BUILD DIR
+####################################
+
+DATA_DIR = str(Path(os.getenv("DATA_DIR", "./data")).resolve())
+FRONTEND_BUILD_DIR = str(Path(os.getenv("FRONTEND_BUILD_DIR", "../build")))
 
 ####################################
-# ENV (dev,test,prod)
+# File Upload DIR
 ####################################
 
-ENV = os.environ.get("ENV", "dev")
+UPLOAD_DIR = f"{DATA_DIR}/uploads"
+Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
 
 ####################################
 # OLLAMA_API_BASE_URL
@@ -107,7 +110,7 @@ if WEBUI_AUTH and WEBUI_JWT_SECRET_KEY == "":
 # RAG
 ####################################
 
-CHROMA_DATA_PATH = "./data/vector_db"
+CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db"
 EMBED_MODEL = "all-MiniLM-L6-v2"
 CHROMA_CLIENT = chromadb.PersistentClient(
     path=CHROMA_DATA_PATH, settings=Settings(allow_reset=True)

+ 6 - 2
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, FRONTEND_BUILD_DIR
 
 
 class SPAStaticFiles(StaticFiles):
@@ -58,4 +58,8 @@ 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=FRONTEND_BUILD_DIR, html=True),
+    name="spa-static-files",
+)