From 83adaae09bbb7af1142f12454b113852dcb0caae Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Thu, 8 Aug 2024 12:55:03 -0700
Subject: [PATCH] allow installing from test.pypi.org

---
 llama_toolchain/distribution/install_distribution.sh | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/llama_toolchain/distribution/install_distribution.sh b/llama_toolchain/distribution/install_distribution.sh
index 80727725d..7c6cc97f4 100755
--- a/llama_toolchain/distribution/install_distribution.sh
+++ b/llama_toolchain/distribution/install_distribution.sh
@@ -8,6 +8,7 @@
 
 LLAMA_MODELS_DIR=${LLAMA_MODELS_DIR:-}
 LLAMA_TOOLCHAIN_DIR=${LLAMA_TOOLCHAIN_DIR:-}
+USE_TEST_PYPI=${USE_TEST_PYPI:-}
 
 set -euo pipefail
 
@@ -56,6 +57,12 @@ ensure_conda_env_python310() {
   eval "$(conda shell.bash hook)"
   conda deactivate && conda activate "${env_name}"
 
+  PIP_ARGS=""
+  if [ "$USE_TEST_PYPI" = "1" ]; then
+    # these packages are damaged in test-pypi, so install them first
+    pip install fastapi libcst
+    PIP_ARGS="--extra-index-url https://test.pypi.org/simple/"
+  fi
   # Re-installing llama-toolchain in the new conda environment
   if [ -n "$LLAMA_TOOLCHAIN_DIR" ]; then
     if [ ! -d "$LLAMA_TOOLCHAIN_DIR" ]; then
@@ -66,7 +73,7 @@ ensure_conda_env_python310() {
     echo "Installing from LLAMA_TOOLCHAIN_DIR: $LLAMA_TOOLCHAIN_DIR"
     pip install -e "$LLAMA_TOOLCHAIN_DIR"
   else
-    pip install llama-toolchain
+    pip install $PIP_ARGS llama-toolchain
   fi
 
   if [ -n "$LLAMA_MODELS_DIR" ]; then
@@ -83,7 +90,7 @@ ensure_conda_env_python310() {
   # Install pip dependencies
   if [ -n "$pip_dependencies" ]; then
     echo "Installing pip dependencies: $pip_dependencies"
-    pip install $pip_dependencies
+    pip install $PIP_ARGS $pip_dependencies
   fi
 }