From b5df8657e8db09716e781cabe1fcd0777f02ad2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 8 Jan 2026 18:22:42 +0100 Subject: [PATCH] chore: add flag to not include ml libs in docker image (#130) --- docker/standalone/Dockerfile | 31 ++++++++++++++++++++++++++----- hindsight-api/pyproject.toml | 7 ++++--- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/docker/standalone/Dockerfile b/docker/standalone/Dockerfile index f555d3e1..beca1653 100644 --- a/docker/standalone/Dockerfile +++ b/docker/standalone/Dockerfile @@ -2,19 +2,24 @@ # Supports building API-only, Control Plane-only, or both # # Build args: -# INCLUDE_API=true/false - Include API (default: true) -# INCLUDE_CP=true/false - Include Control Plane (default: true) -# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: true) +# INCLUDE_API=true/false - Include API (default: true) +# INCLUDE_CP=true/false - Include Control Plane (default: true) +# INCLUDE_LOCAL_MODELS=true/false - Include local ML models for embeddings/reranking (default: true) +# Set to false when using external providers (TEI, OpenAI, Cohere) +# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: true) +# Only effective when INCLUDE_LOCAL_MODELS=true # # Examples: # docker build -t hindsight . # Both (standalone) # docker build -t hindsight-api --build-arg INCLUDE_CP=false . # API only # docker build -t hindsight-cp --build-arg INCLUDE_API=false . # Control Plane only # docker build -t hindsight --build-arg PRELOAD_ML_MODELS=false . # Skip ML model preload +# docker build -t hindsight --build-arg INCLUDE_LOCAL_MODELS=false . # Skip local ML deps (for external providers) ARG INCLUDE_API=true ARG INCLUDE_CP=true ARG PRELOAD_ML_MODELS=true +ARG INCLUDE_LOCAL_MODELS=true # ============================================================================= # Stage: API Builder @@ -22,6 +27,7 @@ ARG PRELOAD_ML_MODELS=true FROM python:3.11-slim AS api-builder ARG INCLUDE_API +ARG INCLUDE_LOCAL_MODELS RUN if [ "$INCLUDE_API" != "true" ]; then echo "Skipping API build" && exit 0; fi WORKDIR /app @@ -40,6 +46,15 @@ COPY hindsight-api/README.md ./api/ WORKDIR /app/api +# Remove local ML model dependencies if INCLUDE_LOCAL_MODELS=false +# This creates a smaller image when using external providers (TEI, OpenAI, Cohere) +RUN if [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then \ + echo "Removing local-models dependencies (sentence-transformers, torch, transformers)..." && \ + sed -i '/"sentence-transformers/d' pyproject.toml && \ + sed -i '/"transformers/d' pyproject.toml && \ + sed -i '/"torch/d' pyproject.toml; \ + fi + # Sync dependencies (will create lock file if needed) RUN uv sync @@ -153,8 +168,10 @@ USER hindsight ENV PATH="/app/api/.venv/bin:${PATH}" # Pre-download ML models to avoid runtime download (conditional) +# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true ARG PRELOAD_ML_MODELS -RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \ +ARG INCLUDE_LOCAL_MODELS +RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \ /app/api/.venv/bin/python -c "\ from sentence_transformers import SentenceTransformer, CrossEncoder; \ print('Downloading embedding model...'); \ @@ -162,6 +179,7 @@ SentenceTransformer('BAAI/bge-small-en-v1.5'); \ print('Downloading cross-encoder model...'); \ CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \ print('Models cached successfully')"; \ + elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \ else echo "Skipping ML model preload"; fi EXPOSE 8888 @@ -258,8 +276,10 @@ USER hindsight ENV PATH="/app/api/.venv/bin:${PATH}" # Pre-download ML models to avoid runtime download (conditional) +# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true ARG PRELOAD_ML_MODELS -RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \ +ARG INCLUDE_LOCAL_MODELS +RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \ /app/api/.venv/bin/python -c "\ from sentence_transformers import SentenceTransformer, CrossEncoder; \ print('Downloading embedding model...'); \ @@ -267,6 +287,7 @@ SentenceTransformer('BAAI/bge-small-en-v1.5'); \ print('Downloading cross-encoder model...'); \ CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \ print('Models cached successfully')"; \ + elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \ else echo "Skipping ML model preload"; fi EXPOSE 8888 9999 diff --git a/hindsight-api/pyproject.toml b/hindsight-api/pyproject.toml index 5c35d652..1bf76ed1 100644 --- a/hindsight-api/pyproject.toml +++ b/hindsight-api/pyproject.toml @@ -14,7 +14,6 @@ dependencies = [ "openai>=1.0.0", "pydantic>=2.0.0", "rich>=13.0.0", - "sentence-transformers>=3.0.0,<3.3.0", "langchain-text-splitters>=0.3.0", "fastapi[standard]>=0.120.3", "uvicorn>=0.38.0", @@ -24,8 +23,6 @@ dependencies = [ "pgvector>=0.4.1", "greenlet>=3.2.4", "psycopg2-binary>=2.9.11", - "transformers>=4.30.0,<4.46.0", - "torch>=2.0.0", "tiktoken>=0.12.0", "httpx>=0.27.0", "fastmcp>=2.3.0", @@ -40,6 +37,10 @@ dependencies = [ "anthropic>=0.40.0", "typer>=0.9.0", "cohere>=5.0.0", + # Local ML models for embeddings/reranking - can be excluded in Docker with INCLUDE_LOCAL_MODELS=false + "sentence-transformers>=3.0.0,<3.3.0", + "transformers>=4.30.0,<4.46.0", + "torch>=2.0.0", ] [project.optional-dependencies]