From 93bf3e8a1f87760e04d6a18b2e27bae0f5d5da0e Mon Sep 17 00:00:00 2001
From: UncleCode <unclecode@kidocode.com>
Date: Fri, 29 Nov 2024 20:08:09 +0800
Subject: [PATCH]   Refactor Dockerfile and clean up main.py   - Enhanced
 Dockerfile for platform-specific installations     - Added ARG for
 TARGETPLATFORM and BUILDPLATFORM     - Improved GPU support conditional on
 TARGETPLATFORM   - Removed static pages mounting in main.py   - Streamlined
 code structure to improve maintainability

---
 Dockerfile | 25 ++++++++++++++++---------
 main.py    |  4 ----
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index bd71deae..2997590a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,9 @@
 # syntax=docker/dockerfile:1.4
 
-# Build arguments
+ARG TARGETPLATFORM
+ARG BUILDPLATFORM
+
+# Other build arguments
 ARG PYTHON_VERSION=3.10
 
 # Base stage with system dependencies
@@ -63,13 +66,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && rm -rf /var/lib/apt/lists/*
 
 # GPU support if enabled and architecture is supported
-RUN if [ "$ENABLE_GPU" = "true" ] && [ "$(dpkg --print-architecture)" != "arm64" ] ; then \
-        apt-get update && apt-get install -y --no-install-recommends \
-        nvidia-cuda-toolkit \
-        && rm -rf /var/lib/apt/lists/* ; \
-    else \
-        echo "Skipping NVIDIA CUDA Toolkit installation (unsupported architecture or GPU disabled)"; \
-    fi
+RUN if [ "$ENABLE_GPU" = "true" ] && [ "$TARGETPLATFORM" = "linux/amd64" ] ; then \
+    apt-get update && apt-get install -y --no-install-recommends \
+    nvidia-cuda-toolkit \
+    && rm -rf /var/lib/apt/lists/* ; \
+else \
+    echo "Skipping NVIDIA CUDA Toolkit installation (unsupported platform or GPU disabled)"; \
+fi
 
 # Create and set working directory
 WORKDIR /app
@@ -120,7 +123,11 @@ RUN pip install --no-cache-dir \
 RUN mkdocs build
 
 # Install Playwright and browsers
-RUN playwright install
+RUN if [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
+    playwright install chromium; \
+    elif [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+    playwright install chromium; \
+    fi
 
 # Expose port
 EXPOSE 8000 11235 9222 8080
diff --git a/main.py b/main.py
index 6d217410..d6c792e8 100644
--- a/main.py
+++ b/main.py
@@ -340,9 +340,6 @@ app.add_middleware(
     allow_headers=["*"],  # Allows all headers
 )
 
-# Mount the pages directory as a static directory
-app.mount("/pages", StaticFiles(directory=__location__ + "/pages"), name="pages")
-
 # API token security
 security = HTTPBearer()
 CRAWL4AI_API_TOKEN = os.getenv("CRAWL4AI_API_TOKEN") or "test_api_code"
@@ -364,7 +361,6 @@ if os.path.exists(__location__ + "/site"):
     app.mount("/mkdocs", StaticFiles(directory="site", html=True), name="mkdocs")
 
 site_templates = Jinja2Templates(directory=__location__ + "/site")
-templates = Jinja2Templates(directory=__location__ + "/pages")
 
 crawler_service = CrawlerService()