From 343458479d5c664c01a37a4af859304a85f51594 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 11 Nov 2024 18:40:13 -0800
Subject: [PATCH] Make sure TEST_PYPI_VERSION is used in docker builds

---
 llama_stack/distribution/build_container.sh | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index e5ec5b4e2..d0874d99f 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -78,7 +78,16 @@ if [ -n "$LLAMA_STACK_DIR" ]; then
   # rebuild. This is just for development convenience.
   add_to_docker "RUN pip install --no-cache -e $stack_mount"
 else
-  add_to_docker "RUN pip install --no-cache llama-stack"
+  if [ -n "$TEST_PYPI_VERSION" ]; then
+    # these packages are damaged in test-pypi, so install them first
+    add_to_docker "RUN pip install fastapi libcst"
+    add_to_docker <<EOF
+RUN pip install --no-cache --extra-index-url https://test.pypi.org/simple/ \
+  llama-models==$TEST_PYPI_VERSION llama-stack==$TEST_PYPI_VERSION
+EOF
+  else
+    add_to_docker "RUN pip install --no-cache llama-stack"
+  fi
 fi
 
 if [ -n "$LLAMA_MODELS_DIR" ]; then