diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 230ca34ac..2730ae174 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -97,7 +97,7 @@ else
     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
+  llama-models==$TEST_PYPI_VERSION llama-stack-client==$TEST_PYPI_VERSION llama-stack==$TEST_PYPI_VERSION
 EOF
   else
     add_to_docker "RUN pip install --no-cache llama-stack"