From cec5fdd14430455153802ee3a8abeecf31e9653d Mon Sep 17 00:00:00 2001
From: Jan-Timo Hesse <jan-timo.hesse@dlr.de>
Date: Mon, 5 Aug 2024 09:42:16 +0200
Subject: [PATCH] Added EXTRA_MODULES argument

---
 Dockerfile       | 4 +++-
 backend/start.sh | 4 ++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 8078bf0ea..f1ea4e064 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -38,6 +38,7 @@ ARG USE_OLLAMA
 ARG USE_CUDA_VER
 ARG USE_EMBEDDING_MODEL
 ARG USE_RERANKING_MODEL
+ARG EXTRA_MODULES
 ARG UID
 ARG GID
 
@@ -49,7 +50,8 @@ ENV ENV=prod \
     USE_CUDA_DOCKER=${USE_CUDA} \
     USE_CUDA_DOCKER_VER=${USE_CUDA_VER} \
     USE_EMBEDDING_MODEL_DOCKER=${USE_EMBEDDING_MODEL} \
-    USE_RERANKING_MODEL_DOCKER=${USE_RERANKING_MODEL}
+    USE_RERANKING_MODEL_DOCKER=${USE_RERANKING_MODEL} \
+    EXTRA_MODULES_DOCKER=${EXTRA_MODULES}
 
 ## Basis URL Config ##
 ENV OLLAMA_BASE_URL="/ollama" \
diff --git a/backend/start.sh b/backend/start.sh
index 16a004e45..b499736b8 100755
--- a/backend/start.sh
+++ b/backend/start.sh
@@ -30,6 +30,10 @@ if [[ "${USE_CUDA_DOCKER,,}" == "true" ]]; then
   export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/python3.11/site-packages/torch/lib:/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib"
 fi
 
+if [ -n "$EXTRA_MODULES_DOCKER" ]; then
+  echo "Loading extra modules: $EXTRA_MODULES_DOCKER"
+  uv pip install --system $EXTRA_MODULES_DOCKER --no-cache-dir
+fi
 
 # Check if SPACE_ID is set, if so, configure for space
 if [ -n "$SPACE_ID" ]; then