diff --git a/Containerfile.add_llama_to_lightspeed b/Containerfile.add_llama_to_lightspeed
index 46fcf97..49896c2 100644
--- a/Containerfile.add_llama_to_lightspeed
+++ b/Containerfile.add_llama_to_lightspeed
@@ -13,7 +13,7 @@ RUN cd /app-root/ && python3.12 -m pip install .
 
 RUN python3.12 -m pip install pyyaml pyaml
 
-RUN python3.12 -m pip install litellm
+RUN python3.12 -m pip install litellm[proxy]
 
 RUN python3.12 -m pip install sqlalchemy
 
diff --git a/Containerfile.assisted-chat b/Containerfile.assisted-chat
index 8058721..021c42b 100644
--- a/Containerfile.assisted-chat
+++ b/Containerfile.assisted-chat
@@ -4,7 +4,7 @@ FROM quay.io/lightspeed-core/lightspeed-stack@sha256:2cf68a775df9157689e95222e41
 
 RUN python3 -m ensurepip --default-pip && pip install --upgrade pip
 
-RUN python3 -m pip install pyyaml pyaml litellm sqlalchemy mcp psycopg2-binary
+RUN python3 -m pip install pyyaml pyaml litellm[proxy] sqlalchemy mcp psycopg2-binary
 
 USER 1001