From a018df2734595080112d61c28b8c3cae5f6cca05 Mon Sep 17 00:00:00 2001
From: Tang Ziya <tcztzy@gmail.Com>
Date: Mon, 20 May 2024 11:12:03 +0800
Subject: [PATCH] infra: build

---
 backend/apps/web/internal/db.py |  6 ++++--
 backend/config.py               | 10 ++++++++--
 backend/open_webui/__init__.py  | 17 +++++++++++++++--
 hatch_build.py                  | 21 +++++++++++++++++++++
 pyproject.toml                  |  9 +++++++--
 5 files changed, 55 insertions(+), 8 deletions(-)
 create mode 100644 hatch_build.py

diff --git a/backend/apps/web/internal/db.py b/backend/apps/web/internal/db.py
index 136e3fafc..a6051de50 100644
--- a/backend/apps/web/internal/db.py
+++ b/backend/apps/web/internal/db.py
@@ -1,7 +1,7 @@
 from peewee import *
 from peewee_migrate import Router
 from playhouse.db_url import connect
-from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL
+from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
 import os
 import logging
 
@@ -18,6 +18,8 @@ else:
 
 DB = connect(DATABASE_URL)
 log.info(f"Connected to a {DB.__class__.__name__} database.")
-router = Router(DB, migrate_dir="apps/web/internal/migrations", logger=log)
+router = Router(
+    DB, migrate_dir=BACKEND_DIR / "apps" / "web" / "internal" / "migrations", logger=log
+)
 router.run()
 DB.connect(reuse_if_open=True)
diff --git a/backend/config.py b/backend/config.py
index b44b2b64e..9ba059008 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -1,6 +1,8 @@
 import os
 import sys
 import logging
+import importlib.metadata
+import pkgutil
 import chromadb
 from chromadb import Settings
 from base64 import b64encode
@@ -92,7 +94,10 @@ ENV = os.environ.get("ENV", "dev")
 try:
     PACKAGE_DATA = json.loads((BASE_DIR / "package.json").read_text())
 except:
-    PACKAGE_DATA = {"version": "0.0.0"}
+    try:
+        PACKAGE_DATA = {"version": importlib.metadata.version("open-webui")}
+    except importlib.metadata.PackageNotFoundError:
+        PACKAGE_DATA = {"version": "0.0.0"}
 
 VERSION = PACKAGE_DATA["version"]
 
@@ -119,7 +124,8 @@ def parse_section(section):
 try:
     changelog_content = (BASE_DIR / "CHANGELOG.md").read_text()
 except:
-    changelog_content = ""
+    changelog_content = (pkgutil.get_data("open_webui", "CHANGELOG.md") or b"").decode()
+
 
 # Convert markdown content to HTML
 html_content = markdown.markdown(changelog_content)
diff --git a/backend/open_webui/__init__.py b/backend/open_webui/__init__.py
index 8f22af8ca..1defac824 100644
--- a/backend/open_webui/__init__.py
+++ b/backend/open_webui/__init__.py
@@ -6,11 +6,11 @@ from pathlib import Path
 import typer
 import uvicorn
 
-import main
-
 app = typer.Typer()
 
 KEY_FILE = Path.cwd() / ".webui_secret_key"
+if (frontend_build_dir := Path(__file__).parent / "frontend").exists():
+    os.environ["FRONTEND_BUILD_DIR"] = str(frontend_build_dir)
 
 
 @app.command()
@@ -40,8 +40,21 @@ def serve(
                 "/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib",
             ]
         )
+    import main  # we need set environment variables before importing main
+
     uvicorn.run(main.app, host=host, port=port, forwarded_allow_ips="*")
 
 
+@app.command()
+def dev(
+    host: str = "0.0.0.0",
+    port: int = 8080,
+    reload: bool = True,
+):
+    uvicorn.run(
+        "main:app", host=host, port=port, reload=reload, forwarded_allow_ips="*"
+    )
+
+
 if __name__ == "__main__":
     app()
diff --git a/hatch_build.py b/hatch_build.py
new file mode 100644
index 000000000..2fa9e4805
--- /dev/null
+++ b/hatch_build.py
@@ -0,0 +1,21 @@
+# noqa: INP001
+import shutil
+import subprocess
+from sys import stderr
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+
+class CustomBuildHook(BuildHookInterface):
+    def initialize(self, version, build_data):
+        super().initialize(version, build_data)
+        stderr.write(">>> Building Open Webui frontend\n")
+        npm = shutil.which("npm")
+        if npm is None:
+            raise RuntimeError(
+                "NodeJS `npm` is required for building Open Webui but it was not found"
+            )
+        stderr.write("### npm install\n")
+        subprocess.run([npm, "install"], check=True)  # noqa: S603
+        stderr.write("\n### npm run build\n")
+        subprocess.run([npm, "run", "build"], check=True)  # noqa: S603
diff --git a/pyproject.toml b/pyproject.toml
index 37d96b16c..93889d8b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,6 +73,8 @@ dynamic = ["version"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Communications :: Chat",
     "Topic :: Multimedia",
 ]
@@ -95,9 +97,9 @@ allow-direct-references = true
 path = "package.json"
 pattern = '"version":\s*"(?P<version>[^"]+)"'
 
+[tool.hatch.build.hooks.custom]  # keep this for reading hooks from `hatch_build.py`
+
 [tool.hatch.build.targets.wheel]
-ignore-vcs = true
-packages = ["backend"]
 sources = ["backend"]
 exclude = [
     ".dockerignore",
@@ -107,4 +109,7 @@ exclude = [
     "requirements.txt",
     "start.sh",
     "start_windows.bat",
+    "webui.db",
+    "chroma.sqlite3",
 ]
+force-include = { "CHANGELOG.md" = "open_webui/CHANGELOG.md", build = "open_webui/frontend" }